transformer weight

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -220,13 +220,14 @@ def load_transformer_weights(
     num_layers: int = 48,
     scan_layers: bool = True,
     subfolder: str = "transformer",
+    filename: str = None,
 ):
   device = jax.local_devices(backend=device)[0]
   max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}")
 
   with jax.default_device(device):
     # Support sharded loading
-    tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device)
+    tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, filename=filename)
 
     flax_state_dict = {}
     cpu = jax.local_devices(backend="cpu")[0]

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -168,12 +168,15 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
     else:
       params = restored_checkpoint["ltx2_state"]
   else:
+    filename = "ltx-2.3-22b-dev.safetensors" if getattr(config, "model_name", "") == "ltx2.3" else None
+    subfolder = "" if getattr(config, "model_name", "") == "ltx2.3" else subfolder
     params = load_transformer_weights(
         config.pretrained_model_name_or_path,
         params,  # eval_shapes
         "cpu",
         scan_layers=getattr(config, "scan_layers", True),
         subfolder=subfolder,
+        filename=filename,
     )
 
   params = jax.tree_util.tree_map_with_path(