vocoder weight

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -382,6 +382,8 @@ def load_vocoder_weights(
     if filename and pt_key.startswith("vocoder."):
       pt_key = pt_key[len("vocoder."):]
     key = rename_for_ltx2_vocoder(pt_key)
+    if filename == "ltx-2.3-22b-dev.safetensors":
+      key = key.replace("resblocks_", "resnets.")
     parts = key.split(".")
 
     if parts[-1] == "weight":
@@ -398,6 +400,9 @@ def load_vocoder_weights(
         tensor = tensor.transpose(2, 0, 1)[::-1, :, :]
       else:
         tensor = tensor.transpose(2, 1, 0)
+    
+    if "mel_stft" in flax_key and ("forward_basis" in flax_key or "inverse_basis" in flax_key):
+      tensor = tensor.transpose(2, 1, 0)
 
     flax_state_dict[flax_key] = jax.device_put(tensor, device=cpu)
 

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -546,6 +546,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
         )
         bwe_generator = Vocoder(
             upsample_initial_channel=512,
+            upsample_kernel_sizes=[12, 11, 4, 4, 4],
             rngs=rngs,
             dtype=jnp.float32,
         )