transformer and vocoder keys

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -93,7 +93,19 @@ def rename_for_ltx2_transformer(key):
   if "to_out_0" in key:
     key = key.replace("to_out_0", "to_out")
 
+  # Add missing mappings
+  key = key.replace("av_ca_video_scale_shift_adaln_single", "av_cross_attn_video_scale_shift")
+  key = key.replace("av_ca_a2v_gate_adaln_single", "av_cross_attn_video_a2v_gate")
+  key = key.replace("av_ca_audio_scale_shift_adaln_single", "av_cross_attn_audio_scale_shift")
+  key = key.replace("av_ca_v2a_gate_adaln_single", "av_cross_attn_audio_v2a_gate")
+  key = key.replace("scale_shift_table_a2v_ca_video", "video_a2v_cross_attn_scale_shift_table")
+  key = key.replace("scale_shift_table_a2v_ca_audio", "audio_a2v_cross_attn_scale_shift_table")
+
   # LTX-2.3 specific mappings
+  # Handle substrings before they are replaced by shorter patterns below
+  key = key.replace("audio_prompt_adaln_single", "audio_prompt_adaln")
+  key = key.replace("prompt_adaln_single", "prompt_adaln")
+
   if "prompt_adaln" in key:
     key = key.replace("prompt_adaln", "caption_projection")
   if "audio_prompt_adaln" in key:
@@ -337,6 +349,8 @@ def rename_for_ltx2_vocoder(key):
   key = key.replace("ups.", "upsamplers.")
   key = key.replace("resblocks", "resnets")
   key = key.replace("conv_post", "conv_out")
+  key = key.replace("conv_pre", "conv_in")
+  key = key.replace("act_post", "act_out")
 
   # LTX-2.3 specific mappings for Vocoder
   if "downsample" in key and "lowpass" not in key:

src/maxdiffusion/models/ltx2/vocoder_bwe_ltx2.py

@@ -384,9 +384,9 @@ def __init__(
         dtype=self.dtype,
     )
 
-    self.ups = nnx.List()
+    self.upsamplers = nnx.List()
     for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-      self.ups.append(
+      self.upsamplers.append(
           nnx.ConvTranspose(
               in_features=upsample_initial_channel // (2**i),
               out_features=upsample_initial_channel // (2 ** (i + 1)),
@@ -398,11 +398,11 @@ def __init__(
           )
       )
 
-    self.resblocks = nnx.List()
+    self.resnets = nnx.List()
     for i in range(len(upsample_rates)):
       ch = upsample_initial_channel // (2 ** (i + 1))
       for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes):
-        self.resblocks.append(AMPBlock1(ch, kernel_size, dilations, activation=activation, rngs=rngs))
+        self.resnets.append(AMPBlock1(ch, kernel_size, dilations, activation=activation, rngs=rngs))
 
     final_channels = upsample_initial_channel // (2 ** len(upsample_rates))
     self.act_out = Activation1d(final_channels, SnakeBeta(final_channels, rngs=rngs), rngs=rngs)
@@ -429,14 +429,14 @@ def __call__(self, hidden_states: Array, time_last: bool = False) -> Array:
     hidden_states = self.conv_in(hidden_states)
 
     for i in range(self.num_upsamples):
-      hidden_states = self.ups[i](hidden_states)
+      hidden_states = self.upsamplers[i](hidden_states)
 
       start = i * self.num_kernels
       end = (i + 1) * self.num_kernels
 
       res_sum = 0.0
       for j in range(start, end):
-        res_sum = res_sum + self.resblocks[j](hidden_states)
+        res_sum = res_sum + self.resnets[j](hidden_states)
 
       hidden_states = res_sum / self.num_kernels