Update port_tiny_to_deep.py

port_tiny_to_deep.py

@@ -2,7 +2,7 @@
 # TinyFlux → TinyFlux-Deep Porting Script
 # ============================================================================
 # Expands: 3 single + 3 double → 25 single + 15 double
-# Heads: 2 → 8 (old heads become first and last)
+# Heads: 2 → 4 (doubles heads, hidden 256 → 512)
 # Freezes ported layers, trains new ones
 # ============================================================================
 
@@ -12,6 +12,7 @@ from safetensors.torch import load_file, save_file
 from huggingface_hub import hf_hub_download, HfApi
 from dataclasses import dataclass
 from copy import deepcopy
+from typing import Tuple
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.bfloat16
@@ -21,34 +22,71 @@ DTYPE = torch.bfloat16
 # ============================================================================
 @dataclass
 class TinyFluxConfig:
-    """Original small config"""
+    """Original small config - matches TinyFlux model on hub (hidden=768, 6 heads)"""
+    # Core dimensions (detected from hub: 768 hidden, 6 heads)
     hidden_size: int = 768
-    num_attention_heads: int = 2
-    attention_head_dim: int = 128
-    num_single_blocks: int = 3
-    num_double_blocks: int = 3
-    mlp_ratio: float = 4.0
-    t5_embed_dim: int = 768
-    clip_embed_dim: int = 768
+    num_attention_heads: int = 6
+    attention_head_dim: int = 128  # 6 * 128 = 768
+    
+    # Input/output
     in_channels: int = 16
-    axes_dims: tuple = (16, 24, 24)
-    theta: int = 10000
+    patch_size: int = 1
+    
+    # Text encoder interfaces
+    joint_attention_dim: int = 768
+    pooled_projection_dim: int = 768
+    
+    # Layers
+    num_double_layers: int = 3
+    num_single_layers: int = 3
+    
+    # MLP
+    mlp_ratio: float = 4.0
+    
+    # RoPE
+    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
+    
+    # Misc
+    guidance_embeds: bool = True
 
 
 @dataclass  
 class TinyFluxDeepConfig:
-    """Expanded deep config"""
-    hidden_size: int = 768          # Same
-    num_attention_heads: int = 8    # 2 → 8 (6 new heads)
-    attention_head_dim: int = 128   # Same (so attention dim = 8*128 = 1024)
-    num_single_blocks: int = 25     # 3 → 25 (more singles like original Flux)
-    num_double_blocks: int = 15     # 3 → 15
-    mlp_ratio: float = 4.0          # Same
-    t5_embed_dim: int = 768         # Same
-    clip_embed_dim: int = 768       # Same
-    in_channels: int = 16           # Same
-    axes_dims: tuple = (16, 24, 24) # Same
-    theta: int = 10000              # Same
+    """
+    Expanded deep config - matches TinyFlux model attribute names exactly.
+    
+    Original TinyFlux: hidden_size=256, 2 heads (256/128=2)
+    Deep variant: hidden_size=512, 4 heads (4*128=512) - double heads
+    """
+    # Core dimensions
+    hidden_size: int = 512          # 4 heads * 128 head_dim
+    num_attention_heads: int = 4    # 2 → 4 (double the heads)
+    attention_head_dim: int = 128   # Same (required for RoPE)
+    
+    # Input/output
+    in_channels: int = 16
+    patch_size: int = 1
+    
+    # Text encoder interfaces
+    joint_attention_dim: int = 768   # T5 embed dim
+    pooled_projection_dim: int = 768  # CLIP embed dim
+    
+    # Layers (uses _layers not _blocks)
+    num_double_layers: int = 15     # 3 → 15
+    num_single_layers: int = 25     # 3 → 25 (more singles like original Flux)
+    
+    # MLP
+    mlp_ratio: float = 4.0
+    
+    # RoPE (must sum to head_dim=128)
+    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
+    
+    # Misc
+    guidance_embeds: bool = True
+    
+    def __post_init__(self):
+        assert self.num_attention_heads * self.attention_head_dim == self.hidden_size, \
+            f"heads ({self.num_attention_heads}) * head_dim ({self.attention_head_dim}) != hidden ({self.hidden_size})"
 
 
 # ============================================================================
@@ -84,71 +122,118 @@ DOUBLE_FROZEN = {0, 4, 7, 10, 14}  # These positions are frozen
 # ============================================================================
 # WEIGHT EXPANSION UTILITIES
 # ============================================================================
-def expand_qkv_weights(old_weight, old_heads=2, new_heads=8, head_dim=128):
+def expand_qkv_weights(old_weight, old_hidden=768, new_hidden=1536, head_dim=128):
     """
-    Expand QKV projection weights from 2 heads to 8 heads.
-    Old heads go to positions 0 and 7, middle heads initialized randomly.
+    Expand QKV projection weights when increasing hidden size / head count.
+    QKV weight shape: (3 * num_heads * head_dim, hidden_size) = (3 * hidden_size, hidden_size)
     
-    QKV weight shape: (in_features, 3 * num_heads * head_dim)
+    Strategy: Copy old weights to corresponding positions, random init new heads.
+    Old heads are spread evenly across new head positions.
     """
-    in_features = old_weight.shape[0]
-    old_qkv_dim = 3 * old_heads * head_dim  # 3 * 2 * 128 = 768
-    new_qkv_dim = 3 * new_heads * head_dim  # 3 * 8 * 128 = 3072
+    old_qkv_dim = old_weight.shape[0]  # 3 * old_hidden
+    new_qkv_dim = 3 * new_hidden
+    
+    old_heads = old_hidden // head_dim
+    new_heads = new_hidden // head_dim
     
     # Initialize new weights
-    new_weight = torch.zeros(in_features, new_qkv_dim, dtype=old_weight.dtype, device=old_weight.device)
-    # Small random init for new heads
+    new_weight = torch.zeros(new_qkv_dim, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
     nn.init.xavier_uniform_(new_weight)
-    new_weight *= 0.1  # Scale down random init
+    new_weight *= 0.02  # Scale down random init
     
-    # For each of Q, K, V
+    # For each of Q, K, V: copy old heads to first N positions
     for qkv_idx in range(3):
-        old_start = qkv_idx * old_heads * head_dim
-        new_start = qkv_idx * new_heads * head_dim
+        old_start = qkv_idx * old_hidden
+        new_start = qkv_idx * new_hidden
         
-        # Copy old head 0 → new head 0
-        old_h0_start = old_start
-        old_h0_end = old_start + head_dim
-        new_h0_start = new_start
-        new_h0_end = new_start + head_dim
-        new_weight[:, new_h0_start:new_h0_end] = old_weight[:, old_h0_start:old_h0_end]
-        
-        # Copy old head 1 → new head 7 (last)
-        old_h1_start = old_start + head_dim
-        old_h1_end = old_start + 2 * head_dim
-        new_h7_start = new_start + 7 * head_dim
-        new_h7_end = new_start + 8 * head_dim
-        new_weight[:, new_h7_start:new_h7_end] = old_weight[:, old_h1_start:old_h1_end]
+        # Copy all old heads to first old_heads positions of new
+        for h in range(old_heads):
+            old_h_start = old_start + h * head_dim
+            old_h_end = old_h_start + head_dim
+            new_h_start = new_start + h * head_dim
+            new_h_end = new_h_start + head_dim
+            # Copy weights, input dim goes to first old_hidden columns
+            new_weight[new_h_start:new_h_end, :old_hidden] = old_weight[old_h_start:old_h_end, :]
     
     return new_weight
 
 
-def expand_out_proj_weights(old_weight, old_heads=2, new_heads=8, head_dim=128):
-    """
-    Expand output projection weights from 2 heads to 8 heads.
+def expand_qkv_bias(old_bias, old_hidden=768, new_hidden=1536, head_dim=128):
+    """Expand QKV bias from old_hidden to new_hidden."""
+    new_qkv_dim = 3 * new_hidden
+    new_bias = torch.zeros(new_qkv_dim, dtype=old_bias.dtype, device=old_bias.device)
     
-    Out proj weight shape: (num_heads * head_dim, out_features)
-    """
-    out_features = old_weight.shape[1]
-    old_attn_dim = old_heads * head_dim  # 2 * 128 = 256
-    new_attn_dim = new_heads * head_dim  # 8 * 128 = 1024
+    old_heads = old_hidden // head_dim
+    
+    # Copy old biases to first old_heads positions for each of Q, K, V
+    for qkv_idx in range(3):
+        old_start = qkv_idx * old_hidden
+        new_start = qkv_idx * new_hidden
+        new_bias[new_start:new_start + old_hidden] = old_bias[old_start:old_start + old_hidden]
     
+    return new_bias
+
+
+def expand_out_proj_weights(old_weight, old_hidden=768, new_hidden=1536, head_dim=128):
+    """
+    Expand output projection weights.
+    Out proj weight shape: (hidden_size, num_heads * head_dim) = (hidden_size, hidden_size)
+    """
     # Initialize new weights
-    new_weight = torch.zeros(new_attn_dim, out_features, dtype=old_weight.dtype, device=old_weight.device)
+    new_weight = torch.zeros(new_hidden, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
     nn.init.xavier_uniform_(new_weight)
-    new_weight *= 0.1
+    new_weight *= 0.02
+    
+    # Copy old weights to top-left corner
+    new_weight[:old_hidden, :old_hidden] = old_weight
     
-    # Copy old head 0 → new head 0
-    new_weight[0:head_dim, :] = old_weight[0:head_dim, :]
+    return new_weight
+
+
+def expand_out_proj_bias(old_bias, old_hidden=768, new_hidden=1536):
+    """Expand output projection bias."""
+    new_bias = torch.zeros(new_hidden, dtype=old_bias.dtype, device=old_bias.device)
+    new_bias[:old_hidden] = old_bias
+    return new_bias
+
+
+def expand_linear_hidden(old_weight, old_hidden=768, new_hidden=1536, expand_in=True, expand_out=True):
+    """
+    Expand a linear layer weight from old_hidden to new_hidden.
+    """
+    old_out, old_in = old_weight.shape
+    
+    new_out = new_hidden if expand_out else old_out
+    new_in = new_hidden if expand_in else old_in
+    
+    new_weight = torch.zeros(new_out, new_in, dtype=old_weight.dtype, device=old_weight.device)
+    nn.init.xavier_uniform_(new_weight)
+    new_weight *= 0.02
     
-    # Copy old head 1 → new head 7
-    new_weight[7*head_dim:8*head_dim, :] = old_weight[head_dim:2*head_dim, :]
+    # Copy old weights to top-left corner
+    copy_out = old_hidden if expand_out else old_out
+    copy_in = old_hidden if expand_in else old_in
+    new_weight[:copy_out, :copy_in] = old_weight[:copy_out, :copy_in]
     
     return new_weight
 
 
-def port_single_block_weights(old_state, old_idx, new_state, new_idx, expand_heads=True):
-    """Port weights from old single block to new single block."""
+def expand_bias(old_bias, old_hidden=768, new_hidden=1536):
+    """Expand bias from old_hidden to new_hidden."""
+    new_bias = torch.zeros(new_hidden, dtype=old_bias.dtype, device=old_bias.device)
+    new_bias[:old_hidden] = old_bias
+    return new_bias
+
+
+def expand_norm(old_weight, old_hidden=768, new_hidden=1536):
+    """Expand RMSNorm weight from old_hidden to new_hidden."""
+    new_weight = torch.ones(new_hidden, dtype=old_weight.dtype, device=old_weight.device)
+    new_weight[:old_hidden] = old_weight
+    return new_weight
+
+
+def port_single_block_weights(old_state, old_idx, new_state, new_idx, old_hidden=256, new_hidden=1024):
+    """Port weights from old single block to new single block with dimension expansion."""
     old_prefix = f"single_blocks.{old_idx}"
     new_prefix = f"single_blocks.{new_idx}"
     
@@ -159,24 +244,119 @@ def port_single_block_weights(old_state, old_idx, new_state, new_idx, expand_hea
         new_key = old_key.replace(old_prefix, new_prefix)
         old_weight = old_state[old_key]
         
-        # Handle attention head expansion
-        if expand_heads:
-            if "attn.qkv.weight" in old_key:
-                new_state[new_key] = expand_qkv_weights(old_weight)
-                print(f"  Expanded QKV: {old_key} → {new_key}")
-                continue
-            elif "attn.out_proj.weight" in old_key:
-                new_state[new_key] = expand_out_proj_weights(old_weight)
-                print(f"  Expanded out_proj: {old_key} → {new_key}")
-                continue
+        # Attention QKV
+        if "attn.qkv.weight" in old_key:
+            new_state[new_key] = expand_qkv_weights(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded QKV weight: {old_key}")
+        elif "attn.qkv.bias" in old_key:
+            new_state[new_key] = expand_qkv_bias(old_weight)
+            print(f"  Expanded QKV bias: {old_key}")
+        
+        # Attention output projection
+        elif "attn.out_proj.weight" in old_key:
+            new_state[new_key] = expand_out_proj_weights(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded out_proj weight: {old_key}")
+        elif "attn.out_proj.bias" in old_key:
+            new_state[new_key] = expand_out_proj_bias(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded out_proj bias: {old_key}")
+        
+        # MLP layers (hidden → 4*hidden → hidden)
+        elif "mlp.fc1.weight" in old_key:
+            # fc1: hidden → 4*hidden
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_weight = torch.zeros(new_mlp_hidden, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_mlp_hidden, :old_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded MLP fc1 weight: {old_key}")
+        elif "mlp.fc1.bias" in old_key:
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_bias = torch.zeros(new_mlp_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            new_bias[:old_mlp_hidden] = old_weight
+            new_state[new_key] = new_bias
+            print(f"  Expanded MLP fc1 bias: {old_key}")
+        elif "mlp.fc2.weight" in old_key:
+            # fc2: 4*hidden → hidden
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_weight = torch.zeros(new_hidden, new_mlp_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_hidden, :old_mlp_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded MLP fc2 weight: {old_key}")
+        elif "mlp.fc2.bias" in old_key:
+            new_state[new_key] = expand_bias(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded MLP fc2 bias: {old_key}")
+        
+        # AdaLayerNorm modulation linear (norm.linear) - outputs 3*hidden for single blocks
+        elif "norm.linear.weight" in old_key:
+            # Shape: (3*old_hidden, old_hidden) → (3*new_hidden, new_hidden)
+            old_out = old_hidden * 3
+            new_out = new_hidden * 3
+            new_weight = torch.zeros(new_out, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_out, :old_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded AdaLN linear weight: {old_key} ({old_out},{old_hidden})→({new_out},{new_hidden})")
+        elif "norm.linear.bias" in old_key:
+            old_out = old_hidden * 3
+            new_out = new_hidden * 3
+            new_bias = torch.zeros(new_out, dtype=old_weight.dtype, device=old_weight.device)
+            new_bias[:old_out] = old_weight
+            new_state[new_key] = new_bias
+            print(f"  Expanded AdaLN linear bias: {old_key} ({old_out})→({new_out})")
+        
+        # RMSNorm inside AdaLN (norm.norm.weight) or standalone norm
+        elif "norm.norm.weight" in old_key or "norm2.weight" in old_key:
+            new_state[new_key] = expand_norm(old_weight, old_hidden, new_hidden)
+            print(f"  Expanded RMSNorm weight: {old_key}")
+        
+        # Generic normalization layers - check actual sizes
+        elif "norm" in old_key and "weight" in old_key:
+            old_size = old_weight.shape[0]
+            new_key_shape = new_state.get(new_key, torch.empty(0)).shape
+            if len(new_key_shape) > 0:
+                new_size = new_key_shape[0]
+                if old_size == new_size:
+                    new_state[new_key] = old_weight.clone()
+                    print(f"  Direct copy norm weight: {old_key} ({old_size})")
+                else:
+                    new_weight = torch.ones(new_size, dtype=old_weight.dtype, device=old_weight.device)
+                    copy_size = min(old_size, new_size)
+                    new_weight[:copy_size] = old_weight[:copy_size]
+                    new_state[new_key] = new_weight
+                    print(f"  Padded norm weight: {old_key} ({old_size}→{new_size})")
+        elif "norm" in old_key and "bias" in old_key:
+            old_size = old_weight.shape[0]
+            new_key_shape = new_state.get(new_key, torch.empty(0)).shape
+            if len(new_key_shape) > 0:
+                new_size = new_key_shape[0]
+                if old_size == new_size:
+                    new_state[new_key] = old_weight.clone()
+                    print(f"  Direct copy norm bias: {old_key} ({old_size})")
+                else:
+                    new_bias = torch.zeros(new_size, dtype=old_weight.dtype, device=old_weight.device)
+                    copy_size = min(old_size, new_size)
+                    new_bias[:copy_size] = old_weight[:copy_size]
+                    new_state[new_key] = new_bias
+                    print(f"  Padded norm bias: {old_key} ({old_size}→{new_size})")
         
-        # Direct copy for other weights
-        new_state[new_key] = old_weight.clone()
-        print(f"  Copied: {old_key} → {new_key}")
+        # Direct copy for anything else (shouldn't be much)
+        else:
+            if old_weight.shape == new_state.get(new_key, torch.empty(0)).shape:
+                new_state[new_key] = old_weight.clone()
+                print(f"  Direct copy: {old_key}")
+            else:
+                print(f"  SKIP (shape mismatch): {old_key}")
 
 
-def port_double_block_weights(old_state, old_idx, new_state, new_idx, expand_heads=True):
-    """Port weights from old double block to new double block."""
+def port_double_block_weights(old_state, old_idx, new_state, new_idx, old_hidden=256, new_hidden=1024):
+    """Port weights from old double block to new double block with dimension expansion."""
     old_prefix = f"double_blocks.{old_idx}"
     new_prefix = f"double_blocks.{new_idx}"
     
@@ -187,40 +367,211 @@ def port_double_block_weights(old_state, old_idx, new_state, new_idx, expand_hea
         new_key = old_key.replace(old_prefix, new_prefix)
         old_weight = old_state[old_key]
         
-        # Handle attention head expansion for joint attention
-        if expand_heads:
-            if any(x in old_key for x in ["img_qkv.weight", "txt_qkv.weight"]):
-                new_state[new_key] = expand_qkv_weights(old_weight)
-                print(f"  Expanded QKV: {old_key} → {new_key}")
-                continue
-            elif any(x in old_key for x in ["img_out.weight", "txt_out.weight"]):
-                new_state[new_key] = expand_out_proj_weights(old_weight)
-                print(f"  Expanded out_proj: {old_key} → {new_key}")
-                continue
+        # Joint attention QKV (img and txt)
+        if any(x in old_key for x in ["img_qkv.weight", "txt_qkv.weight"]):
+            new_state[new_key] = expand_qkv_weights(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded QKV weight: {old_key}")
+        elif any(x in old_key for x in ["img_qkv.bias", "txt_qkv.bias"]):
+            new_state[new_key] = expand_qkv_bias(old_weight)
+            print(f"  Expanded QKV bias: {old_key}")
+        
+        # Joint attention output projections
+        elif any(x in old_key for x in ["img_out.weight", "txt_out.weight"]):
+            new_state[new_key] = expand_out_proj_weights(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded out_proj weight: {old_key}")
+        elif any(x in old_key for x in ["img_out.bias", "txt_out.bias"]):
+            new_state[new_key] = expand_out_proj_bias(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded out_proj bias: {old_key}")
+        
+        # MLP layers
+        elif "mlp" in old_key and "fc1.weight" in old_key:
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_weight = torch.zeros(new_mlp_hidden, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_mlp_hidden, :old_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded MLP fc1 weight: {old_key}")
+        elif "mlp" in old_key and "fc1.bias" in old_key:
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_bias = torch.zeros(new_mlp_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            new_bias[:old_mlp_hidden] = old_weight
+            new_state[new_key] = new_bias
+            print(f"  Expanded MLP fc1 bias: {old_key}")
+        elif "mlp" in old_key and "fc2.weight" in old_key:
+            old_mlp_hidden = old_hidden * 4
+            new_mlp_hidden = new_hidden * 4
+            new_weight = torch.zeros(new_hidden, new_mlp_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_hidden, :old_mlp_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded MLP fc2 weight: {old_key}")
+        elif "mlp" in old_key and "fc2.bias" in old_key:
+            new_state[new_key] = expand_bias(old_weight, old_hidden=old_hidden, new_hidden=new_hidden)
+            print(f"  Expanded MLP fc2 bias: {old_key}")
+        
+        # AdaLayerNormZero modulation linear - outputs 6*hidden (img_norm1, txt_norm1)
+        elif ("img_norm1.linear" in old_key or "txt_norm1.linear" in old_key) and "weight" in old_key:
+            old_out = old_hidden * 6
+            new_out = new_hidden * 6
+            new_weight = torch.zeros(new_out, new_hidden, dtype=old_weight.dtype, device=old_weight.device)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_out, :old_hidden] = old_weight
+            new_state[new_key] = new_weight
+            print(f"  Expanded AdaLN linear weight: {old_key}")
+        elif ("img_norm1.linear" in old_key or "txt_norm1.linear" in old_key) and "bias" in old_key:
+            old_out = old_hidden * 6
+            new_out = new_hidden * 6
+            new_bias = torch.zeros(new_out, dtype=old_weight.dtype, device=old_weight.device)
+            new_bias[:old_out] = old_weight
+            new_state[new_key] = new_bias
+            print(f"  Expanded AdaLN linear bias: {old_key}")
+        
+        # RMSNorm inside AdaLN (img_norm1.norm, txt_norm1.norm) or standalone (img_norm2, txt_norm2)
+        elif any(x in old_key for x in ["_norm1.norm.weight", "_norm2.weight"]):
+            new_state[new_key] = expand_norm(old_weight, old_hidden, new_hidden)
+            print(f"  Expanded RMSNorm weight: {old_key}")
         
-        # Direct copy
-        new_state[new_key] = old_weight.clone()
-        print(f"  Copied: {old_key} → {new_key}")
+        # Generic normalization layers - check actual sizes
+        elif "norm" in old_key and "weight" in old_key:
+            old_size = old_weight.shape[0]
+            new_key_shape = new_state.get(new_key, torch.empty(0)).shape
+            if len(new_key_shape) > 0:
+                new_size = new_key_shape[0]
+                if old_size == new_size:
+                    new_state[new_key] = old_weight.clone()
+                    print(f"  Direct copy norm weight: {old_key} ({old_size})")
+                else:
+                    new_weight = torch.ones(new_size, dtype=old_weight.dtype, device=old_weight.device)
+                    copy_size = min(old_size, new_size)
+                    new_weight[:copy_size] = old_weight[:copy_size]
+                    new_state[new_key] = new_weight
+                    print(f"  Padded norm weight: {old_key} ({old_size}→{new_size})")
+        elif "norm" in old_key and "bias" in old_key:
+            old_size = old_weight.shape[0]
+            new_key_shape = new_state.get(new_key, torch.empty(0)).shape
+            if len(new_key_shape) > 0:
+                new_size = new_key_shape[0]
+                if old_size == new_size:
+                    new_state[new_key] = old_weight.clone()
+                    print(f"  Direct copy norm bias: {old_key} ({old_size})")
+                else:
+                    new_bias = torch.zeros(new_size, dtype=old_weight.dtype, device=old_weight.device)
+                    copy_size = min(old_size, new_size)
+                    new_bias[:copy_size] = old_weight[:copy_size]
+                    new_state[new_key] = new_bias
+                    print(f"  Padded norm bias: {old_key} ({old_size}→{new_size})")
+        
+        # Direct copy for matching shapes
+        else:
+            if old_weight.shape == new_state.get(new_key, torch.empty(0)).shape:
+                new_state[new_key] = old_weight.clone()
+                print(f"  Direct copy: {old_key}")
+            else:
+                print(f"  SKIP (shape mismatch): {old_key}")
 
 
-def port_non_block_weights(old_state, new_state, old_heads=2, new_heads=8):
-    """Port weights that aren't in single/double blocks."""
-    head_dim = 128
+def port_non_block_weights(old_state, new_state, old_hidden=256, new_hidden=1024):
+    """Port weights that aren't in single/double blocks with dimension expansion."""
     
     for old_key, old_weight in old_state.items():
         # Skip block weights (handled separately)
         if "single_blocks" in old_key or "double_blocks" in old_key:
             continue
         
-        # These can be copied directly (same dimensions)
-        direct_copy_keys = [
-            "img_in", "txt_in", "time_in", "vector_in", "guidance_in",
-            "final_norm", "final_linear", "rope"
-        ]
+        # Skip buffers that will be recomputed
+        if any(x in old_key for x in ["sin_basis", "freqs_"]):
+            print(f"  Skip buffer: {old_key}")
+            continue
+        
+        # img_in: in_channels → hidden
+        if "img_in.weight" in old_key:
+            new_weight = torch.zeros(new_hidden, old_weight.shape[1], dtype=old_weight.dtype)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_hidden, :] = old_weight
+            new_state[old_key] = new_weight
+            print(f"  Expanded: {old_key}")
+        elif "img_in.bias" in old_key:
+            new_state[old_key] = expand_bias(old_weight, old_hidden, new_hidden)
+            print(f"  Expanded: {old_key}")
+        
+        # txt_in: joint_attention_dim → hidden
+        elif "txt_in.weight" in old_key:
+            new_weight = torch.zeros(new_hidden, old_weight.shape[1], dtype=old_weight.dtype)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:old_hidden, :] = old_weight
+            new_state[old_key] = new_weight
+            print(f"  Expanded: {old_key}")
+        elif "txt_in.bias" in old_key:
+            new_state[old_key] = expand_bias(old_weight, old_hidden, new_hidden)
+            print(f"  Expanded: {old_key}")
         
-        if any(k in old_key for k in direct_copy_keys):
-            new_state[old_key] = old_weight.clone()
+        # time_in, guidance_in: MLPEmbedder (hidden → hidden)
+        elif any(x in old_key for x in ["time_in", "guidance_in"]):
+            if "fc1.weight" in old_key:
+                new_weight = torch.zeros(new_hidden, new_hidden, dtype=old_weight.dtype)
+                nn.init.xavier_uniform_(new_weight)
+                new_weight *= 0.02
+                new_weight[:old_hidden, :old_hidden] = old_weight
+                new_state[old_key] = new_weight
+                print(f"  Expanded: {old_key}")
+            elif "fc1.bias" in old_key:
+                new_state[old_key] = expand_bias(old_weight, old_hidden, new_hidden)
+                print(f"  Expanded: {old_key}")
+            elif "fc2.weight" in old_key:
+                new_weight = torch.zeros(new_hidden, new_hidden, dtype=old_weight.dtype)
+                nn.init.xavier_uniform_(new_weight)
+                new_weight *= 0.02
+                new_weight[:old_hidden, :old_hidden] = old_weight
+                new_state[old_key] = new_weight
+                print(f"  Expanded: {old_key}")
+            elif "fc2.bias" in old_key:
+                new_state[old_key] = expand_bias(old_weight, old_hidden, new_hidden)
+                print(f"  Expanded: {old_key}")
+        
+        # vector_in: pooled_projection_dim → hidden
+        elif "vector_in" in old_key:
+            if "weight" in old_key:
+                new_weight = torch.zeros(new_hidden, old_weight.shape[1], dtype=old_weight.dtype)
+                nn.init.xavier_uniform_(new_weight)
+                new_weight *= 0.02
+                new_weight[:old_hidden, :] = old_weight
+                new_state[old_key] = new_weight
+                print(f"  Expanded: {old_key}")
+            elif "bias" in old_key:
+                new_state[old_key] = expand_bias(old_weight, old_hidden, new_hidden)
+                print(f"  Expanded: {old_key}")
+        
+        # final_norm: RMSNorm(hidden)
+        elif "final_norm" in old_key:
+            if "weight" in old_key:
+                new_state[old_key] = expand_norm(old_weight, old_hidden, new_hidden)
+                print(f"  Expanded: {old_key}")
+        
+        # final_linear: hidden → in_channels
+        elif "final_linear.weight" in old_key:
+            new_weight = torch.zeros(old_weight.shape[0], new_hidden, dtype=old_weight.dtype)
+            nn.init.xavier_uniform_(new_weight)
+            new_weight *= 0.02
+            new_weight[:, :old_hidden] = old_weight
+            new_state[old_key] = new_weight
+            print(f"  Expanded: {old_key}")
+        elif "final_linear.bias" in old_key:
+            new_state[old_key] = old_weight.clone()  # output dim unchanged
             print(f"  Direct copy: {old_key}")
+        
+        # RoPE - skip, will be recomputed
+        elif "rope" in old_key:
+            print(f"  Skip RoPE: {old_key}")
+        
+        else:
+            print(f"  Unknown non-block key: {old_key}")
 
 
 # ============================================================================
@@ -247,14 +598,31 @@ def port_tinyflux_to_deep(old_weights_path, new_model):
         print("Stripping _orig_mod prefix...")
         old_state = {k.replace("_orig_mod.", ""): v for k, v in old_state.items()}
     
-    # Get new model's state dict as template
+    # Get new model's state dict as template FIRST
     new_state = new_model.state_dict()
     frozen_params = set()
     
+    # Auto-detect old hidden size from weights
+    if "final_norm.weight" in old_state:
+        old_hidden = old_state["final_norm.weight"].shape[0]
+    elif "img_in.weight" in old_state:
+        old_hidden = old_state["img_in.weight"].shape[0]
+    else:
+        old_hidden = 256  # Default for TinyFlux
+    
+    # Get new hidden size from new model's state dict
+    if "final_norm.weight" in new_state:
+        new_hidden = new_state["final_norm.weight"].shape[0]
+    else:
+        new_hidden = 512  # Default for TinyFlux-Deep
+    
+    print(f"Detected old hidden size: {old_hidden}")
+    print(f"New hidden size: {new_hidden}")
+    
     print("\n" + "="*60)
     print("Porting non-block weights...")
     print("="*60)
-    port_non_block_weights(old_state, new_state)
+    port_non_block_weights(old_state, new_state, old_hidden=old_hidden, new_hidden=new_hidden)
     
     print("\n" + "="*60)
     print("Porting single blocks (3 → 25)...")
@@ -262,7 +630,7 @@ def port_tinyflux_to_deep(old_weights_path, new_model):
     for old_idx, new_positions in SINGLE_MAPPING.items():
         for new_idx in new_positions:
             print(f"\nSingle block {old_idx} → {new_idx}:")
-            port_single_block_weights(old_state, old_idx, new_state, new_idx, expand_heads=True)
+            port_single_block_weights(old_state, old_idx, new_state, new_idx, old_hidden=old_hidden, new_hidden=new_hidden)
             # Mark as frozen
             for key in new_state.keys():
                 if f"single_blocks.{new_idx}." in key:
@@ -274,7 +642,7 @@ def port_tinyflux_to_deep(old_weights_path, new_model):
     for old_idx, new_positions in DOUBLE_MAPPING.items():
         for new_idx in new_positions:
             print(f"\nDouble block {old_idx} → {new_idx}:")
-            port_double_block_weights(old_state, old_idx, new_state, new_idx, expand_heads=True)
+            port_double_block_weights(old_state, old_idx, new_state, new_idx, old_hidden=old_hidden, new_hidden=new_hidden)
             # Mark as frozen
             for key in new_state.keys():
                 if f"double_blocks.{new_idx}." in key:
@@ -325,26 +693,50 @@ if __name__ == "__main__":
     print("TinyFlux → TinyFlux-Deep Porting")
     print("="*60)
     
-    # Load old weights from hub
+    # Load old weights from hub FIRST to detect dimensions
     print("\nDownloading TinyFlux weights from hub...")
     old_weights_path = hf_hub_download(
         repo_id="AbstractPhil/tiny-flux",
         filename="model.safetensors"
     )
     
-    # Create new deep model
-    print("\nCreating TinyFlux-Deep model...")
+    # Load and detect old dimensions
+    print("Detecting old model dimensions...")
+    old_state = load_file(old_weights_path)
+    if any(k.startswith("_orig_mod.") for k in old_state.keys()):
+        old_state = {k.replace("_orig_mod.", ""): v for k, v in old_state.items()}
+    
+    # Detect old hidden size
+    old_hidden = old_state["final_norm.weight"].shape[0]
+    head_dim = 128  # Fixed for RoPE
+    old_heads = old_hidden // head_dim
+    
+    print(f"  Old hidden size: {old_hidden}")
+    print(f"  Old attention heads: {old_heads}")
+    print(f"  Head dim: {head_dim}")
+    
+    # Calculate new dimensions (double the heads)
+    new_heads = old_heads * 2  # 6 → 12
+    new_hidden = new_heads * head_dim  # 12 * 128 = 1536
+    
+    print(f"\nNew dimensions:")
+    print(f"  New hidden size: {new_hidden}")
+    print(f"  New attention heads: {new_heads}")
+    
+    # Create deep config with detected dimensions
     deep_config = TinyFluxDeepConfig()
+    deep_config.hidden_size = new_hidden
+    deep_config.num_attention_heads = new_heads
     
+    print("\nCreating TinyFlux-Deep model...")
     # You need to define TinyFlux class first (run model cell)
-    # This assumes TinyFlux accepts the config
     deep_model = TinyFlux(deep_config).to(DTYPE)
     
     print(f"\nDeep model config:")
     print(f"  Hidden size: {deep_config.hidden_size}")
     print(f"  Attention heads: {deep_config.num_attention_heads}")
-    print(f"  Single blocks: {deep_config.num_single_blocks}")
-    print(f"  Double blocks: {deep_config.num_double_blocks}")
+    print(f"  Single layers: {deep_config.num_single_layers}")
+    print(f"  Double layers: {deep_config.num_double_layers}")
     
     # Port weights
     new_state, frozen_params = port_tinyflux_to_deep(old_weights_path, deep_model)
@@ -384,14 +776,14 @@ if __name__ == "__main__":
         "hidden_size": deep_config.hidden_size,
         "num_attention_heads": deep_config.num_attention_heads,
         "attention_head_dim": deep_config.attention_head_dim,
-        "num_single_blocks": deep_config.num_single_blocks,
-        "num_double_blocks": deep_config.num_double_blocks,
+        "num_single_layers": deep_config.num_single_layers,
+        "num_double_layers": deep_config.num_double_layers,
         "mlp_ratio": deep_config.mlp_ratio,
-        "t5_embed_dim": deep_config.t5_embed_dim,
-        "clip_embed_dim": deep_config.clip_embed_dim,
+        "joint_attention_dim": deep_config.joint_attention_dim,
+        "pooled_projection_dim": deep_config.pooled_projection_dim,
         "in_channels": deep_config.in_channels,
-        "axes_dims": list(deep_config.axes_dims),
-        "theta": deep_config.theta,
+        "axes_dims_rope": list(deep_config.axes_dims_rope),
+        "guidance_embeds": deep_config.guidance_embeds,
     }
     with open("config_deep.json", "w") as f:
         json.dump(config_dict, f, indent=2)