weights loading for vocoder

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -354,3 +354,53 @@ def load_vae_weights(
       jax.clear_caches()
       return flax_state_dict
 
+def rename_for_ltx2_vocoder(key):
+    key = key.replace("ups", "upsamplers")
+    key = key.replace("resblocks", "resnets")
+    key = key.replace("conv_post", "conv_out")
+    return key
+
+
+def load_vocoder_weights(
+    pretrained_model_name_or_path: str,
+    eval_shapes: dict,
+    device: str,
+    hf_download: bool = True,
+    subfolder: str = "vocoder"
+):
+  tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device)
+  
+  flax_state_dict = {}
+  cpu = jax.local_devices(backend="cpu")[0]
+
+  for pt_key, tensor in tensors.items():
+      # Initial renaming
+      key = rename_for_ltx2_vocoder(pt_key)
+      parts = key.split(".")
+      
+      flax_key_parts = []
+      for part in parts:
+          if part.isdigit():
+              flax_key_parts.append(int(part))
+          else:
+              flax_key_parts.append(part)
+              
+      if flax_key_parts[-1] == "weight":
+          flax_key_parts[-1] = "kernel"
+      
+      flax_key = tuple(flax_key_parts)
+      
+      # Transpose weights
+      if flax_key[-1] == "kernel":
+           if "upsamplers" in flax_key:
+               # ConvTranspose: (In, Out, K) -> (K, In, Out)
+               tensor = tensor.transpose(2, 0, 1)
+           else:
+               # Conv: (Out, In, K) -> (K, In, Out)
+               tensor = tensor.transpose(2, 1, 0)
+               
+      flax_state_dict[flax_key] = jax.device_put(tensor, device=cpu)
+      
+  validate_flax_state_dict(eval_shapes, flax_state_dict)
+  return unflatten_dict(flax_state_dict)
+

src/maxdiffusion/tests/test_ltx2_utils.py

@@ -113,7 +113,6 @@ def test_load_vae_weights(self):
         )
 
         print("Validating VAE Weights...")
-        # Filter out dropout/rngs keys from eval_shapes as they are not expected in weights
         filtered_eval_shapes = {}
         flat_eval_shapes = flatten_dict(eval_shapes)
         for k, v in flat_eval_shapes.items():
@@ -126,5 +125,79 @@ def test_load_vae_weights(self):
         validate_flax_state_dict(unflatten_dict(filtered_eval_shapes), flatten_dict(loaded_weights))
         print("VAE Weights Validated Successfully!")
 
+    def test_load_vocoder_weights(self):
+        from maxdiffusion.models.ltx2.vocoder_ltx2 import LTX2Vocoder
+        from unittest import mock
+        import torch
+        from maxdiffusion.models.ltx2.ltx2_utils import load_vocoder_weights
+        
+        config = {
+          "hidden_channels": 1024,
+          "in_channels": 128,
+          "leaky_relu_negative_slope": 0.1,
+          "out_channels": 2,
+          "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+          "resnet_kernel_sizes": [3, 7, 11],
+          "upsample_factors": [6, 5, 2, 2, 2],
+          "upsample_kernel_sizes": [16, 15, 8, 4, 4],
+          "rngs": nnx.Rngs(0)
+        }
+        with jax.default_device(jax.devices("cpu")[0]):
+            model = LTX2Vocoder(**config)
+        state = nnx.state(model)
+        eval_shapes = state.to_pure_dict()
+        
+        # Create dummy PyTorch weights
+        pt_weights = {}
+        flat_shapes = flatten_dict(eval_shapes)
+        
+        for key, value in flat_shapes.items():
+            # key is tuple of strings/ints, e.g. ('conv_in', 'kernel')
+            # Map to PyTorch key
+            pt_key_parts = []
+            is_upsampler = "upsamplers" in [str(k) for k in key]
+            is_kernel = False
+            
+            for part in key:
+                if str(part) == "upsamplers":
+                    pt_key_parts.append("ups")
+                elif str(part) == "resnets":
+                    pt_key_parts.append("resblocks")
+                elif str(part) == "conv_out":
+                    pt_key_parts.append("conv_post")
+                elif str(part) == "kernel":
+                    pt_key_parts.append("weight")
+                    is_kernel = True
+                else:
+                    pt_key_parts.append(str(part))
+            
+            pt_key = ".".join(pt_key_parts)
+            
+            # Create tensor with PyTorch shape
+            shape = value.shape
+            if is_kernel:
+                if is_upsampler:
+                    # Flax (K, I, O) -> PyTorch (I, O, K)
+                    pt_shape = (shape[1], shape[2], shape[0])
+                else:
+                    # Flax (K, I, O) -> PyTorch (O, I, K)
+                    pt_shape = (shape[2], shape[1], shape[0])
+            else:
+                pt_shape = shape
+                
+            pt_weights[pt_key] = torch.randn(pt_shape)
+
+        with mock.patch("maxdiffusion.models.ltx2.ltx2_utils.load_sharded_checkpoint", return_value=pt_weights):
+            loaded_weights = load_vocoder_weights(
+                pretrained_model_name_or_path="dummy",
+                eval_shapes=eval_shapes,
+                device=self.device,
+                hf_download=False
+            )
+            
+        # Validate
+        validate_flax_state_dict(eval_shapes, flatten_dict(loaded_weights))
+        print("Vocoder Weights Validated Successfully!")
+
 if __name__ == "__main__":
     unittest.main()