params names changed in transformer file

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -107,11 +107,8 @@ def __init__(
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         self.audio_norm1 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -123,81 +120,66 @@ def __init__(
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         # 2. Prompt Cross-Attention
         self.norm2 = nnx.RMSNorm(self.dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
         self.attn2 = LTX2Attention(
             rngs=rngs,
             query_dim=dim,
-            cross_attention_dim=cross_attention_dim,
+            context_dim=cross_attention_dim,
             heads=num_attention_heads,
             dim_head=attention_head_dim,
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         self.audio_norm2 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
         self.audio_attn2 = LTX2Attention(
             rngs=rngs,
             query_dim=audio_dim,
-            cross_attention_dim=audio_cross_attention_dim,
+            context_dim=audio_cross_attention_dim,
             heads=audio_num_attention_heads,
             dim_head=audio_attention_head_dim,
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
         self.audio_to_video_norm = nnx.RMSNorm(dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
         self.audio_to_video_attn = LTX2Attention(
             rngs=rngs,
             query_dim=dim,
-            cross_attention_dim=audio_dim,
+            context_dim=audio_dim,
             heads=audio_num_attention_heads,
             dim_head=audio_attention_head_dim,
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         self.video_to_audio_norm = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
         self.video_to_audio_attn = LTX2Attention(
             rngs=rngs,
             query_dim=audio_dim,
-            cross_attention_dim=dim,
+            context_dim=dim,
             heads=audio_num_attention_heads,
             dim_head=audio_attention_head_dim,
             dropout=0.0,
             bias=attention_bias,
             out_bias=attention_out_bias,
-            qk_norm=qk_norm,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-            dtype=dtype,
-            param_dtype=weights_dtype
+            eps=norm_eps,
+            dtype=dtype
         )
 
         # 4. Feed Forward