inspect connectors

Author: prishajain1

inspect_connectors.py

@@ -0,0 +1,30 @@
+
+import torch
+from safetensors.torch import load_file
+from huggingface_hub import snapshot_download
+import os
+
+def inspect_connectors():
+    repo_id = "Lightricks/LTX-2"
+    subfolder = "connectors"
+    
+    print(f"Downloading {subfolder} from {repo_id}...")
+    try:
+        path = snapshot_download(repo_id, allow_patterns=[f"{subfolder}/*"])
+        print(f"Downloaded to {path}")
+        
+        folder_path = os.path.join(path, subfolder)
+        files = [f for f in os.listdir(folder_path) if f.endswith(".safetensors")]
+        
+        for f in files:
+            print(f"\nScanning {f}...")
+            file_path = os.path.join(folder_path, f)
+            state_dict = load_file(file_path)
+            for key, value in state_dict.items():
+                print(f"{key}: {value.shape}")
+                
+    except Exception as e:
+        print(f"Error: {e}")
+
+if __name__ == "__main__":
+    inspect_connectors()

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -372,7 +372,6 @@ def load_vocoder_weights(
   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(".")
 
@@ -388,13 +387,10 @@ def load_vocoder_weights(
 
       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)