transformer weight loading

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -113,20 +113,20 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
   # 1. Load config.
   if restored_checkpoint:
     ltx2_config = restored_checkpoint["ltx2_config"]
-    elif getattr(config, "model_name", "") == "ltx2.3":
-      ltx2_config = {
-          "in_channels": 128,
-          "num_attention_heads": 32,
-          "attention_head_dim": 128,
-          "cross_attention_dim": 4096,
-          "audio_in_channels": 128,
-          "audio_num_attention_heads": 32,
-          "audio_attention_head_dim": 64,
-          "audio_cross_attention_dim": 2048,
-          "num_layers": 48,
-          "caption_channels": 8192,
-          "audio_caption_channels": 4096,
-      }
+  elif getattr(config, "model_name", "") == "ltx2.3":
+    ltx2_config = {
+        "in_channels": 128,
+        "num_attention_heads": 32,
+        "attention_head_dim": 128,
+        "cross_attention_dim": 4096,
+        "audio_in_channels": 128,
+        "audio_num_attention_heads": 32,
+        "audio_attention_head_dim": 64,
+        "audio_cross_attention_dim": 2048,
+        "num_layers": 48,
+        "caption_channels": 8192,
+        "audio_caption_channels": 4096,
+    }
   else:
     ltx2_config = LTX2VideoTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder=subfolder)