Create colab_inference_lailah_early.py

colab_inference_lailah_early.py

@@ -0,0 +1,1057 @@
+"""
+    TinyFlux-Lailah Inference
+    Loads the model code, the weights, and runs the inference based on the settings below.
+    Set up with only EULER for now.
+
+    No guarantees for any of this to work.
+
+    It's pretty bad in it's current phases, just check on it later if you're interested.
+    LICENSE: MIT
+"""
+
+
+POSITIVE_PROMPT = "woman" # @param {type:"string"}
+NEGATIVE_PROMPT = "" # @param {type:"string"}
+STEPS = 50 # @param {type:"integer"}
+CFG_GUIDANCE = 5 # @param {type: "number"}
+FLUX_SHIFT = 3 # @param {type: "number"}
+SEED = 420 # @param {type: "integer"}
+OUTPUT_PATH = "output.png" # @param {type:"string"}
+WIDTH = 512 # @param {type: "integer"}
+HEIGHT = 512 # @param {type: "integer"}
+
+# Model loading
+HF_REPO = "AbstractPhil/tiny-flux-deep"  # @param {type:"string"}
+# "hub", "hub:step_XXXXX", "local:/path/to/weights.safetensors"
+LOAD_FROM = "hub:step_293750"  # @param {type:"string"} 
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+
+#@title Preview (updates in-place)
+from IPython.display import display, Image as IPyImage, update_display
+from PIL import Image as PIL
+import numpy as np, io
+
+_PREVIEW_DISPLAY_ID = "tf_preview"
+
+preview_size = min(512, max(WIDTH, HEIGHT) // 2)
+
+
+def _pil_to_png_bytes(img: PIL) -> bytes:
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+def init_preview(square: int = 256):
+    """Show a black placeholder square once."""
+    black = PIL.fromarray(np.zeros((square, square, 3), dtype=np.uint8))
+    display(IPyImage(data=_pil_to_png_bytes(black)), display_id=_PREVIEW_DISPLAY_ID)
+
+def set_preview_from_pil(img: PIL, square: int = 256):
+    """Update the preview in-place with a PIL image."""
+    im = img.convert("RGB").copy()
+    im.thumbnail((square, square), resample=PIL.Resampling.LANCZOS)
+    # pad to square (so it stays a square widget)
+    canvas = PIL.fromarray(np.zeros((square, square, 3), dtype=np.uint8))
+    x = (square - im.size[0]) // 2
+    y = (square - im.size[1]) // 2
+    canvas.paste(im, (x, y))
+    update_display(IPyImage(data=_pil_to_png_bytes(canvas)), display_id=_PREVIEW_DISPLAY_ID)
+
+def set_preview_from_path(path: str, square: int = 256):
+    """Update preview from an image file path."""
+    set_preview_from_pil(PIL.open(path), square=square)
+
+# initialize placeholder
+init_preview(square=preview_size)
+#set_preview_from_pil(image, square=preview_size)
+
+
+
+"""
+TinyFlux-Deep: Deeper variant with 15 double + 25 single blocks.
+
+Config derived from checkpoint step_285625.safetensors:
+  - hidden_size: 512
+  - num_attention_heads: 4
+  - attention_head_dim: 128
+  - num_double_layers: 15
+  - num_single_layers: 25
+  - Uses biases in MLP
+  - Old RoPE format with cached freqs buffers
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, List
+
+@dataclass
+class TinyFluxDeepConfig:
+    """Configuration for TinyFlux-Deep model."""
+    hidden_size: int = 512
+    num_attention_heads: int = 4
+    attention_head_dim: int = 128
+
+    in_channels: int = 16
+    patch_size: int = 1
+
+    joint_attention_dim: int = 768
+    pooled_projection_dim: int = 768
+
+    num_double_layers: int = 15
+    num_single_layers: int = 25
+
+    mlp_ratio: float = 4.0
+    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
+    guidance_embeds: bool = True
+
+    def __post_init__(self):
+        assert self.num_attention_heads * self.attention_head_dim == self.hidden_size
+        assert sum(self.axes_dims_rope) == self.attention_head_dim
+
+
+# =============================================================================
+# Normalization
+# =============================================================================
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
+        super().__init__()
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.register_parameter('weight', None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        out = (x * norm).type_as(x)
+        if self.weight is not None:
+            out = out * self.weight
+        return out
+
+
+# =============================================================================
+# RoPE - Old format with cached frequency buffers (checkpoint compatible)
+# =============================================================================
+
+class EmbedND(nn.Module):
+    """
+    Original TinyFlux RoPE with cached frequency buffers.
+    Matches checkpoint format with rope.freqs_0, rope.freqs_1, rope.freqs_2
+    """
+
+    def __init__(self, theta: float = 10000.0, axes_dim: Tuple[int, int, int] = (16, 56, 56)):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        
+        # Register frequency buffers (matches checkpoint keys rope.freqs_*)
+        for i, dim in enumerate(axes_dim):
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+            self.register_buffer(f'freqs_{i}', freqs, persistent=True)
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            ids: (N, 3) position indices [temporal, height, width]
+        Returns:
+            rope: (N, 1, head_dim) interleaved [cos, sin, cos, sin, ...]
+        """
+        device = ids.device
+        n_axes = ids.shape[-1]
+        emb_list = []
+
+        for i in range(n_axes):
+            freqs = getattr(self, f'freqs_{i}').to(device)
+            pos = ids[:, i].float()
+            angles = pos.unsqueeze(-1) * freqs.unsqueeze(0)  # (N, dim/2)
+            
+            # Interleave cos and sin
+            cos = angles.cos()
+            sin = angles.sin()
+            emb = torch.stack([cos, sin], dim=-1).flatten(-2)  # (N, dim)
+            emb_list.append(emb)
+
+        rope = torch.cat(emb_list, dim=-1)  # (N, head_dim)
+        return rope.unsqueeze(1)  # (N, 1, head_dim)
+
+
+def apply_rotary_emb_old(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Apply rotary embeddings (old interleaved format).
+    
+    Args:
+        x: (B, H, N, D) query or key tensor
+        freqs_cis: (N, 1, D) interleaved [cos0, sin0, cos1, sin1, ...]
+    Returns:
+        Rotated tensor of same shape
+    """
+    # freqs_cis is (N, 1, D) with interleaved cos/sin
+    freqs = freqs_cis.squeeze(1)  # (N, D)
+    
+    # Split interleaved cos/sin
+    cos = freqs[:, 0::2].repeat_interleave(2, dim=-1)  # (N, D)
+    sin = freqs[:, 1::2].repeat_interleave(2, dim=-1)  # (N, D)
+    
+    cos = cos[None, None, :, :].to(x.device)  # (1, 1, N, D)
+    sin = sin[None, None, :, :].to(x.device)
+
+    # Split into real/imag pairs and rotate
+    x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(-2)
+
+    return (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+
+# =============================================================================
+# Embeddings
+# =============================================================================
+
+class MLPEmbedder(nn.Module):
+    """MLP for embedding scalars (timestep, guidance)."""
+
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(256, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        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)
+        return self.mlp(emb)
+
+
+# =============================================================================
+# AdaLayerNorm
+# =============================================================================
+
+class AdaLayerNormZero(nn.Module):
+    """AdaLN-Zero for double-stream blocks (6 params)."""
+
+    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):
+        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)
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class AdaLayerNormZeroSingle(nn.Module):
+    """AdaLN-Zero for single-stream blocks (3 params)."""
+
+    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):
+        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)
+        return x, gate
+
+
+# =============================================================================
+# Attention (original format - no Q/K norm, matches checkpoint)
+# =============================================================================
+
+class Attention(nn.Module):
+    """Multi-head attention (original TinyFlux format, no Q/K norm)."""
+
+    def __init__(self, hidden_size: int, num_heads: int, head_dim: int, use_bias: bool = False):
+        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=use_bias)
+        self.out_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=use_bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, N, _ = x.shape
+
+        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, H, N, D)
+
+        # Apply RoPE
+        if rope is not None:
+            q = apply_rotary_emb_old(q, rope)
+            k = apply_rotary_emb_old(k, rope)
+
+        # Scaled dot-product attention
+        attn = F.scaled_dot_product_attention(q, k, v)
+        out = attn.transpose(1, 2).reshape(B, N, -1)
+        return self.out_proj(out)
+
+
+class JointAttention(nn.Module):
+    """Joint attention for double-stream blocks (original format)."""
+
+    def __init__(self, hidden_size: int, num_heads: int, head_dim: int, use_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.scale = head_dim ** -0.5
+
+        self.txt_qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=use_bias)
+        self.img_qkv = nn.Linear(hidden_size, 3 * num_heads * head_dim, bias=use_bias)
+
+        self.txt_out = nn.Linear(num_heads * head_dim, hidden_size, bias=use_bias)
+        self.img_out = nn.Linear(num_heads * head_dim, hidden_size, bias=use_bias)
+
+    def forward(
+        self,
+        txt: torch.Tensor,
+        img: torch.Tensor,
+        rope: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        B, L, _ = txt.shape
+        _, N, _ = img.shape
+
+        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 only
+        if rope is not None:
+            img_q = apply_rotary_emb_old(img_q, rope)
+            img_k = apply_rotary_emb_old(img_k, rope)
+
+        # Concatenate for joint attention
+        k = torch.cat([txt_k, img_k], dim=2)
+        v = torch.cat([txt_v, img_v], dim=2)
+
+        txt_out = F.scaled_dot_product_attention(txt_q, k, v)
+        txt_out = txt_out.transpose(1, 2).reshape(B, L, -1)
+
+        img_out = F.scaled_dot_product_attention(img_q, k, v)
+        img_out = img_out.transpose(1, 2).reshape(B, N, -1)
+
+        return self.txt_out(txt_out), self.img_out(img_out)
+
+
+# =============================================================================
+# MLP (with bias - matches checkpoint)
+# =============================================================================
+
+class MLP(nn.Module):
+    """Feed-forward network with GELU activation and biases."""
+
+    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, bias=True)  # bias=True for checkpoint compat
+        self.act = nn.GELU(approximate='tanh')
+        self.fc2 = nn.Linear(mlp_hidden, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.act(self.fc1(x)))
+
+
+# =============================================================================
+# Transformer Blocks
+# =============================================================================
+
+class DoubleStreamBlock(nn.Module):
+    """Double-stream transformer block."""
+
+    def __init__(self, config: TinyFluxDeepConfig):
+        super().__init__()
+        hidden = config.hidden_size
+        heads = config.num_attention_heads
+        head_dim = config.attention_head_dim
+
+        self.img_norm1 = AdaLayerNormZero(hidden)
+        self.txt_norm1 = AdaLayerNormZero(hidden)
+
+        self.attn = JointAttention(hidden, heads, head_dim, use_bias=False)
+
+        self.img_norm2 = RMSNorm(hidden)
+        self.txt_norm2 = RMSNorm(hidden)
+
+        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,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        img_normed, img_gate_msa, img_shift_mlp, img_scale_mlp, img_gate_mlp = self.img_norm1(img, vec)
+        txt_normed, txt_gate_msa, txt_shift_mlp, txt_scale_mlp, txt_gate_mlp = self.txt_norm1(txt, vec)
+
+        txt_attn_out, img_attn_out = self.attn(txt_normed, img_normed, rope)
+
+        txt = txt + txt_gate_msa.unsqueeze(1) * txt_attn_out
+        img = img + img_gate_msa.unsqueeze(1) * img_attn_out
+
+        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."""
+
+    def __init__(self, config: TinyFluxDeepConfig):
+        super().__init__()
+        hidden = config.hidden_size
+        heads = config.num_attention_heads
+        head_dim = config.attention_head_dim
+
+        self.norm = AdaLayerNormZeroSingle(hidden)
+        self.attn = Attention(hidden, heads, head_dim, use_bias=False)
+        self.mlp = MLP(hidden, config.mlp_ratio)
+        self.norm2 = RMSNorm(hidden)
+
+    def forward(
+        self,
+        txt: torch.Tensor,
+        img: torch.Tensor,
+        vec: torch.Tensor,
+        rope: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        L = txt.shape[1]
+
+        x = torch.cat([txt, img], dim=1)
+
+        x_normed, gate = self.norm(x, vec)
+        x = x + gate.unsqueeze(1) * self.attn(x_normed, rope)
+        x = x + self.mlp(self.norm2(x))
+
+        txt, img = x.split([L, x.shape[1] - L], dim=1)
+        return txt, img
+
+
+# =============================================================================
+# Main Model
+# =============================================================================
+
+class TinyFluxDeep(nn.Module):
+    """TinyFlux-Deep: 15 double + 25 single blocks."""
+
+    def __init__(self, config: Optional[TinyFluxDeepConfig] = None):
+        super().__init__()
+        self.config = config or TinyFluxDeepConfig()
+        cfg = self.config
+
+        # Input projections (with bias to match checkpoint)
+        self.img_in = nn.Linear(cfg.in_channels, cfg.hidden_size, bias=True)
+        self.txt_in = nn.Linear(cfg.joint_attention_dim, cfg.hidden_size, bias=True)
+
+        # Conditioning
+        self.time_in = MLPEmbedder(cfg.hidden_size)
+        self.vector_in = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(cfg.pooled_projection_dim, cfg.hidden_size, bias=True)
+        )
+        if cfg.guidance_embeds:
+            self.guidance_in = MLPEmbedder(cfg.hidden_size)
+
+        # RoPE (old format with cached freqs)
+        self.rope = EmbedND(theta=10000.0, axes_dim=cfg.axes_dims_rope)
+
+        # Transformer blocks
+        self.double_blocks = nn.ModuleList([
+            DoubleStreamBlock(cfg) for _ in range(cfg.num_double_layers)
+        ])
+        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, bias=True)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        def _init(module):
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+        self.apply(_init)
+        nn.init.zeros_(self.final_linear.weight)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        pooled_projections: torch.Tensor,
+        timestep: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt_ids: Optional[torch.Tensor] = None,
+        guidance: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B = hidden_states.shape[0]
+        L = encoder_hidden_states.shape[1]
+        N = hidden_states.shape[1]
+
+        # Input projections
+        img = self.img_in(hidden_states)
+        txt = self.txt_in(encoder_hidden_states)
+
+        # Conditioning
+        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)
+
+        # Handle img_ids shape
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]  # (N, 3)
+
+        # Compute RoPE for image positions
+        img_rope = self.rope(img_ids)  # (N, 1, head_dim)
+
+        # Double-stream blocks
+        for block in self.double_blocks:
+            txt, img = block(txt, img, vec, img_rope)
+
+        # Build full sequence RoPE for single-stream
+        if txt_ids is None:
+            txt_ids = torch.zeros(L, 3, device=img_ids.device, dtype=img_ids.dtype)
+        elif txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        all_ids = torch.cat([txt_ids, img_ids], dim=0)
+        full_rope = self.rope(all_ids)
+
+        # Single-stream blocks
+        for block in self.single_blocks:
+            txt, img = block(txt, img, vec, full_rope)
+
+        # 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."""
+        img_ids = torch.zeros(height * width, 3, device=device)
+        for i in range(height):
+            for j in range(width):
+                idx = i * width + j
+                img_ids[idx, 0] = 0
+                img_ids[idx, 1] = i
+                img_ids[idx, 2] = j
+        return img_ids
+
+    @staticmethod
+    def create_txt_ids(text_len: int, device: torch.device) -> torch.Tensor:
+        """Create text position IDs."""
+        txt_ids = torch.zeros(text_len, 3, device=device)
+        txt_ids[:, 0] = torch.arange(text_len, device=device)
+        return txt_ids
+
+    def count_parameters(self) -> dict:
+        """Count parameters by component."""
+        counts = {}
+        counts['img_in'] = sum(p.numel() for p in self.img_in.parameters())
+        counts['txt_in'] = sum(p.numel() for p in self.txt_in.parameters())
+        counts['time_in'] = sum(p.numel() for p in self.time_in.parameters())
+        counts['vector_in'] = sum(p.numel() for p in self.vector_in.parameters())
+        if hasattr(self, 'guidance_in'):
+            counts['guidance_in'] = sum(p.numel() for p in self.guidance_in.parameters())
+        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())
+        counts['total'] = sum(p.numel() for p in self.parameters())
+        return counts
+
+
+# =============================================================================
+# Test
+# =============================================================================
+
+def test_model():
+    """Test TinyFlux-Deep model."""
+    print("=" * 60)
+    print("TinyFlux-Deep Test")
+    print("=" * 60)
+
+    config = TinyFluxDeepConfig()
+    model = TinyFluxDeep(config)
+
+    counts = model.count_parameters()
+    print(f"\nConfig:")
+    print(f"  hidden_size: {config.hidden_size}")
+    print(f"  num_attention_heads: {config.num_attention_heads}")
+    print(f"  attention_head_dim: {config.attention_head_dim}")
+    print(f"  num_double_layers: {config.num_double_layers}")
+    print(f"  num_single_layers: {config.num_single_layers}")
+    
+    print(f"\nParameters:")
+    for name, count in counts.items():
+        print(f"  {name}: {count:,}")
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = model.to(device)
+
+    B, H, W = 2, 64, 64
+    L = 77
+
+    hidden_states = torch.randn(B, H * W, config.in_channels, device=device)
+    encoder_hidden_states = torch.randn(B, L, config.joint_attention_dim, device=device)
+    pooled_projections = torch.randn(B, config.pooled_projection_dim, device=device)
+    timestep = torch.rand(B, device=device)
+    img_ids = TinyFluxDeep.create_img_ids(B, H, W, device)
+    txt_ids = TinyFluxDeep.create_txt_ids(L, device)
+    guidance = torch.ones(B, device=device) * 3.5
+
+    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,
+            txt_ids=txt_ids,
+            guidance=guidance,
+        )
+
+    print(f"\nOutput shape: {output.shape}")
+    print(f"Output range: [{output.min():.4f}, {output.max():.4f}]")
+    print("\n✓ Forward pass successful!")
+
+
+#if __name__ == "__main__":
+#    test_model()
+
+# ============================================================================
+# TinyFlux-Deep Inference Cell - Euler Discrete Flow Matching
+# ============================================================================
+# Run the model cell before this one (defines TinyFluxDeep, TinyFluxDeepConfig)
+# Loads from: AbstractPhil/tiny-flux-deep or local checkpoint
+# ============================================================================
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from transformers import T5EncoderModel, T5Tokenizer, CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL
+from PIL import Image
+import numpy as np
+import os
+
+
+# Generation settings
+NUM_STEPS = STEPS
+GUIDANCE_SCALE = CFG_GUIDANCE
+SHIFT = FLUX_SHIFT
+
+# ============================================================================
+# LOAD TEXT ENCODERS
+# ============================================================================
+print("Loading text encoders...")
+t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
+t5_enc = T5EncoderModel.from_pretrained("google/flan-t5-base", torch_dtype=DTYPE).to(DEVICE).eval()
+
+clip_tok = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+clip_enc = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=DTYPE).to(DEVICE).eval()
+
+# ============================================================================
+# LOAD VAE
+# ============================================================================
+print("Loading Flux VAE...")
+vae = AutoencoderKL.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    subfolder="vae",
+    torch_dtype=DTYPE
+).to(DEVICE).eval()
+
+# ============================================================================
+# LOAD TINYFLUX-DEEP MODEL
+# ============================================================================
+print(f"Loading TinyFlux-Deep from: {LOAD_FROM}")
+
+# Use TinyFluxDeep (512 hidden, 4 heads, 15 double, 25 single)
+config = TinyFluxDeepConfig()
+model = TinyFluxDeep(config).to(DEVICE).to(DTYPE)
+
+# Deprecated keys that may exist in old checkpoints but aren't needed
+DEPRECATED_KEYS = {'time_in.sin_basis', 'guidance_in.sin_basis'}
+
+
+def load_weights(path):
+    """Load weights from .safetensors or .pt file."""
+    if path.endswith(".safetensors"):
+        state_dict = load_file(path)
+    elif path.endswith(".pt"):
+        ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
+        if isinstance(ckpt, dict):
+            if "model" in ckpt:
+                state_dict = ckpt["model"]
+            elif "state_dict" in ckpt:
+                state_dict = ckpt["state_dict"]
+            else:
+                state_dict = ckpt
+        else:
+            state_dict = ckpt
+    else:
+        try:
+            state_dict = load_file(path)
+        except:
+            state_dict = torch.load(path, map_location=DEVICE, weights_only=False)
+    
+    # Strip "_orig_mod." prefix from keys (added by torch.compile)
+    if any(k.startswith("_orig_mod.") for k in state_dict.keys()):
+        print("  Stripping torch.compile prefix...")
+        state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
+    
+    return state_dict
+
+
+def load_model_weights(model, weights, source_name):
+    """Load weights with verbose reporting."""
+    # Filter out deprecated keys
+    filtered_weights = {k: v for k, v in weights.items() if k not in DEPRECATED_KEYS}
+    deprecated_found = [k for k in weights.keys() if k in DEPRECATED_KEYS]
+    
+    if deprecated_found:
+        print(f"  ✓ Ignored deprecated keys: {deprecated_found}")
+    
+    missing, unexpected = model.load_state_dict(filtered_weights, strict=False)
+    
+    if missing:
+        print(f"  ⚠ Missing keys: {missing[:10]}{'...' if len(missing) > 10 else ''}")
+    if unexpected:
+        print(f"  ⚠ Unexpected keys: {unexpected[:10]}{'...' if len(unexpected) > 10 else ''}")
+    if not missing and not unexpected:
+        print(f"  ✓ All weights loaded successfully")
+    
+    print(f"✓ Loaded from {source_name}")
+
+
+if LOAD_FROM == "hub":
+    try:
+        weights_path = hf_hub_download(repo_id=HF_REPO, filename="model.safetensors")
+    except:
+        weights_path = hf_hub_download(repo_id=HF_REPO, filename="model.pt")
+    weights = load_weights(weights_path)
+    load_model_weights(model, weights, HF_REPO)
+
+elif LOAD_FROM.startswith("hub:"):
+    ckpt_name = LOAD_FROM[4:]
+    for ext in [".safetensors", ".pt", ""]:
+        try:
+            if ckpt_name.endswith((".safetensors", ".pt")):
+                filename = ckpt_name if "/" in ckpt_name else f"checkpoints/{ckpt_name}"
+            else:
+                filename = f"checkpoints/{ckpt_name}{ext}"
+            weights_path = hf_hub_download(repo_id=HF_REPO, filename=filename)
+            weights = load_weights(weights_path)
+            load_model_weights(model, weights, f"{HF_REPO}/{filename}")
+            break
+        except Exception as e:
+            continue
+    else:
+        raise ValueError(f"Could not find checkpoint: {ckpt_name}")
+
+elif LOAD_FROM.startswith("local:"):
+    weights_path = LOAD_FROM[6:]
+    weights = load_weights(weights_path)
+    load_model_weights(model, weights, weights_path)
+
+else:
+    raise ValueError(f"Unknown LOAD_FROM: {LOAD_FROM}")
+
+model.eval()
+print(f"Model params: {sum(p.numel() for p in model.parameters()):,}")
+
+# ============================================================================
+# ENCODING FUNCTIONS
+# ============================================================================
+@torch.inference_mode()
+def encode_prompt(prompt: str, max_length: int = 128):
+    """Encode prompt with flan-t5-base and CLIP-L."""
+    t5_in = t5_tok(
+        prompt, 
+        max_length=max_length, 
+        padding="max_length", 
+        truncation=True, 
+        return_tensors="pt"
+    ).to(DEVICE)
+    t5_out = t5_enc(
+        input_ids=t5_in.input_ids, 
+        attention_mask=t5_in.attention_mask
+    ).last_hidden_state
+    
+    clip_in = clip_tok(
+        prompt, 
+        max_length=77, 
+        padding="max_length", 
+        truncation=True, 
+        return_tensors="pt"
+    ).to(DEVICE)
+    clip_out = clip_enc(
+        input_ids=clip_in.input_ids, 
+        attention_mask=clip_in.attention_mask
+    )
+    clip_pooled = clip_out.pooler_output
+    
+    return t5_out.to(DTYPE), clip_pooled.to(DTYPE)
+
+
+# ============================================================================
+# FLOW MATCHING HELPERS
+# ============================================================================
+def flux_shift(t, s=SHIFT):
+    """Flux timestep shift - biases towards higher t (closer to data)."""
+    return s * t / (1 + (s - 1) * t)
+
+
+# ============================================================================
+# EULER DISCRETE FLOW MATCHING SAMPLER
+# ============================================================================
+@torch.inference_mode()
+def euler_sample(
+    model,
+    prompt: str,
+    negative_prompt: str = "",
+    num_steps: int = 28,
+    guidance_scale: float = 3.5,
+    height: int = 512,
+    width: int = 512,
+    seed: int = None,
+):
+    """
+    Euler discrete sampler for rectified flow matching.
+    
+    Flow Matching formulation:
+        x_t = (1 - t) * noise + t * data
+        At t=0: noise, At t=1: data
+        Velocity v = data - noise (constant)
+        
+    Sampling: Integrate from t=0 (noise) to t=1 (data)
+    """
+    if seed is not None:
+        torch.manual_seed(seed)
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
+    else:
+        generator = None
+    
+    H_lat = height // 8
+    W_lat = width // 8
+    C_lat = 16
+    
+    # Encode prompts
+    t5_cond, clip_cond = encode_prompt(prompt)
+    if guidance_scale > 1.0 and negative_prompt is not None:
+        t5_uncond, clip_uncond = encode_prompt(negative_prompt)
+    else:
+        t5_uncond, clip_uncond = None, None
+    
+    # Start from pure noise (t=0)
+    x = torch.randn(1, H_lat * W_lat, C_lat, device=DEVICE, dtype=DTYPE, generator=generator)
+    
+    # Create image position IDs
+    img_ids = TinyFluxDeep.create_img_ids(1, H_lat, W_lat, DEVICE)
+    
+    # Timesteps: 0 → 1 with flux shift
+    t_linear = torch.linspace(0, 1, num_steps + 1, device=DEVICE, dtype=DTYPE)
+    timesteps = flux_shift(t_linear, s=SHIFT)
+    
+    print(f"Sampling with {num_steps} Euler steps (t: 0→1, shifted)...")
+    
+    for i in range(num_steps):
+        t_curr = timesteps[i]
+        t_next = timesteps[i + 1]
+        dt = t_next - t_curr
+        
+        t_batch = t_curr.unsqueeze(0)
+        guidance_embed = torch.tensor([guidance_scale], device=DEVICE, dtype=DTYPE)
+        
+        # Predict velocity
+        v_cond = model(
+            hidden_states=x,
+            encoder_hidden_states=t5_cond,
+            pooled_projections=clip_cond,
+            timestep=t_batch,
+            img_ids=img_ids,
+            guidance=guidance_embed,
+        )
+        
+        # Classifier-free guidance
+        if guidance_scale > 1.0 and t5_uncond is not None:
+            v_uncond = model(
+                hidden_states=x,
+                encoder_hidden_states=t5_uncond,
+                pooled_projections=clip_uncond,
+                timestep=t_batch,
+                img_ids=img_ids,
+                guidance=guidance_embed,
+            )
+            v = v_uncond + guidance_scale * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        
+        # Euler step: x_{t+dt} = x_t + v * dt
+        x = x + v * dt
+        
+        if (i + 1) % max(1, num_steps // 5) == 0 or i == num_steps - 1:
+            print(f"  Step {i+1}/{num_steps}, t={t_next.item():.3f}")
+    
+    # Reshape: (1, H*W, C) -> (1, C, H, W)
+    latents = x.reshape(1, H_lat, W_lat, C_lat).permute(0, 3, 1, 2)
+    
+    return latents
+
+
+# ============================================================================
+# DECODE LATENTS TO IMAGE
+# ============================================================================
+@torch.inference_mode()
+def decode_latents(latents):
+    """Decode VAE latents to PIL Image."""
+    latents = latents / vae.config.scaling_factor
+    image = vae.decode(latents.to(vae.dtype)).sample
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image[0].float().permute(1, 2, 0).cpu().numpy()
+    image = (image * 255).astype(np.uint8)
+    return Image.fromarray(image)
+
+
+# ============================================================================
+# MAIN GENERATION FUNCTION
+# ============================================================================
+def generate(
+    prompt: str = POSITIVE_PROMPT,
+    negative_prompt: str = NEGATIVE_PROMPT,
+    num_steps: int = NUM_STEPS,
+    guidance_scale: float = GUIDANCE_SCALE,
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    seed: int = SEED,
+    save_path: str = OUTPUT_PATH,
+):
+    """
+    Generate an image from a text prompt.
+    
+    Args:
+        prompt: Text description of desired image
+        negative_prompt: What to avoid (empty string for none)
+        num_steps: Number of Euler steps (20-50 recommended)
+        guidance_scale: CFG scale (1.0=none, 3-7 typical)
+        height: Output height in pixels (divisible by 8)
+        width: Output width in pixels (divisible by 8)
+        seed: Random seed (None for random)
+        save_path: Path to save image (None to skip)
+    
+    Returns:
+        PIL.Image
+    """
+    #print(f"\nGenerating: '{prompt}'")
+    #print(f"Settings: {num_steps} steps, cfg={guidance_scale}, {width}x{height}, seed={seed}")
+    
+    latents = euler_sample(
+        model=model,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_steps=num_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        seed=seed,
+    )
+    
+    #print("Decoding latents...")
+    image = decode_latents(latents)
+    
+    if save_path:
+        image.save(save_path)
+        #print(f"✓ Saved to {save_path}")
+
+        set_preview_from_pil(image, square=512)
+    
+    print("✓ Done!")
+    return image
+
+
+# ============================================================================
+# BATCH GENERATION
+# ============================================================================
+def generate_batch(
+    prompts: list,
+    negative_prompt: str = "",
+    num_steps: int = NUM_STEPS,
+    guidance_scale: float = GUIDANCE_SCALE,
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    seed: int = SEED,
+    output_dir: str = "./outputs",
+):
+    """Generate multiple images."""
+    os.makedirs(output_dir, exist_ok=True)
+    images = []
+    
+    for i, prompt in enumerate(prompts):
+        img_seed = seed + i if seed is not None else None
+        image = generate(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_steps=num_steps,
+            guidance_scale=guidance_scale,
+            height=height,
+            width=width,
+            seed=img_seed,
+            save_path=os.path.join(output_dir, f"{i:03d}.png"),
+        )
+        images.append(image)
+    
+    return images
+
+
+# ============================================================================
+# QUICK TEST
+# ============================================================================
+#print("\n" + "="*60)
+#print("TinyFlux-Deep Inference Ready!")
+#print("="*60)
+#print(f"Config: {config.hidden_size} hidden, {config.num_attention_heads} heads")
+#print(f"        {config.num_double_layers} double, {config.num_single_layers} single layers")
+#print(f"Total:  {sum(p.numel() for p in model.parameters()):,} parameters")
+
+# Example usage:
+image = generate()
+#image