Update model.py

model.py

@@ -1,17 +1,21 @@
 """
 TinyFlux: A /12 scaled Flux architecture for experimentation.
+OPTIMIZED VERSION - Flash Attention, vectorized RoPE, caching
 
 Architecture:
   - hidden: 256 (3072/12)
-  - num_heads: 2 (24/12)  
+  - num_heads: 2 (24/12)
   - head_dim: 128 (preserved for RoPE compatibility)
   - in_channels: 16 (Flux VAE output channels)
   - double_layers: 3
   - single_layers: 3
-  
-Text Encoders (runtime):
-  - flan-t5-base (768 dim) → txt_in projects to hidden
-  - CLIP-L (768 dim pooled) → vector_in projects to hidden
+
+Optimizations:
+  - Flash Attention (F.scaled_dot_product_attention)
+  - Vectorized RoPE with precomputed frequencies
+  - Vectorized img_ids creation (no Python loops)
+  - Caching for img_ids and RoPE embeddings
+  - Precomputed sinusoidal embeddings
 """
 
 import torch
@@ -19,7 +23,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import math
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Dict
 
 
 @dataclass
@@ -29,28 +33,28 @@ class TinyFluxConfig:
     hidden_size: int = 256
     num_attention_heads: int = 2
     attention_head_dim: int = 128  # Preserved for RoPE
-    
+
     # Input/output (Flux VAE has 16 channels)
-    in_channels: int = 16  # Flux VAE output channels
-    patch_size: int = 1    # No 2x2 patchification, raw latent tokens
-    
-    # Text encoder interfaces (runtime encoding)
-    joint_attention_dim: int = 768   # flan-t5-base output dim
+    in_channels: int = 16
+    patch_size: int = 1
+
+    # Text encoder interfaces
+    joint_attention_dim: int = 768  # flan-t5-base output dim
     pooled_projection_dim: int = 768  # CLIP-L pooled dim
-    
+
     # Layers
     num_double_layers: int = 3
     num_single_layers: int = 3
-    
+
     # MLP
     mlp_ratio: float = 4.0
-    
+
     # RoPE (must sum to head_dim)
     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})"
@@ -60,6 +64,7 @@ class TinyFluxConfig:
 
 class RMSNorm(nn.Module):
     """Root Mean Square Layer Normalization."""
+
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.eps = eps
@@ -71,43 +76,43 @@ class RMSNorm(nn.Module):
 
 
 class RotaryEmbedding(nn.Module):
-    """Rotary Position Embedding for 2D + temporal."""
+    """Rotary Position Embedding - OPTIMIZED with precomputed frequencies."""
+
     def __init__(self, dim: int, axes_dims: Tuple[int, int, int], theta: float = 10000.0):
         super().__init__()
         self.dim = dim
-        self.axes_dims = axes_dims  # (temporal, height, width)
+        self.axes_dims = axes_dims
         self.theta = theta
-        
+
+        # Precompute frequencies for each axis (no loop at runtime)
+        for i, axis_dim in enumerate(axes_dims):
+            freqs = 1.0 / (theta ** (torch.arange(0, axis_dim, 2).float() / axis_dim))
+            self.register_buffer(f'freqs_{i}', freqs)
+
     def forward(self, ids: torch.Tensor, dtype: torch.dtype = None) -> torch.Tensor:
         """
         ids: (B, N, 3) - temporal, height, width indices
-        dtype: output dtype (defaults to ids.dtype, but use model dtype for bf16)
         Returns: (B, N, dim) rotary embeddings
         """
         B, N, _ = ids.shape
-        device = ids.device
-        # Compute in float32 for precision, cast at the end
-        compute_dtype = torch.float32
         output_dtype = dtype if dtype is not None else ids.dtype
-        
-        embeddings = []
-        dim_offset = 0
-        
-        for axis_idx, axis_dim in enumerate(self.axes_dims):
-            # Compute frequencies for this axis
-            freqs = 1.0 / (self.theta ** (torch.arange(0, axis_dim, 2, device=device, dtype=compute_dtype) / axis_dim))
-            # Get positions for this axis
-            positions = ids[:, :, axis_idx].to(compute_dtype)  # (B, N)
-            # Outer product: (B, N) x (axis_dim/2) -> (B, N, axis_dim/2)
-            angles = positions.unsqueeze(-1) * freqs.unsqueeze(0).unsqueeze(0)
-            # Interleave sin/cos
-            emb = torch.stack([angles.cos(), angles.sin()], dim=-1)  # (B, N, axis_dim/2, 2)
-            emb = emb.flatten(-2)  # (B, N, axis_dim)
-            embeddings.append(emb)
-            dim_offset += axis_dim
-            
-        result = torch.cat(embeddings, dim=-1)  # (B, N, dim)
-        return result.to(output_dtype)
+
+        # Extract positions for each axis
+        pos0 = ids[:, :, 0:1].float()  # (B, N, 1)
+        pos1 = ids[:, :, 1:2].float()
+        pos2 = ids[:, :, 2:3].float()
+
+        # Compute angles (broadcasting: (B, N, 1) * (axis_dim/2,) -> (B, N, axis_dim/2))
+        angles0 = pos0 * self.freqs_0
+        angles1 = pos1 * self.freqs_1
+        angles2 = pos2 * self.freqs_2
+
+        # Stack sin/cos and flatten for each axis
+        emb0 = torch.stack([angles0.cos(), angles0.sin()], dim=-1).flatten(-2)
+        emb1 = torch.stack([angles1.cos(), angles1.sin()], dim=-1).flatten(-2)
+        emb2 = torch.stack([angles2.cos(), angles2.sin()], dim=-1).flatten(-2)
+
+        return torch.cat([emb0, emb1, emb2], dim=-1).to(output_dtype)
 
 
 def apply_rope(x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor:
@@ -115,28 +120,28 @@ def apply_rope(x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor:
     # x: (B, heads, N, head_dim)
     # rope: (B, N, head_dim)
     B, H, N, D = x.shape
-    
-    # Ensure rope matches x dtype
+
     rope = rope.to(x.dtype).unsqueeze(1)  # (B, 1, N, D)
-    
+
     # Split into pairs
     x_pairs = x.reshape(B, H, N, D // 2, 2)
     rope_pairs = rope.reshape(B, 1, N, D // 2, 2)
-    
+
     cos = rope_pairs[..., 0]
     sin = rope_pairs[..., 1]
-    
+
     x0 = x_pairs[..., 0]
     x1 = x_pairs[..., 1]
-    
+
     out0 = x0 * cos - x1 * sin
     out1 = x1 * cos + x0 * sin
-    
+
     return torch.stack([out0, out1], dim=-1).flatten(-2)
 
 
 class MLPEmbedder(nn.Module):
-    """MLP for embedding scalars (timestep, guidance)."""
+    """MLP for embedding scalars - OPTIMIZED with precomputed basis."""
+
     def __init__(self, hidden_size: int):
         super().__init__()
         self.mlp = nn.Sequential(
@@ -144,38 +149,31 @@ class MLPEmbedder(nn.Module):
             nn.SiLU(),
             nn.Linear(hidden_size, hidden_size),
         )
-        
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Sinusoidal embedding first
+        # Precompute sinusoidal basis
         half_dim = 128
         emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=x.device, dtype=x.dtype) * -emb)
-        emb = x.unsqueeze(-1) * emb.unsqueeze(0)
-        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)  # (B, 256)
+        emb = torch.exp(torch.arange(half_dim) * -emb)
+        self.register_buffer('sin_basis', emb)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Use precomputed basis
+        emb = x.unsqueeze(-1) * self.sin_basis.to(x.dtype)
+        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
         return self.mlp(emb)
 
 
 class AdaLayerNormZero(nn.Module):
-    """
-    AdaLN-Zero for double-stream blocks.
-    Outputs 6 modulation params: (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp)
-    """
+    """AdaLN-Zero for double-stream blocks."""
+
     def __init__(self, hidden_size: int):
         super().__init__()
         self.silu = nn.SiLU()
         self.linear = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
         self.norm = RMSNorm(hidden_size)
-        
+
     def forward(
-        self, x: torch.Tensor, emb: torch.Tensor
+            self, x: torch.Tensor, emb: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            x: hidden states (B, N, D)
-            emb: conditioning embedding (B, D)
-        Returns:
-            (normed_x, gate_msa, shift_mlp, scale_mlp, gate_mlp)
-        """
         emb_out = self.linear(self.silu(emb))
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb_out.chunk(6, dim=-1)
         x = self.norm(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
@@ -183,26 +181,17 @@ class AdaLayerNormZero(nn.Module):
 
 
 class AdaLayerNormZeroSingle(nn.Module):
-    """
-    AdaLN-Zero for single-stream blocks.
-    Outputs 3 modulation params: (shift, scale, gate)
-    """
+    """AdaLN-Zero for single-stream blocks."""
+
     def __init__(self, hidden_size: int):
         super().__init__()
         self.silu = nn.SiLU()
         self.linear = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
         self.norm = RMSNorm(hidden_size)
-        
+
     def forward(
-        self, x: torch.Tensor, emb: torch.Tensor
+            self, x: torch.Tensor, emb: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Args:
-            x: hidden states (B, N, D)
-            emb: conditioning embedding (B, D)
-        Returns:
-            (normed_x, gate)
-        """
         emb_out = self.linear(self.silu(emb))
         shift, scale, gate = emb_out.chunk(3, dim=-1)
         x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
@@ -210,256 +199,212 @@ class AdaLayerNormZeroSingle(nn.Module):
 
 
 class Attention(nn.Module):
-    """Multi-head attention with optional RoPE."""
+    """Multi-head attention - OPTIMIZED with Flash Attention."""
+
     def __init__(self, hidden_size: int, num_heads: int, head_dim: int):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.scale = head_dim ** -0.5
-        
+
         self.qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=False)
         self.out_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
-        
+
     def forward(
-        self, 
-        x: torch.Tensor, 
-        rope: Optional[torch.Tensor] = None,
-        mask: Optional[torch.Tensor] = None
+            self,
+            x: torch.Tensor,
+            rope: Optional[torch.Tensor] = None,
+            mask: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
         B, N, _ = x.shape
         dtype = x.dtype
-        
-        # Ensure RoPE matches input dtype
+
         if rope is not None:
             rope = rope.to(dtype)
-        
+
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
         q, k, v = qkv.permute(2, 0, 3, 1, 4)  # 3 x (B, heads, N, head_dim)
-        
+
         if rope is not None:
             q = apply_rope(q, rope)
             k = apply_rope(k, rope)
-        
-        # Scaled dot-product attention
-        attn = (q @ k.transpose(-2, -1)) * self.scale
-        if mask is not None:
-            attn = attn.masked_fill(mask == 0, float('-inf'))
-        attn = attn.softmax(dim=-1)
-        
-        out = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+
+        # Flash Attention - faster and memory efficient
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=self.scale)
+        out = out.transpose(1, 2).reshape(B, N, -1)
         return self.out_proj(out)
 
 
 class JointAttention(nn.Module):
-    """Joint attention for double-stream blocks (separate Q,K,V for txt and img)."""
+    """Joint attention - OPTIMIZED with Flash Attention."""
+
     def __init__(self, hidden_size: int, num_heads: int, head_dim: int):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.scale = head_dim ** -0.5
-        
-        # Separate projections for text and image
+
         self.txt_qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=False)
         self.img_qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=False)
-        
+
         self.txt_out = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
         self.img_out = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
-        
+
     def forward(
-        self,
-        txt: torch.Tensor,
-        img: torch.Tensor,
-        rope: Optional[torch.Tensor] = None,
+            self,
+            txt: torch.Tensor,
+            img: torch.Tensor,
+            rope: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         B, L, _ = txt.shape
         _, N, _ = img.shape
-        
-        # Ensure consistent dtype (use img dtype as reference)
+
         dtype = img.dtype
         txt = txt.to(dtype)
         if rope is not None:
             rope = rope.to(dtype)
-        
+
         # Compute Q, K, V for both streams
         txt_qkv = self.txt_qkv(txt).reshape(B, L, 3, self.num_heads, self.head_dim)
         img_qkv = self.img_qkv(img).reshape(B, N, 3, self.num_heads, self.head_dim)
-        
+
         txt_q, txt_k, txt_v = txt_qkv.permute(2, 0, 3, 1, 4)
         img_q, img_k, img_v = img_qkv.permute(2, 0, 3, 1, 4)
-        
-        # Apply RoPE to image queries/keys only (text doesn't have positions)
+
+        # Apply RoPE to image only
         if rope is not None:
             img_q = apply_rope(img_q, rope)
             img_k = apply_rope(img_k, rope)
-        
-        # Concatenate keys and values for joint attention
-        k = torch.cat([txt_k, img_k], dim=2)  # (B, heads, L+N, head_dim)
+
+        # Concatenate for joint attention
+        k = torch.cat([txt_k, img_k], dim=2)
         v = torch.cat([txt_v, img_v], dim=2)
-        
-        # Text attends to all
-        txt_attn = (txt_q @ k.transpose(-2, -1)) * self.scale
-        txt_attn = txt_attn.softmax(dim=-1)
-        txt_out = (txt_attn @ v).transpose(1, 2).reshape(B, L, -1)
-        
-        # Image attends to all
-        img_attn = (img_q @ k.transpose(-2, -1)) * self.scale
-        img_attn = img_attn.softmax(dim=-1)
-        img_out = (img_attn @ v).transpose(1, 2).reshape(B, N, -1)
-        
+
+        # Flash Attention for both streams
+        txt_out = F.scaled_dot_product_attention(txt_q, k, v, scale=self.scale)
+        img_out = F.scaled_dot_product_attention(img_q, k, v, scale=self.scale)
+
+        txt_out = txt_out.transpose(1, 2).reshape(B, L, -1)
+        img_out = img_out.transpose(1, 2).reshape(B, N, -1)
+
         return self.txt_out(txt_out), self.img_out(img_out)
 
 
 class MLP(nn.Module):
     """Feed-forward network."""
+
     def __init__(self, hidden_size: int, mlp_ratio: float = 4.0):
         super().__init__()
         mlp_hidden = int(hidden_size * mlp_ratio)
         self.fc1 = nn.Linear(hidden_size, mlp_hidden)
         self.act = nn.GELU(approximate='tanh')
         self.fc2 = nn.Linear(mlp_hidden, hidden_size)
-        
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.fc2(self.act(self.fc1(x)))
 
 
 class DoubleStreamBlock(nn.Module):
-    """
-    Double-stream transformer block (MMDiT style).
-    Text and image have separate weights but attend to each other.
-    Uses AdaLN-Zero with 6 modulation params per stream.
-    """
+    """Double-stream transformer block (MMDiT style)."""
+
     def __init__(self, config: TinyFluxConfig):
         super().__init__()
         hidden = config.hidden_size
         heads = config.num_attention_heads
         head_dim = config.attention_head_dim
-        mlp_hidden = int(hidden * config.mlp_ratio)
-        
-        # AdaLN-Zero for each stream (outputs 6 params each)
+
         self.img_norm1 = AdaLayerNormZero(hidden)
         self.txt_norm1 = AdaLayerNormZero(hidden)
-        
-        # Joint attention (separate QKV projections)
         self.attn = JointAttention(hidden, heads, head_dim)
-        
-        # Second norm for MLP (not adaptive, uses params from norm1)
         self.img_norm2 = RMSNorm(hidden)
         self.txt_norm2 = RMSNorm(hidden)
-        
-        # MLPs
         self.img_mlp = MLP(hidden, config.mlp_ratio)
         self.txt_mlp = MLP(hidden, config.mlp_ratio)
-        
+
     def forward(
-        self,
-        txt: torch.Tensor,
-        img: torch.Tensor,
-        vec: torch.Tensor,
-        rope: Optional[torch.Tensor] = None,
+            self,
+            txt: torch.Tensor,
+            img: torch.Tensor,
+            vec: torch.Tensor,
+            rope: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # Image stream: norm + modulation, get MLP params for later
         img_normed, img_gate_msa, img_shift_mlp, img_scale_mlp, img_gate_mlp = self.img_norm1(img, vec)
-        
-        # Text stream: norm + modulation, get MLP params for later
         txt_normed, txt_gate_msa, txt_shift_mlp, txt_scale_mlp, txt_gate_mlp = self.txt_norm1(txt, vec)
-        
-        # Joint attention
+
         txt_attn_out, img_attn_out = self.attn(txt_normed, img_normed, rope)
-        
-        # Residual with gate
+
         txt = txt + txt_gate_msa.unsqueeze(1) * txt_attn_out
         img = img + img_gate_msa.unsqueeze(1) * img_attn_out
-        
-        # MLP with modulation (using params from norm1)
+
         txt_mlp_in = self.txt_norm2(txt) * (1 + txt_scale_mlp.unsqueeze(1)) + txt_shift_mlp.unsqueeze(1)
         img_mlp_in = self.img_norm2(img) * (1 + img_scale_mlp.unsqueeze(1)) + img_shift_mlp.unsqueeze(1)
-        
+
         txt = txt + txt_gate_mlp.unsqueeze(1) * self.txt_mlp(txt_mlp_in)
         img = img + img_gate_mlp.unsqueeze(1) * self.img_mlp(img_mlp_in)
-        
+
         return txt, img
 
 
 class SingleStreamBlock(nn.Module):
-    """
-    Single-stream transformer block.
-    Text and image are concatenated and share weights.
-    Uses AdaLN-Zero with 3 modulation params (no separate MLP modulation).
-    """
+    """Single-stream transformer block."""
+
     def __init__(self, config: TinyFluxConfig):
         super().__init__()
         hidden = config.hidden_size
         heads = config.num_attention_heads
         head_dim = config.attention_head_dim
-        mlp_hidden = int(hidden * config.mlp_ratio)
-        
-        # AdaLN-Zero (outputs 3 params: shift, scale, gate)
+
         self.norm = AdaLayerNormZeroSingle(hidden)
-        
-        # Combined QKV + MLP projection (Flux fuses these)
-        # Linear attention: QKV projection
         self.attn = Attention(hidden, heads, head_dim)
-        
-        # MLP
         self.mlp = MLP(hidden, config.mlp_ratio)
-        
-        # Pre-MLP norm (not modulated in single-stream)
         self.norm2 = RMSNorm(hidden)
-        
+
     def forward(
-        self,
-        txt: torch.Tensor,
-        img: torch.Tensor,
-        vec: torch.Tensor,
-        txt_rope: Optional[torch.Tensor] = None,
-        img_rope: Optional[torch.Tensor] = None,
+            self,
+            txt: torch.Tensor,
+            img: torch.Tensor,
+            vec: torch.Tensor,
+            txt_rope: Optional[torch.Tensor] = None,
+            img_rope: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         L = txt.shape[1]
-        
-        # Concatenate txt and img
+
         x = torch.cat([txt, img], dim=1)
-        
-        # Concatenate RoPE (zeros for text positions)
+
         if img_rope is not None:
             B, N, D = img_rope.shape
             txt_rope_zeros = torch.zeros(B, L, D, device=img_rope.device, dtype=img_rope.dtype)
             rope = torch.cat([txt_rope_zeros, img_rope], dim=1)
         else:
             rope = None
-        
-        # Norm + modulation (only 3 params for single stream)
+
         x_normed, gate = self.norm(x, vec)
-        
-        # Attention with gated residual
         x = x + gate.unsqueeze(1) * self.attn(x_normed, rope)
-        
-        # MLP (no separate modulation in single-stream Flux)
         x = x + self.mlp(self.norm2(x))
-        
-        # Split back
+
         txt, img = x.split([L, x.shape[1] - L], dim=1)
         return txt, img
 
 
+# Global cache for img_ids (they don't change for same resolution)
+_IMG_IDS_CACHE: Dict[Tuple, torch.Tensor] = {}
+
+
 class TinyFlux(nn.Module):
     """
     TinyFlux: A scaled-down Flux diffusion transformer.
-    
-    Scaling: /12 from original Flux
-      - hidden: 3072 → 256
-      - heads: 24 → 2
-      - head_dim: 128 (preserved)
-      - in_channels: 16 (Flux VAE)
+    OPTIMIZED with Flash Attention, vectorized ops, and caching.
     """
+
     def __init__(self, config: Optional[TinyFluxConfig] = None):
         super().__init__()
         self.config = config or TinyFluxConfig()
         cfg = self.config
-        
+
         # Input projections
         self.img_in = nn.Linear(cfg.in_channels, cfg.hidden_size)
         self.txt_in = nn.Linear(cfg.joint_attention_dim, cfg.hidden_size)
-        
+
         # Conditioning projections
         self.time_in = MLPEmbedder(cfg.hidden_size)
         self.vector_in = nn.Sequential(
@@ -468,10 +413,10 @@ class TinyFlux(nn.Module):
         )
         if cfg.guidance_embeds:
             self.guidance_in = MLPEmbedder(cfg.hidden_size)
-        
+
         # RoPE
         self.rope = RotaryEmbedding(cfg.attention_head_dim, cfg.axes_dims_rope)
-        
+
         # Transformer blocks
         self.double_blocks = nn.ModuleList([
             DoubleStreamBlock(cfg) for _ in range(cfg.num_double_layers)
@@ -479,13 +424,16 @@ class TinyFlux(nn.Module):
         self.single_blocks = nn.ModuleList([
             SingleStreamBlock(cfg) for _ in range(cfg.num_single_layers)
         ])
-        
+
         # Output
         self.final_norm = RMSNorm(cfg.hidden_size)
         self.final_linear = nn.Linear(cfg.hidden_size, cfg.in_channels)
-        
+
+        # RoPE cache
+        self._rope_cache: Dict[Tuple, torch.Tensor] = {}
+
         self._init_weights()
-        
+
     def _init_weights(self):
         """Initialize weights."""
         def _init(module):
@@ -493,68 +441,78 @@ class TinyFlux(nn.Module):
                 nn.init.xavier_uniform_(module.weight)
                 if module.bias is not None:
                     nn.init.zeros_(module.bias)
+
         self.apply(_init)
-        
-        # Zero-init output projection for residual
         nn.init.zeros_(self.final_linear.weight)
-        
+
     def forward(
-        self,
-        hidden_states: torch.Tensor,       # (B, N, in_channels) - image patches
-        encoder_hidden_states: torch.Tensor,  # (B, L, joint_attention_dim) - T5 tokens
-        pooled_projections: torch.Tensor,  # (B, pooled_projection_dim) - CLIP pooled
-        timestep: torch.Tensor,            # (B,) - diffusion timestep
-        img_ids: torch.Tensor,             # (B, N, 3) - image position ids
-        guidance: Optional[torch.Tensor] = None,  # (B,) - guidance scale
+            self,
+            hidden_states: torch.Tensor,
+            encoder_hidden_states: torch.Tensor,
+            pooled_projections: torch.Tensor,
+            timestep: torch.Tensor,
+            img_ids: torch.Tensor,
+            guidance: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """
-        Forward pass.
-        
-        Returns:
-            Predicted noise/velocity of shape (B, N, in_channels)
-        """
+        """Forward pass."""
         # Input projections
-        img = self.img_in(hidden_states)  # (B, N, hidden)
-        txt = self.txt_in(encoder_hidden_states)  # (B, L, hidden)
-        
+        img = self.img_in(hidden_states)
+        txt = self.txt_in(encoder_hidden_states)
+
         # Conditioning vector
         vec = self.time_in(timestep)
         vec = vec + self.vector_in(pooled_projections)
         if self.config.guidance_embeds and guidance is not None:
             vec = vec + self.guidance_in(guidance)
-        
-        # RoPE for image positions (match model dtype)
+
+        # RoPE for image positions
         img_rope = self.rope(img_ids, dtype=img.dtype)
-        
+
         # Double-stream blocks
         for block in self.double_blocks:
             txt, img = block(txt, img, vec, img_rope)
-        
+
         # Single-stream blocks
         for block in self.single_blocks:
             txt, img = block(txt, img, vec, img_rope=img_rope)
-        
-        # Output (image only)
+
+        # Output
         img = self.final_norm(img)
         img = self.final_linear(img)
-        
+
         return img
-    
+
     @staticmethod
     def create_img_ids(batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
-        """Create image position IDs for RoPE."""
-        # height, width are in latent space (image_size / 8)
-        img_ids = torch.zeros(batch_size, height * width, 3, device=device)
+        """Create image position IDs - VECTORIZED (no Python loops)."""
+        global _IMG_IDS_CACHE
+        
+        # Check cache first
+        cache_key = (batch_size, height, width, device)
+        if cache_key in _IMG_IDS_CACHE:
+            return _IMG_IDS_CACHE[cache_key]
+
+        # Vectorized creation using meshgrid
+        h_ids = torch.arange(height, device=device, dtype=torch.float32)
+        w_ids = torch.arange(width, device=device, dtype=torch.float32)
+
+        grid_h, grid_w = torch.meshgrid(h_ids, w_ids, indexing='ij')
+
+        # Stack: (H*W, 3) with [temporal=0, height, width]
+        img_ids = torch.stack([
+            torch.zeros(height * width, device=device),  # temporal
+            grid_h.flatten(),
+            grid_w.flatten(),
+        ], dim=-1)
+
+        # Expand for batch
+        img_ids = img_ids.unsqueeze(0).expand(batch_size, -1, -1)
         
-        for i in range(height):
-            for j in range(width):
-                idx = i * width + j
-                img_ids[:, idx, 0] = 0  # temporal (always 0 for images)
-                img_ids[:, idx, 1] = i  # height
-                img_ids[:, idx, 2] = j  # width
-                
+        # Cache it
+        _IMG_IDS_CACHE[cache_key] = img_ids
+
         return img_ids
-    
+
     def count_parameters(self) -> dict:
         """Count parameters by component."""
         counts = {}
@@ -567,72 +525,61 @@ class TinyFlux(nn.Module):
         counts['double_blocks'] = sum(p.numel() for p in self.double_blocks.parameters())
         counts['single_blocks'] = sum(p.numel() for p in self.single_blocks.parameters())
         counts['final'] = sum(p.numel() for p in self.final_norm.parameters()) + \
-                         sum(p.numel() for p in self.final_linear.parameters())
+                          sum(p.numel() for p in self.final_linear.parameters())
         counts['total'] = sum(p.numel() for p in self.parameters())
         return counts
 
 
 def test_tiny_flux():
-    """Quick test of the model."""
+    """Quick test of the optimized model."""
     print("=" * 60)
-    print("TinyFlux Model Test")
+    print("TinyFlux OPTIMIZED Model Test")
     print("=" * 60)
-    
+
     config = TinyFluxConfig()
     print(f"\nConfig:")
     print(f"  hidden_size: {config.hidden_size}")
     print(f"  num_heads: {config.num_attention_heads}")
     print(f"  head_dim: {config.attention_head_dim}")
-    print(f"  in_channels: {config.in_channels}")
-    print(f"  double_layers: {config.num_double_layers}")
-    print(f"  single_layers: {config.num_single_layers}")
-    print(f"  joint_attention_dim: {config.joint_attention_dim}")
-    print(f"  pooled_projection_dim: {config.pooled_projection_dim}")
-    
+
     model = TinyFlux(config)
-    
-    # Count parameters
+
     counts = model.count_parameters()
-    print(f"\nParameters:")
-    for name, count in counts.items():
-        print(f"  {name}: {count:,} ({count/1e6:.2f}M)")
-    
-    # Test forward pass
+    print(f"\nParameters: {counts['total']:,} ({counts['total'] / 1e6:.2f}M)")
+
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     model = model.to(device)
-    
-    batch_size = 2
-    latent_h, latent_w = 64, 64  # 512x512 image / 8
+
+    batch_size = 4
+    latent_h, latent_w = 64, 64
     num_patches = latent_h * latent_w
     text_len = 77
-    
-    # Create dummy inputs
+
     hidden_states = torch.randn(batch_size, num_patches, config.in_channels, device=device)
     encoder_hidden_states = torch.randn(batch_size, text_len, config.joint_attention_dim, device=device)
     pooled_projections = torch.randn(batch_size, config.pooled_projection_dim, device=device)
     timestep = torch.rand(batch_size, device=device)
     img_ids = TinyFlux.create_img_ids(batch_size, latent_h, latent_w, device)
     guidance = torch.ones(batch_size, device=device) * 3.5
-    
-    print(f"\nInput shapes:")
-    print(f"  hidden_states: {hidden_states.shape}")
-    print(f"  encoder_hidden_states: {encoder_hidden_states.shape}")
-    print(f"  pooled_projections: {pooled_projections.shape}")
-    print(f"  img_ids: {img_ids.shape}")
-    
-    # Forward pass
+
+    # Warmup
     with torch.no_grad():
-        output = model(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            pooled_projections=pooled_projections,
-            timestep=timestep,
-            img_ids=img_ids,
-            guidance=guidance,
-        )
-    
-    print(f"\nOutput shape: {output.shape}")
-    print(f"Output range: [{output.min():.4f}, {output.max():.4f}]")
+        for _ in range(3):
+            _ = model(hidden_states, encoder_hidden_states, pooled_projections, timestep, img_ids, guidance)
+
+    # Benchmark
+    if device == 'cuda':
+        torch.cuda.synchronize()
+        import time
+        start = time.time()
+        with torch.no_grad():
+            for _ in range(10):
+                output = model(hidden_states, encoder_hidden_states, pooled_projections, timestep, img_ids, guidance)
+        torch.cuda.synchronize()
+        elapsed = (time.time() - start) / 10
+        print(f"\nAverage forward pass: {elapsed*1000:.2f}ms")
+
+    print(f"Output shape: {output.shape}")
     print("\n✓ Forward pass successful!")