headers

Author: Serenagu525

src/maxdiffusion/generate_ltx_video.py

@@ -112,6 +112,3 @@ def main(argv: Sequence[str]) -> None:
 
 if __name__ == "__main__":
   app.run(main)
-
-
-

src/maxdiffusion/max_utils.py

@@ -271,7 +271,7 @@ def create_device_mesh(config, devices=None, logging=True):
     max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
 
     return mesh
-    
+
   try:
     num_slices = 1 + max([d.slice_index for d in devices])
   except:
@@ -303,66 +303,9 @@ def create_device_mesh(config, devices=None, logging=True):
   if logging:
     max_logging.log(f"Decided on mesh: {mesh}")
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
   return mesh
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 def unbox_logicallypartioned_trainstate(boxed_train_state: train_state.TrainState):
   """Unboxes the flax.LogicallyPartitioned pieces in a train state.
 
@@ -685,4 +628,4 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     initialize_jax_for_gpu()
     max_logging.log("Jax distributed system initialized on GPU!")
   else:
-    jax.distributed.initialize()
+    jax.distributed.initialize()

src/maxdiffusion/models/ltx_video/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025 Lightricks Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Lightricks/LTX-Video/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This implementation is based on the Torch version available at:
+# https://github.com/Lightricks/LTX-Video/tree/main

src/maxdiffusion/models/ltx_video/transformers/__init__.py

@@ -0,0 +1,15 @@
+"""
+ Copyright 2025 Google LLC
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      https://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""

src/maxdiffusion/models/ltx_video/transformers/attention.py

@@ -452,7 +452,7 @@ def __call__(
       deterministic: bool = True,
       **cross_attention_kwargs,
   ) -> jnp.ndarray:
-    cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters} #noqa: F821
+    cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}  # noqa: F821
     assert cross_attention_kwargs.get("scale", None) is None, "Not supported"
 
     input_axis_names = ("activation_batch", "activation_length", "activation_embed")

src/maxdiffusion/models/ltx_video/transformers/transformer3d.py

@@ -112,7 +112,7 @@ def scale_shift_table_init(key):
       self.transformer_blocks = RepeatableLayer(
           RemattedBasicTransformerBlock,
           num_layers=self.num_layers,
-          module_init_kwargs=dict( #noqa: C408
+          module_init_kwargs=dict(  # noqa: C408
               dim=self.inner_dim,
               num_attention_heads=self.num_attention_heads,
               attention_head_dim=self.attention_head_dim,

src/maxdiffusion/models/ltx_video/utils/conversion_script_instruction.md

@@ -1,13 +1,10 @@
 ### Transformer Pytorch Weight Downloading and Jax Weight Loading Instructions:
-1. Create new tansformers_pytorch folder under models/ltx_video.
-2. Move files from LTX repo, specifically, attention.py, embeddings.py, symmetric_patchifier.py, and transformer3d.py into the newly created folder. See here: https://github.com/Lightricks/LTX-Video/tree/main/ltx_video/models/transformers 
-3. Rename transformer3d.py to transformer_pt.py to distinguish from the pytorch version. Change classname to Transformer3DModel_PT. Also change classname in line "transformer = Transformer3DModel.from_config(transformer_config)"
-4. Weight Downloading and Conversion
+1. Weight Downloading and Conversion
     - If first time running (no local safetensors): \
     In the src/maxdiffusion/models/ltx_video/utils folder, run python convert_torch_weights_to_jax.py --download_ckpt_path [location to download safetensors] --output_dir [location to save jax ckpt] --transformer_config_path ../xora_v1.2-13B-balanced-128.json.
     - If already have local pytorch checkpoint: \
     Replace the --download_ckpt_path with --local_ckpt_path and add corresponding location
-5. Restoring Jax Weights into transformer:
+2. Restoring Jax Weights into transformer:
     - Replace the "ckpt_path" in src/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json with jax ckpt path.
     - Run python src/maxdiffusion/generate_ltx_video.py src/maxdiffusion/configs/ltx_video.yml in the outer repo folder.
 

src/maxdiffusion/models/ltx_video/utils/convert_torch_weights_to_jax.py

@@ -1,88 +1,103 @@
+# Copyright 2025 Lightricks Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Lightricks/LTX-Video/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This implementation is based on the Torch version available at:
+# https://github.com/Lightricks/LTX-Video/tree/main
 import argparse
 import json
 from typing import Any, Dict, Optional
 
 
-
 import jax
 import jax.numpy as jnp
 from flax.training import train_state
 import optax
 import orbax.checkpoint as ocp
 from safetensors.torch import load_file
 import requests
-import shutil
 from urllib.parse import urljoin
 
-# from maxdiffusion.models.ltx_video.transformers_pytorch.transformer import Transformer3DModel
 from maxdiffusion.models.ltx_video.transformers.transformer3d import Transformer3DModel as JaxTranformer3DModel
 from maxdiffusion.models.ltx_video.utils.torch_compat import torch_statedict_to_jax
 
 from huggingface_hub import hf_hub_download
 import os
 import importlib
+
+
 def download_and_move_files(github_base_url, base_path, target_folder_name, files_to_move, module_to_import):
-    """
-    Downloads files from a GitHub repo, moves them to a local folder, and then dynamically imports a module.
+  """
+  Downloads files from a GitHub repo, moves them to a local folder, and then dynamically imports a module.
+
+  Args:
+      github_base_url (str): The base URL of the GitHub repo.
+      base_path (str): The base path where the new folder will be created.
+      target_folder_name (str): The name of the folder to create.
+      files_to_move (list): A list of file names to download and move.
+      module_to_import (str): The full module path to import.
+  """
 
-    Args:
-        github_base_url (str): The base URL of the GitHub repo.
-        base_path (str): The base path where the new folder will be created.
-        target_folder_name (str): The name of the folder to create.
-        files_to_move (list): A list of file names to download and move.
-        module_to_import (str): The full module path to import.
-    """
+  target_path = os.path.join(base_path, target_folder_name)
 
-    target_path = os.path.join(base_path, target_folder_name)
+  try:
+    # Create the target directory
+    os.makedirs(target_path, exist_ok=True)
+    print(f"Created directory: {target_path}")
 
+    # Download and move files
+    for file_name in files_to_move:
+      file_url = urljoin(github_base_url, file_name)
+      destination_path = os.path.join(target_path, file_name)
+
+      try:
+        response = requests.get(file_url, stream=True)
+        response.raise_for_status()
+
+        with open(destination_path, "wb") as f:
+          for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+
+        print(f"Downloaded and moved: {file_name} -> {destination_path}")
+
+      except requests.exceptions.RequestException as e:
+        print(f"Error downloading {file_name}: {e}")
+      except OSError as e:
+        print(f"Error writing file {file_name}: {e}")
+    print("Files downloaded and moved successfully.")
+
+    # Verify that the folder exists
+    if not os.path.exists(target_path):
+      print(f"Error: Target folder {target_path} does not exist after files download.")
+    # Dynamically import the module
     try:
-        # Create the target directory
-        os.makedirs(target_path, exist_ok=True)
-        print(f"Created directory: {target_path}")
-
-        # Download and move files
-        for file_name in files_to_move:
-            file_url = urljoin(github_base_url, file_name)
-            destination_path = os.path.join(target_path, file_name)
-
-            try:
-                response = requests.get(file_url, stream=True)
-                response.raise_for_status()
-
-                with open(destination_path, 'wb') as f:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        f.write(chunk)
-
-                print(f"Downloaded and moved: {file_name} -> {destination_path}")
-
-            except requests.exceptions.RequestException as e:
-                print(f"Error downloading {file_name}: {e}")
-                return # Stop if there is an error.
-            except OSError as e:
-                print(f"Error writing file {file_name}: {e}")
-                return # Stop if there is an error.
-        print("Files downloaded and moved successfully.")
-
-        # Verify that the folder exists
-        if not os.path.exists(target_path):
-          print(f"Error: Target folder {target_path} does not exist after files download.")
-        # Dynamically import the module
-        try:
-          imported_module = importlib.import_module(module_to_import)
-          print(f"Module '{module_to_import}' imported successfully.")
-          # Access the class
-          transformer_class = getattr(imported_module, "Transformer3DModel")
-          print(f"Class 'Transformer3DModel' accessed successfully: {transformer_class}")
-          return transformer_class
-        except ImportError as e:
-            print(f"Error importing module '{module_to_import}': {e}")
-        except AttributeError as e:
-            print(f"Error accessing class 'Transformer3DModel': {e}")
-
-    except OSError as e:
-        print(f"Error during file system operation: {e}")
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
+      imported_module = importlib.import_module(module_to_import)
+      print(f"Module '{module_to_import}' imported successfully.")
+      # Access the class
+      transformer_class = getattr(imported_module, "Transformer3DModel")
+      print(f"Class 'Transformer3DModel' accessed successfully: {transformer_class}")
+      return transformer_class
+    except ImportError as e:
+      print(f"Error importing module '{module_to_import}': {e}")
+    except AttributeError as e:
+      print(f"Error accessing class 'Transformer3DModel': {e}")
+
+  except OSError as e:
+    print(f"Error during file system operation: {e}")
+  except Exception as e:
+    print(f"An unexpected error occurred: {e}")
+
+
 class Checkpointer:
   """
   Checkpointer - to load and store JAX checkpoints
@@ -204,13 +219,13 @@ def main(args):
     )
   print("Downloading files from GitHub...")
   github_url = "https://raw.githubusercontent.com/Lightricks/LTX-Video/main/ltx_video/models/transformers/"
-  ltx_repo_path = "../"  
+  ltx_repo_path = "../"
   target_folder = "transformers_pytorch"
   files = ["attention.py", "embeddings.py", "symmetric_patchifier.py", "transformer3d.py"]
   module_path = "maxdiffusion.models.ltx_video.transformers_pytorch.transformer3d"
 
   Transformer3DModel = download_and_move_files(github_url, ltx_repo_path, target_folder, files, module_path)
-  
+
   print("Loading safetensors, flush = True")
   weight_file = "ltxv-13b-0.9.7-dev.safetensors"
 

src/maxdiffusion/models/ltx_video/utils/diffusers_config_mapping.py

@@ -1,3 +1,19 @@
+# Copyright 2025 Lightricks Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Lightricks/LTX-Video/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This implementation is based on the Torch version available at:
+# https://github.com/Lightricks/LTX-Video/tree/main
 def make_hashable_key(dict_key):
   def convert_value(value):
     if isinstance(value, list):

src/maxdiffusion/models/ltx_video/utils/skip_layer_strategy.py

@@ -1,3 +1,19 @@
+# Copyright 2025 Lightricks Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Lightricks/LTX-Video/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This implementation is based on the Torch version available at:
+# https://github.com/Lightricks/LTX-Video/tree/main
 from enum import Enum, auto