connectors

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -410,6 +410,8 @@ def rename_for_ltx2_connector(key):
         # Handle FF
         key = key.replace("ff.net.0.proj", "ff.proj1")
         key = key.replace("ff.net.2", "ff.proj2")
+        # Handle to_out
+        key = key.replace("to_out.0", "to_out")
 
     # Validation/Weight suffix
     if key.endswith(".weight"):

src/maxdiffusion/models/ltx2/text_encoders/text_encoders_ltx2.py

@@ -38,7 +38,7 @@ def __init__(
       # Feature Extractor Config
       gemma_dim: int = 3840,  # Gemma-3-12b
       gemma_layers: int = 49,  # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
-      projection_dim: int = 3840,  # LTX-2 conditioning dim
+      projection_dim: int = 4096,  # LTX-2 conditioning dim
       # Connector Config
       connector_heads: int = 32,
       connector_head_dim: int = 128,
@@ -98,7 +98,7 @@ def __init__(
       # Feature Extractor Config (Shared)
       gemma_dim: int = 3840,  # Gemma-3-12b
       gemma_layers: int = 49,  # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
-      projection_dim: int = 3840,
+      projection_dim: int = 4096,
       # Connector Config
       connector_heads: int = 32,
       connector_head_dim: int = 128,