transformer weight loading

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -117,6 +117,8 @@ def rename_for_ltx2_transformer(key):
   if "audio_text_proj_in" in key:
     key = key.replace("audio_text_proj_in", "feature_extractor.audio_linear")
 
+  key = key.replace("k_norm", "norm_k")
+  key = key.replace("q_norm", "norm_q")
   key = key.replace("adaln_single", "time_embed")
   return key
 

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -590,6 +590,7 @@ def __init__(
       norm_elementwise_affine: bool = False,
       norm_eps: float = 1e-6,
       caption_channels: int = 3840,
+      audio_caption_channels: int = None,
       attention_bias: bool = True,
       attention_out_bias: bool = True,
       rope_theta: float = 10000.0,
@@ -643,6 +644,7 @@ def __init__(
     self.norm_elementwise_affine = norm_elementwise_affine
     self.norm_eps = norm_eps
     self.caption_channels = caption_channels
+    self.audio_caption_channels = audio_caption_channels or caption_channels
     self.attention_bias = attention_bias
     self.attention_out_bias = attention_out_bias
     self.rope_theta = rope_theta
@@ -703,7 +705,7 @@ def __init__(
       )
       self.audio_caption_projection = NNXCombinedTimestepTextProjEmbeddings(
           rngs=rngs,
-          in_features=self.caption_channels,
+          in_features=self.audio_caption_channels,
           hidden_size=audio_inner_dim,
           embedding_dim=audio_inner_dim,
           dtype=self.dtype,
@@ -719,7 +721,7 @@ def __init__(
       )
       self.audio_caption_projection = NNXPixArtAlphaTextProjection(
           rngs=rngs,
-          in_features=self.caption_channels,
+          in_features=self.audio_caption_channels,
           hidden_size=audio_inner_dim,
           dtype=self.dtype,
           weights_dtype=self.weights_dtype,

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -113,18 +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,
-    }
+    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)