Add iris_model.py

iris_model.py

@@ -0,0 +1,1246 @@
+"""
+IRIS: Iterative Recurrent Image Synthesis
+==========================================
+A novel architecture for mobile-first high-quality image generation.
+
+Key innovations:
+1. Wavelet-Frequency Latent Space (Haar DWT + lightweight VAE)
+2. Recurrent Depth Core (Prelude-Core-Coda with shared weights)
+3. Gated Recurrent Fourier Mixer (GRFM) — novel token mixing
+4. Manhattan Spatial Decay — learned 2D inductive bias
+5. Rectified Flow training with consistency distillation support
+6. Adaptive compute budget (4-16 iterations, same model)
+
+Author: IRIS Research
+License: Apache 2.0
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from dataclasses import dataclass, field
+
+
+# ============================================================================
+# Configuration
+# ============================================================================
+
+@dataclass
+class IRISConfig:
+    """Configuration for IRIS model."""
+    # Latent space
+    latent_channels: int = 16          # Channels in latent space
+    latent_spatial: int = 32           # Spatial dim of latent (for 512px with 16x compression)
+    
+    # Model dimensions
+    hidden_dim: int = 512              # Main hidden dimension
+    num_heads: int = 8                 # Number of attention heads
+    head_dim: int = 64                 # Dimension per head
+    ffn_ratio: float = 2.667          # FFN expansion ratio (SwiGLU-adjusted)
+    
+    # Architecture structure
+    num_prelude_blocks: int = 2        # Prelude blocks (unique weights)
+    num_core_layers: int = 4           # Layers WITHIN each core iteration
+    num_coda_blocks: int = 2           # Coda blocks (unique weights)
+    default_iterations: int = 8        # Default core iterations at inference
+    max_iterations: int = 16           # Maximum core iterations
+    
+    # GRFM settings
+    fourier_num_blocks: int = 8        # Block-diagonal blocks in Fourier MLP
+    sparsity_threshold: float = 0.01   # Soft-shrinkage lambda
+    recurrence_dim: int = 256          # Dimension for gated recurrence pathway
+    manhattan_window: int = 16         # Windowed Manhattan decay (for efficiency)
+    
+    # Cross-attention
+    text_dim: int = 768                # CLIP-L/14 text embedding dim
+    max_text_tokens: int = 77          # Maximum text sequence length
+    
+    # Patch embedding
+    patch_size: int = 2                # Patches in latent space (2×2)
+    
+    # Conditioning
+    num_timesteps: int = 1000          # Noise schedule discretization
+    
+    # VAE
+    vae_channels: list = field(default_factory=lambda: [32, 64, 128, 256])
+    
+    # Training
+    dropout: float = 0.0
+    
+    @property
+    def vae_latent_channels(self) -> int:
+        """VAE latent channels must match generator latent channels."""
+        return self.latent_channels
+    
+    @property
+    def num_patches(self) -> int:
+        return (self.latent_spatial // self.patch_size) ** 2
+    
+    @property
+    def patch_dim(self) -> int:
+        return self.latent_channels * self.patch_size * self.patch_size
+
+
+# ============================================================================
+# Wavelet Transforms (Haar)
+# ============================================================================
+
+class HaarDWT2D(nn.Module):
+    """2D Discrete Wavelet Transform using Haar wavelets.
+    Decomposes x ∈ R^{B,C,H,W} into R^{B,4C,H/2,W/2} (LL, LH, HL, HH subbands).
+    """
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Haar DWT: split into even/odd along both spatial dims
+        x_ll = (x[:, :, 0::2, 0::2] + x[:, :, 0::2, 1::2] + 
+                x[:, :, 1::2, 0::2] + x[:, :, 1::2, 1::2]) * 0.5
+        x_lh = (x[:, :, 0::2, 0::2] + x[:, :, 0::2, 1::2] - 
+                x[:, :, 1::2, 0::2] - x[:, :, 1::2, 1::2]) * 0.5
+        x_hl = (x[:, :, 0::2, 0::2] - x[:, :, 0::2, 1::2] + 
+                x[:, :, 1::2, 0::2] - x[:, :, 1::2, 1::2]) * 0.5
+        x_hh = (x[:, :, 0::2, 0::2] - x[:, :, 0::2, 1::2] - 
+                x[:, :, 1::2, 0::2] + x[:, :, 1::2, 1::2]) * 0.5
+        return torch.cat([x_ll, x_lh, x_hl, x_hh], dim=1)
+
+
+class HaarIDWT2D(nn.Module):
+    """2D Inverse Discrete Wavelet Transform (Haar).
+    Reconstructs x ∈ R^{B,C,H,W} from R^{B,4*(C//4),H/2,W/2}.
+    """
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C4, Hh, Wh = x.shape
+        C = C4 // 4
+        ll, lh, hl, hh = x[:, :C], x[:, C:2*C], x[:, 2*C:3*C], x[:, 3*C:]
+        
+        # Reconstruct 2× spatial resolution
+        H, W = Hh * 2, Wh * 2
+        out = torch.zeros(B, C, H, W, device=x.device, dtype=x.dtype)
+        out[:, :, 0::2, 0::2] = (ll + lh + hl + hh) * 0.5
+        out[:, :, 0::2, 1::2] = (ll + lh - hl - hh) * 0.5
+        out[:, :, 1::2, 0::2] = (ll - lh + hl - hh) * 0.5
+        out[:, :, 1::2, 1::2] = (ll - lh - hl + hh) * 0.5
+        return out
+
+
+# ============================================================================
+# Lightweight Wavelet VAE
+# ============================================================================
+
+class DepthwiseSeparableConv(nn.Module):
+    """Depthwise separable convolution — key mobile optimization."""
+    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1):
+        super().__init__()
+        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size, stride, padding, groups=in_ch)
+        self.pointwise = nn.Conv2d(in_ch, out_ch, 1)
+    
+    def forward(self, x):
+        return self.pointwise(self.depthwise(x))
+
+
+class ResBlock(nn.Module):
+    """Residual block with depthwise separable convolutions."""
+    def __init__(self, channels):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(8, channels)
+        self.conv1 = DepthwiseSeparableConv(channels, channels)
+        self.norm2 = nn.GroupNorm(8, channels)
+        self.conv2 = DepthwiseSeparableConv(channels, channels)
+        # Zero-init final layer for residual learning stability
+        nn.init.zeros_(self.conv2.pointwise.weight)
+        nn.init.zeros_(self.conv2.pointwise.bias)
+    
+    def forward(self, x):
+        h = F.silu(self.norm1(x))
+        h = self.conv1(h)
+        h = F.silu(self.norm2(h))
+        h = self.conv2(h)
+        return x + h
+
+
+class WaveletVAEEncoder(nn.Module):
+    """Lightweight encoder: Haar DWT preprocessing + small convolutional encoder.
+    Input: images R^{B,3,H,W} → Output: latent R^{B,C_latent,H/16,W/16}
+    Compression: 3×H×W → C_latent×(H/16)×(W/16)
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        self.dwt = HaarDWT2D()
+        channels = config.vae_channels
+        latent_ch = config.vae_latent_channels
+        
+        # DWT: 3 channels → 12 channels at H/2 × W/2
+        self.conv_in = nn.Conv2d(12, channels[0], 3, 1, 1)
+        
+        # Downsampling path: H/2→H/4→H/8→H/16
+        self.down_blocks = nn.ModuleList()
+        for i in range(len(channels) - 1):
+            self.down_blocks.append(nn.Sequential(
+                ResBlock(channels[i]),
+                nn.Conv2d(channels[i], channels[i+1], 3, 2, 1),  # 2× downsample
+            ))
+        
+        # Bottleneck
+        self.mid = nn.Sequential(
+            ResBlock(channels[-1]),
+            ResBlock(channels[-1]),
+        )
+        
+        # To latent (mean + logvar)
+        self.norm_out = nn.GroupNorm(8, channels[-1])
+        self.conv_out = nn.Conv2d(channels[-1], 2 * latent_ch, 1)
+    
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Haar DWT preprocessing
+        x = self.dwt(x)  # [B, 12, H/2, W/2]
+        x = self.conv_in(x)
+        
+        for down in self.down_blocks:
+            x = down(x)
+        
+        x = self.mid(x)
+        x = F.silu(self.norm_out(x))
+        x = self.conv_out(x)
+        
+        mean, logvar = x.chunk(2, dim=1)
+        logvar = torch.clamp(logvar, -30.0, 20.0)
+        return mean, logvar
+    
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        mean, logvar = self.forward(x)
+        std = torch.exp(0.5 * logvar)
+        z = mean + std * torch.randn_like(std)
+        return z, mean, logvar
+
+
+class WaveletVAEDecoder(nn.Module):
+    """Tiny decoder: latent → wavelet coefficients → Haar IDWT → image.
+    Designed to be as small as possible for mobile inference.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        channels = list(reversed(config.vae_channels))
+        latent_ch = config.vae_latent_channels
+        self.idwt = HaarIDWT2D()
+        
+        # From latent
+        self.conv_in = nn.Conv2d(latent_ch, channels[0], 3, 1, 1)
+        
+        # Bottleneck
+        self.mid = nn.Sequential(
+            ResBlock(channels[0]),
+        )
+        
+        # Upsampling path
+        self.up_blocks = nn.ModuleList()
+        for i in range(len(channels) - 1):
+            self.up_blocks.append(nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                DepthwiseSeparableConv(channels[i], channels[i+1]),
+                nn.SiLU(),
+                ResBlock(channels[i+1]),
+            ))
+        
+        # To wavelet coefficients (12 channels: 4 subbands × 3 color channels)
+        self.norm_out = nn.GroupNorm(8, channels[-1])
+        self.conv_out = nn.Conv2d(channels[-1], 12, 3, 1, 1)
+    
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        x = self.conv_in(z)
+        x = self.mid(x)
+        
+        for up in self.up_blocks:
+            x = up(x)
+        
+        x = F.silu(self.norm_out(x))
+        x = self.conv_out(x)  # [B, 12, H/2, W/2] wavelet coefficients
+        
+        # Inverse DWT to get image
+        x = self.idwt(x)  # [B, 3, H, W]
+        return x
+
+
+class WaveletVAE(nn.Module):
+    """Complete Wavelet VAE with DWT preprocessing."""
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        self.encoder = WaveletVAEEncoder(config)
+        self.decoder = WaveletVAEDecoder(config)
+        self.config = config
+    
+    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.encoder.encode(x)
+    
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+    
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        z, mean, logvar = self.encode(x)
+        x_recon = self.decode(z)
+        return x_recon, mean, logvar
+
+
+# ============================================================================
+# Conditioning Modules
+# ============================================================================
+
+class TimestepEmbedding(nn.Module):
+    """Sinusoidal timestep embedding with MLP projection."""
+    def __init__(self, dim: int, max_period: int = 10000):
+        super().__init__()
+        self.dim = dim
+        self.max_period = max_period
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, 4 * dim),
+            nn.SiLU(),
+            nn.Linear(4 * dim, dim),
+        )
+    
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.dim // 2
+        freqs = torch.exp(
+            -math.log(self.max_period) * torch.arange(half, device=t.device, dtype=t.dtype) / half
+        )
+        args = t[:, None] * freqs[None, :]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.mlp(embedding)
+
+
+class IterationEmbedding(nn.Module):
+    """Learnable embedding for iteration index within recurrent core."""
+    def __init__(self, max_iterations: int, dim: int):
+        super().__init__()
+        self.embedding = nn.Embedding(max_iterations, dim)
+    
+    def forward(self, i: torch.Tensor) -> torch.Tensor:
+        return self.embedding(i)
+
+
+class AdaLNSingle(nn.Module):
+    """Adaptive Layer Normalization (single shared MLP, per-layer bias).
+    From PixArt-α: saves 27% params vs standard adaLN.
+    
+    Produces (scale, shift, gate) for each sub-layer from a shared condition vector.
+    """
+    def __init__(self, dim: int, num_modulations: int = 6):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, num_modulations * dim)
+        self.num_modulations = num_modulations
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+    
+    def forward(self, c: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """c: [B, D] condition vector → tuple of num_modulations tensors [B, D]."""
+        params = self.linear(self.silu(c))
+        return params.chunk(self.num_modulations, dim=-1)
+
+
+# ============================================================================
+# GRFM: Gated Recurrent Fourier Mixer (Novel Contribution)
+# ============================================================================
+
+class FourierMixingPathway(nn.Module):
+    """Pathway 1: Adaptive Fourier Neural Operator-style global mixing.
+    O(N log N) complexity via FFT. Block-diagonal MLP in frequency domain.
+    """
+    def __init__(self, dim: int, num_blocks: int = 8, sparsity_threshold: float = 0.01):
+        super().__init__()
+        self.dim = dim
+        self.num_blocks = num_blocks
+        self.block_size = dim // num_blocks
+        self.sparsity_threshold = sparsity_threshold
+        
+        # Block-diagonal complex-valued MLP in Fourier domain
+        # Each block: R^{block_size} → R^{block_size}
+        # Using real-valued params for complex ops (split real/imag)
+        self.w1_real = nn.Parameter(torch.randn(num_blocks, self.block_size, self.block_size) * 0.02)
+        self.w1_imag = nn.Parameter(torch.randn(num_blocks, self.block_size, self.block_size) * 0.02)
+        self.w2_real = nn.Parameter(torch.randn(num_blocks, self.block_size, self.block_size) * 0.02)
+        self.w2_imag = nn.Parameter(torch.randn(num_blocks, self.block_size, self.block_size) * 0.02)
+        self.b1 = nn.Parameter(torch.zeros(num_blocks, self.block_size))
+        self.b2 = nn.Parameter(torch.zeros(num_blocks, self.block_size))
+    
+    def complex_matmul(self, x: torch.Tensor, w_real: torch.Tensor, w_imag: torch.Tensor) -> torch.Tensor:
+        """Complex matrix multiplication: (a+bi)(c+di) = (ac-bd) + (ad+bc)i
+        x: [..., num_blocks, block_size] (complex)
+        w: [num_blocks, block_size, block_size] (real)
+        """
+        # Use einsum for proper block-diagonal matmul
+        # x: [B, Hf, Wf, K, bs], w: [K, bs, bs] → out: [B, Hf, Wf, K, bs]
+        out_real = torch.einsum('...ki,kij->...kj', x.real, w_real) - torch.einsum('...ki,kij->...kj', x.imag, w_imag)
+        out_imag = torch.einsum('...ki,kij->...kj', x.real, w_imag) + torch.einsum('...ki,kij->...kj', x.imag, w_real)
+        return torch.complex(out_real, out_imag)
+    
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        B, N, D = x.shape
+        x_2d = x.reshape(B, H, W, D)
+        
+        # 2D Real FFT on spatial dimensions
+        x_freq = torch.fft.rfft2(x_2d, dim=(1, 2), norm='ortho')  # [B, H, W//2+1, D]
+        
+        # Reshape channel dim for block-diagonal MLP: D → (num_blocks, block_size)
+        Hf, Wf = x_freq.shape[1], x_freq.shape[2]
+        x_freq = x_freq.reshape(B, Hf, Wf, self.num_blocks, self.block_size)
+        
+        # Block MLP Layer 1: operates on last dim (block_size)
+        # x_freq: [B, Hf, Wf, num_blocks, block_size]
+        # w1: [num_blocks, block_size, block_size]
+        x_freq = self.complex_matmul(x_freq, self.w1_real, self.w1_imag)
+        x_freq = x_freq + self.b1  # Broadcast bias (real only)
+        x_freq = torch.complex(F.relu(x_freq.real), F.relu(x_freq.imag))
+        
+        # Block MLP Layer 2
+        x_freq = self.complex_matmul(x_freq, self.w2_real, self.w2_imag)
+        x_freq = x_freq + self.b2
+        
+        # Reshape back to [B, Hf, Wf, D]
+        x_freq = x_freq.reshape(B, Hf, Wf, D)
+        
+        # Soft-shrinkage (sparsity in Fourier domain)
+        magnitude = x_freq.abs()
+        shrunk_mag = F.relu(magnitude - self.sparsity_threshold)
+        # Preserve phase, shrink magnitude
+        x_freq = x_freq * (shrunk_mag / (magnitude + 1e-8))
+        
+        # Inverse FFT
+        x_out = torch.fft.irfft2(x_freq, s=(H, W), dim=(1, 2), norm='ortho')
+        return x_out.reshape(B, N, D)
+
+
+class GatedLinearRecurrence(nn.Module):
+    """Pathway 2: Bidirectional Gated Linear Recurrence (RG-LRU inspired).
+    O(N) complexity with O(1) state per position.
+    
+    h_t = a_t * h_{t-1} + sqrt(1 - a_t^2) * (i_t * x_t)
+    where a_t = sigmoid(Λ)^(c * sigmoid(W_a * x_t))
+    """
+    def __init__(self, dim: int, recurrence_dim: int):
+        super().__init__()
+        self.dim = dim
+        self.rec_dim = recurrence_dim
+        
+        # Project to recurrence space
+        self.proj_in = nn.Linear(dim, recurrence_dim * 2)  # Forward + backward
+        
+        # Gating parameters
+        self.W_a = nn.Linear(recurrence_dim, recurrence_dim, bias=False)
+        self.W_x = nn.Linear(recurrence_dim, recurrence_dim, bias=False)
+        self.Lambda = nn.Parameter(torch.randn(recurrence_dim) * 0.5 + 2.0)  # Init for decay ~0.88-0.95
+        self.c = 8.0  # Decay scaling constant (from Griffin)
+        
+        # Output projection
+        self.proj_out = nn.Linear(recurrence_dim * 2, dim)
+    
+    def _scan(self, x: torch.Tensor) -> torch.Tensor:
+        """Sequential scan for a single direction. x: [B, N, rec_dim]"""
+        B, N, D = x.shape
+        
+        # Compute gates (can be parallelized)
+        a_base = torch.sigmoid(self.Lambda)  # [D]
+        r = torch.sigmoid(self.W_a(x))       # [B, N, D] - recurrence gate
+        i = torch.sigmoid(self.W_x(x))       # [B, N, D] - input gate
+        
+        # a_t = a_base^(c * r_t) — data-dependent decay
+        a = a_base.pow(self.c * r)            # [B, N, D]
+        
+        # Normalized input: sqrt(1 - a^2) for variance preservation
+        input_scale = torch.sqrt(1.0 - a * a + 1e-8)
+        scaled_input = input_scale * (i * x)  # [B, N, D]
+        
+        # Sequential recurrence (use parallel scan in production)
+        outputs = []
+        h = torch.zeros(B, D, device=x.device, dtype=x.dtype)
+        for t in range(N):
+            h = a[:, t] * h + scaled_input[:, t]
+            outputs.append(h)
+        
+        return torch.stack(outputs, dim=1)  # [B, N, D]
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, D = x.shape
+        
+        # Project to recurrence space and split for bidirectional
+        x_proj = self.proj_in(x)  # [B, N, 2*rec_dim]
+        x_fwd, x_bwd = x_proj.chunk(2, dim=-1)
+        
+        # Forward and backward scans
+        h_fwd = self._scan(x_fwd)
+        h_bwd = self._scan(x_bwd.flip(1)).flip(1)
+        
+        # Merge bidirectional
+        h = torch.cat([h_fwd, h_bwd], dim=-1)
+        return self.proj_out(h)
+
+
+class ManhattanSpatialGate(nn.Module):
+    """Pathway 3: Manhattan distance spatial decay gating.
+    Provides learned 2D spatial inductive bias with per-head multi-scale receptive fields.
+    Uses windowed computation for efficiency.
+    """
+    def __init__(self, dim: int, num_heads: int, window: int = 16):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window = window
+        
+        # Per-head learnable decay rate
+        # Initialize so gamma ∈ [0.7, 0.95] — multi-scale
+        self.gamma_logit = nn.Parameter(torch.linspace(0.85, 2.94, num_heads))  # sigmoid → [0.7, 0.95]
+        
+        # Value and gate projections
+        self.v_proj = nn.Linear(dim, dim)
+        self.g_proj = nn.Linear(dim, dim)
+        self.o_proj = nn.Linear(dim, dim)
+    
+    def _get_manhattan_mask(self, H: int, W: int, device: torch.device) -> torch.Tensor:
+        """Compute Manhattan distance matrix between all 2D positions."""
+        coords = torch.stack(torch.meshgrid(
+            torch.arange(H, device=device), 
+            torch.arange(W, device=device),
+            indexing='ij'
+        ), dim=-1).reshape(-1, 2).float()  # [N, 2]
+        
+        # Manhattan distance: |x1-x2| + |y1-y2|
+        dist = torch.cdist(coords, coords, p=1)  # [N, N]
+        return dist
+    
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        B, N, D = x.shape
+        
+        # Compute spatial decay
+        gamma = torch.sigmoid(self.gamma_logit)  # [num_heads]
+        manhattan_dist = self._get_manhattan_mask(H, W, x.device)  # [N, N]
+        
+        # Window the distance matrix for efficiency
+        # Only compute decay for positions within window distance
+        decay_mask = (manhattan_dist <= self.window).float()
+        
+        # Per-head decay: gamma_h^dist
+        decay = gamma[:, None, None].pow(manhattan_dist[None, :, :])  # [heads, N, N]
+        decay = decay * decay_mask[None, :, :]
+        
+        # Value and gate
+        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim)
+        g = torch.sigmoid(self.g_proj(x))
+        
+        # Apply spatial decay to values
+        # [B, heads, N, head_dim] = [heads, N, N] @ [B, heads, N, head_dim]
+        v = v.permute(0, 2, 1, 3)  # [B, heads, N, head_dim]
+        out = torch.matmul(decay.unsqueeze(0), v)  # [B, heads, N, head_dim]
+        
+        # Normalize by decay sum
+        decay_sum = decay.sum(dim=-1, keepdim=True).unsqueeze(0) + 1e-8  # [1, heads, N, 1]
+        out = out / decay_sum
+        
+        out = out.permute(0, 2, 1, 3).reshape(B, N, D)  # [B, N, D]
+        out = out * g  # Gating
+        return self.o_proj(out)
+
+
+class GRFM(nn.Module):
+    """Gated Recurrent Fourier Mixer — the core innovation of IRIS.
+    
+    Fuses three complementary pathways:
+    1. Fourier Global Mixing (O(N log N)) — captures textures, patterns
+    2. Gated Linear Recurrence (O(N)) — captures sequential/local dependencies
+    3. Manhattan Spatial Gate — provides 2D inductive bias
+    
+    Pathways are combined via learned adaptive gating.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        D = config.hidden_dim
+        
+        self.fourier = FourierMixingPathway(D, config.fourier_num_blocks, config.sparsity_threshold)
+        self.recurrence = GatedLinearRecurrence(D, config.recurrence_dim)
+        self.spatial = ManhattanSpatialGate(D, config.num_heads, config.manhattan_window)
+        
+        # Adaptive gate: learns to blend Fourier vs Recurrence based on content
+        self.blend_gate = nn.Sequential(
+            nn.Linear(D, D),
+            nn.SiLU(),
+            nn.Linear(D, D),
+            nn.Sigmoid(),
+        )
+        
+        # Spatial pathway weight (smaller contribution, additive)
+        self.spatial_scale = nn.Parameter(torch.tensor(0.1))
+    
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        # Three pathways
+        x_fourier = self.fourier(x, H, W)
+        x_recurrent = self.recurrence(x)
+        x_spatial = self.spatial(x, H, W)
+        
+        # Adaptive blending
+        gate = self.blend_gate(x)  # [B, N, D] values in [0, 1]
+        
+        # Fourier for global structure, recurrence for local detail
+        output = gate * x_fourier + (1 - gate) * x_recurrent
+        
+        # Add spatial bias (small contribution)
+        output = output + self.spatial_scale * x_spatial
+        
+        return output
+
+
+# ============================================================================
+# Cross-Attention (for text conditioning)
+# ============================================================================
+
+class CrossAttention(nn.Module):
+    """Efficient cross-attention for text conditioning.
+    Only 77 text tokens → O(N × 77 × d) per layer, very cheap.
+    """
+    def __init__(self, dim: int, text_dim: 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.q_proj = nn.Linear(dim, num_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(text_dim, num_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(text_dim, num_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * head_dim, dim)
+        
+        # QK normalization for stability (from SANA-Sprint)
+        self.q_norm = nn.RMSNorm(head_dim)
+        self.k_norm = nn.RMSNorm(head_dim)
+    
+    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        B, N, _ = x.shape
+        _, S, _ = context.shape
+        
+        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(context).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(context).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        # QK normalization
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        
+        # Scaled dot-product attention
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        
+        out = out.transpose(1, 2).reshape(B, N, -1)
+        return self.o_proj(out)
+
+
+# ============================================================================
+# Feed-Forward Network (SwiGLU)
+# ============================================================================
+
+class SwiGLUFFN(nn.Module):
+    """SwiGLU Feed-Forward Network — better than GELU for transformers."""
+    def __init__(self, dim: int, ratio: float = 2.667, dropout: float = 0.0):
+        super().__init__()
+        hidden = int(dim * ratio)
+        # Ensure hidden is multiple of 64 for hardware efficiency
+        hidden = ((hidden + 63) // 64) * 64
+        
+        self.w1 = nn.Linear(dim, hidden, bias=False)
+        self.w2 = nn.Linear(dim, hidden, bias=False)  # Gate
+        self.w3 = nn.Linear(hidden, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w3(self.dropout(F.silu(self.w1(x)) * self.w2(x)))
+
+
+# ============================================================================
+# Prelude Block (unique weights, conv-based)
+# ============================================================================
+
+class PreludeBlock(nn.Module):
+    """Lightweight conv-based block for initial feature extraction."""
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=5, padding=2, groups=dim)
+        self.pointwise = nn.Linear(dim, dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = SwiGLUFFN(dim)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Depthwise conv path
+        h = self.norm1(x)
+        h = h.transpose(1, 2)  # [B, D, N]
+        h = self.dwconv(h).transpose(1, 2)  # [B, N, D]
+        h = F.silu(h)
+        h = self.pointwise(h)
+        x = x + h
+        
+        # FFN
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+# ============================================================================
+# Core Block (shared weights, the heart of IRIS)
+# ============================================================================
+
+class CoreLayer(nn.Module):
+    """Single layer within the core block.
+    Contains: GRFM + Cross-Attention + FFN, all with adaLN-Zero conditioning.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        D = config.hidden_dim
+        
+        # Sub-layer 1: GRFM
+        self.norm1 = nn.LayerNorm(D, elementwise_affine=False)
+        self.grfm = GRFM(config)
+        
+        # Sub-layer 2: Cross-Attention
+        self.norm2 = nn.LayerNorm(D, elementwise_affine=False)
+        self.cross_attn = CrossAttention(D, config.text_dim, config.num_heads, config.head_dim)
+        
+        # Sub-layer 3: FFN
+        self.norm3 = nn.LayerNorm(D, elementwise_affine=False)
+        self.ffn = SwiGLUFFN(D, config.ffn_ratio, config.dropout)
+        
+        # adaLN-Zero: 9 modulations (scale1, shift1, gate1, scale2, shift2, gate2, scale3, shift3, gate3)
+        self.adaln = AdaLNSingle(D, num_modulations=9)
+    
+    def _modulate(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+    
+    def forward(self, x: torch.Tensor, c: torch.Tensor, text_tokens: torch.Tensor, 
+                H: int, W: int) -> torch.Tensor:
+        """
+        x: [B, N, D] — token sequence
+        c: [B, D] — conditioning vector (timestep + iteration)
+        text_tokens: [B, S, text_dim] — CLIP text tokens
+        H, W: spatial dimensions of token grid
+        """
+        s1, sh1, g1, s2, sh2, g2, s3, sh3, g3 = self.adaln(c)
+        
+        # GRFM with adaLN-Zero
+        h = self._modulate(self.norm1(x), s1, sh1)
+        h = self.grfm(h, H, W)
+        x = x + g1.unsqueeze(1) * h
+        
+        # Cross-attention with adaLN-Zero
+        h = self._modulate(self.norm2(x), s2, sh2)
+        h = self.cross_attn(h, text_tokens)
+        x = x + g2.unsqueeze(1) * h
+        
+        # FFN with adaLN-Zero
+        h = self._modulate(self.norm3(x), s3, sh3)
+        h = self.ffn(h)
+        x = x + g3.unsqueeze(1) * h
+        
+        return x
+
+
+class CoreBlock(nn.Module):
+    """The shared-weight core block, iterated r times.
+    Contains multiple CoreLayers to give sufficient per-iteration capacity.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            CoreLayer(config) for _ in range(config.num_core_layers)
+        ])
+    
+    def forward(self, x: torch.Tensor, c: torch.Tensor, text_tokens: torch.Tensor,
+                H: int, W: int) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x, c, text_tokens, H, W)
+        return x
+
+
+# ============================================================================
+# Coda Block (unique weights, final refinement)
+# ============================================================================
+
+class LocalWindowAttention(nn.Module):
+    """Window-based local attention for final refinement.
+    Small window (8×8) for efficient local detail refinement.
+    """
+    def __init__(self, dim: int, num_heads: int, head_dim: int, window_size: int = 8):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.window_size = window_size
+        self.scale = head_dim ** -0.5
+        
+        self.qkv = nn.Linear(dim, 3 * num_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * head_dim, dim)
+    
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        B, N, D = x.shape
+        ws = self.window_size
+        
+        # Reshape to 2D and partition into windows
+        x_2d = x.reshape(B, H, W, D)
+        
+        # Pad if necessary
+        pad_h = (ws - H % ws) % ws
+        pad_w = (ws - W % ws) % ws
+        if pad_h > 0 or pad_w > 0:
+            x_2d = F.pad(x_2d, (0, 0, 0, pad_w, 0, pad_h))
+        
+        Hp, Wp = x_2d.shape[1], x_2d.shape[2]
+        nH, nW = Hp // ws, Wp // ws
+        
+        # [B, nH, ws, nW, ws, D] → [B*nH*nW, ws*ws, D]
+        x_win = x_2d.reshape(B, nH, ws, nW, ws, D)
+        x_win = x_win.permute(0, 1, 3, 2, 4, 5).reshape(-1, ws * ws, D)
+        
+        # QKV and attention within windows
+        qkv = self.qkv(x_win).reshape(-1, ws * ws, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(2)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        
+        out = out.transpose(1, 2).reshape(-1, ws * ws, self.num_heads * self.head_dim)
+        out = self.o_proj(out)
+        
+        # Unpartition
+        out = out.reshape(B, nH, nW, ws, ws, D)
+        out = out.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, D)
+        
+        # Remove padding
+        out = out[:, :H, :W, :].reshape(B, N, D)
+        return out
+
+
+class CodaBlock(nn.Module):
+    """Final refinement block with local window attention."""
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        D = config.hidden_dim
+        self.norm1 = nn.LayerNorm(D)
+        self.attn = LocalWindowAttention(D, config.num_heads, config.head_dim, window_size=8)
+        self.norm2 = nn.LayerNorm(D)
+        self.ffn = SwiGLUFFN(D, config.ffn_ratio)
+    
+    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), H, W)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+# ============================================================================
+# IRIS Generator (Main Model)
+# ============================================================================
+
+class IRISGenerator(nn.Module):
+    """
+    IRIS: Iterative Recurrent Image Synthesis
+    
+    The main denoising network with Prelude-Core-Coda structure.
+    Predicts velocity field v for rectified flow training.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        self.config = config
+        D = config.hidden_dim
+        
+        # Patch embedding: latent patches → tokens
+        self.patch_embed = nn.Linear(config.patch_dim, D)
+        
+        # Positional embedding (learned)
+        self.pos_embed = nn.Parameter(torch.randn(1, config.num_patches, D) * 0.02)
+        
+        # Conditioning
+        self.time_embed = TimestepEmbedding(D)
+        self.iter_embed = IterationEmbedding(config.max_iterations, D)
+        self.text_proj = nn.Linear(config.text_dim, D)  # Project CLIP text to model dim
+        
+        # Global text pooling for conditioning
+        self.text_pool_proj = nn.Sequential(
+            nn.Linear(config.text_dim, D),
+            nn.SiLU(),
+            nn.Linear(D, D),
+        )
+        
+        # Prelude (unique weights)
+        self.prelude = nn.ModuleList([PreludeBlock(D) for _ in range(config.num_prelude_blocks)])
+        
+        # Core (shared weights, iterated)
+        self.core = CoreBlock(config)
+        
+        # Long skip connection (from Diffusion-RWKV: linear(cat(shallow, deep)))
+        self.skip_proj = nn.Linear(2 * D, D)
+        
+        # Coda (unique weights)
+        self.coda = nn.ModuleList([CodaBlock(config) for _ in range(config.num_coda_blocks)])
+        
+        # Output projection: tokens → latent patches
+        self.final_norm = nn.LayerNorm(D)
+        self.output_proj = nn.Linear(D, config.patch_dim)
+        
+        # Zero-init output for stable training start
+        nn.init.zeros_(self.output_proj.weight)
+        nn.init.zeros_(self.output_proj.bias)
+        
+        # Precompute patch spatial dimensions
+        self.patch_h = config.latent_spatial // config.patch_size
+        self.patch_w = config.latent_spatial // config.patch_size
+    
+    def patchify(self, z: torch.Tensor) -> torch.Tensor:
+        """Convert latent z [B, C, H, W] → patches [B, N, patch_dim]."""
+        B, C, H, W = z.shape
+        p = self.config.patch_size
+        z = z.reshape(B, C, H // p, p, W // p, p)
+        z = z.permute(0, 2, 4, 1, 3, 5).reshape(B, -1, C * p * p)
+        return z
+    
+    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
+        """Convert patches [B, N, patch_dim] → latent [B, C, H, W]."""
+        B, N, _ = x.shape
+        p = self.config.patch_size
+        C = self.config.latent_channels
+        H = self.patch_h
+        W = self.patch_w
+        x = x.reshape(B, H, W, C, p, p)
+        x = x.permute(0, 3, 1, 4, 2, 5).reshape(B, C, H * p, W * p)
+        return x
+    
+    def forward(
+        self,
+        z_t: torch.Tensor,          # Noisy latent [B, C, H, W]
+        t: torch.Tensor,            # Timestep [B] in [0, 1]
+        text_tokens: torch.Tensor,  # CLIP text embeddings [B, S, text_dim]
+        num_iterations: Optional[int] = None,  # Override iteration count
+    ) -> torch.Tensor:
+        """Predict velocity field v(z_t, t, c) for rectified flow."""
+        B = z_t.shape[0]
+        r = num_iterations or self.config.default_iterations
+        H, W = self.patch_h, self.patch_w
+        
+        # Patchify and embed
+        x = self.patch_embed(self.patchify(z_t)) + self.pos_embed
+        
+        # Timestep conditioning
+        t_emb = self.time_embed(t * self.config.num_timesteps)  # [B, D]
+        
+        # Text conditioning (project to model dim for cross-attention)
+        text_projected = self.text_proj(text_tokens)  # [B, S, D]
+        
+        # Global text pool for adaLN conditioning
+        text_global = self.text_pool_proj(text_tokens.mean(dim=1))  # [B, D]
+        
+        # ============ PRELUDE ============
+        for block in self.prelude:
+            x = block(x)
+        
+        # Save for long skip connection
+        x_shallow = x
+        
+        # ============ CORE (iterated r times) ============
+        for i in range(r):
+            # Iteration-aware conditioning
+            iter_idx = torch.full((B,), i, device=z_t.device, dtype=torch.long)
+            i_emb = self.iter_embed(iter_idx)  # [B, D]
+            
+            # Combined conditioning: timestep + iteration + text global
+            c = t_emb + i_emb + text_global  # [B, D]
+            
+            # Apply shared core block (pass original text_tokens for cross-attention)
+            x = self.core(x, c, text_tokens, H, W)
+        
+        # Long skip connection (from Diffusion-RWKV paper)
+        x = self.skip_proj(torch.cat([x_shallow, x], dim=-1))
+        
+        # ============ CODA ============
+        for block in self.coda:
+            x = block(x, H, W)
+        
+        # Output projection
+        x = self.final_norm(x)
+        x = self.output_proj(x)
+        
+        # Unpatchify to latent shape
+        v_pred = self.unpatchify(x)
+        return v_pred
+
+
+# ============================================================================
+# Full IRIS System
+# ============================================================================
+
+class IRIS(nn.Module):
+    """Complete IRIS system: VAE + Generator.
+    
+    For training: use train_step() which handles noise scheduling.
+    For inference: use generate() which runs the full pipeline.
+    """
+    def __init__(self, config: IRISConfig):
+        super().__init__()
+        self.config = config
+        self.vae = WaveletVAE(config)
+        self.generator = IRISGenerator(config)
+    
+    def encode(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Encode images to latent space."""
+        return self.vae.encode(images)
+    
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode latent to images."""
+        return self.vae.decode(z)
+    
+    def get_velocity_target(self, z_0: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        """Rectified flow velocity target: v = noise - z_0."""
+        return noise - z_0
+    
+    def add_noise(self, z_0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """Rectified flow forward process: z_t = (1-t)*z_0 + t*noise."""
+        t_expand = t[:, None, None, None]
+        return (1 - t_expand) * z_0 + t_expand * noise
+    
+    def sample_timesteps(self, batch_size: int, device: torch.device) -> torch.Tensor:
+        """Sample timesteps from logit-normal distribution (from SD3/RF).
+        Concentrates sampling on intermediate timesteps where learning is hardest.
+        """
+        u = torch.randn(batch_size, device=device)
+        t = torch.sigmoid(u)  # Logit-normal with mean=0, std=1
+        # Clamp to avoid t=0 and t=1
+        t = t.clamp(1e-5, 1 - 1e-5)
+        return t
+    
+    def train_step(
+        self,
+        images: torch.Tensor,
+        text_tokens: torch.Tensor,
+        num_iterations: Optional[int] = None,
+    ) -> dict:
+        """Single training step for rectified flow.
+        
+        Returns dict with loss and diagnostics.
+        """
+        B = images.shape[0]
+        device = images.device
+        
+        # Encode to latent
+        z_0, mean, logvar = self.encode(images)
+        
+        # Sample noise and timesteps
+        noise = torch.randn_like(z_0)
+        t = self.sample_timesteps(B, device)
+        
+        # Create noisy latent
+        z_t = self.add_noise(z_0, noise, t)
+        
+        # Predict velocity
+        # Randomly sample iteration count for training robustness
+        if num_iterations is None:
+            r_choices = [4, 6, 8, 10, 12]
+            r = r_choices[torch.randint(0, len(r_choices), (1,)).item()]
+        else:
+            r = num_iterations
+        
+        v_pred = self.generator(z_t, t, text_tokens, num_iterations=r)
+        v_target = self.get_velocity_target(z_0, noise)
+        
+        # SNR-weighted loss (from Rectified Flow paper)
+        # w(t) = t / (1 - t) — emphasizes high-noise timesteps
+        w = t / (1 - t + 1e-8)
+        w = w[:, None, None, None]
+        
+        # Velocity matching loss
+        velocity_loss = (w * (v_pred - v_target).pow(2)).mean()
+        
+        # VAE KL loss
+        kl_loss = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp()).mean()
+        
+        return {
+            'loss': velocity_loss + 0.001 * kl_loss,
+            'velocity_loss': velocity_loss.item(),
+            'kl_loss': kl_loss.item(),
+            'mean_t': t.mean().item(),
+        }
+    
+    @torch.no_grad()
+    def generate(
+        self,
+        text_tokens: torch.Tensor,
+        num_steps: int = 4,
+        num_iterations: int = 8,
+        cfg_scale: float = 4.0,
+        seed: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Generate images from text conditioning using Euler solver.
+        
+        Args:
+            text_tokens: [B, S, text_dim] CLIP text embeddings
+            num_steps: Number of ODE solver steps (1-50)
+            num_iterations: Core iterations per step (quality budget)
+            cfg_scale: Classifier-free guidance scale
+            seed: Random seed for reproducibility
+        """
+        B, S, _ = text_tokens.shape
+        device = text_tokens.device
+        
+        if seed is not None:
+            torch.manual_seed(seed)
+        
+        # Start from pure noise
+        z = torch.randn(B, self.config.latent_channels, 
+                        self.config.latent_spatial, self.config.latent_spatial,
+                        device=device)
+        
+        # Euler solver for rectified flow ODE: dz/dt = -v(z, t)
+        # Integrate from t=1 (noise) to t=0 (data)
+        dt = 1.0 / num_steps
+        
+        for step in range(num_steps):
+            t_val = 1.0 - step * dt
+            t = torch.full((B,), t_val, device=device)
+            
+            # Predict velocity
+            v = self.generator(z, t, text_tokens, num_iterations=num_iterations)
+            
+            # Classifier-free guidance (if cfg_scale > 1)
+            if cfg_scale > 1.0:
+                null_tokens = torch.zeros_like(text_tokens)
+                v_uncond = self.generator(z, t, null_tokens, num_iterations=num_iterations)
+                v = v_uncond + cfg_scale * (v - v_uncond)
+            
+            # Euler step: z = z - dt * v
+            z = z - dt * v
+        
+        # Decode to image
+        images = self.decode(z)
+        images = images.clamp(-1, 1)
+        return images
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+def count_parameters(model: nn.Module) -> dict:
+    """Count parameters in each component."""
+    counts = {}
+    total = 0
+    for name, module in model.named_children():
+        n = sum(p.numel() for p in module.parameters())
+        counts[name] = n
+        total += n
+    counts['total'] = total
+    
+    # Separate trainable vs frozen
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    counts['trainable'] = trainable
+    return counts
+
+
+def estimate_memory_mb(model: nn.Module, dtype=torch.float16) -> float:
+    """Estimate model memory in MB."""
+    bytes_per_param = 2 if dtype == torch.float16 else 4
+    total_params = sum(p.numel() for p in model.parameters())
+    return total_params * bytes_per_param / (1024 * 1024)
+
+
+def create_iris_small(latent_spatial: int = 32) -> IRIS:
+    """Create IRIS-Small: ~75M generator params, suitable for mobile."""
+    config = IRISConfig(
+        latent_channels=16,
+        latent_spatial=latent_spatial,
+        hidden_dim=512,
+        num_heads=8,
+        head_dim=64,
+        ffn_ratio=2.667,
+        num_prelude_blocks=2,
+        num_core_layers=4,
+        num_coda_blocks=2,
+        default_iterations=8,
+        max_iterations=16,
+        fourier_num_blocks=8,
+        sparsity_threshold=0.01,
+        recurrence_dim=256,
+        manhattan_window=16,
+        text_dim=768,
+        max_text_tokens=77,
+        patch_size=2,
+    )
+    return IRIS(config)
+
+
+def create_iris_tiny(latent_spatial: int = 32) -> IRIS:
+    """Create IRIS-Tiny: ~30M generator params, ultra-mobile."""
+    config = IRISConfig(
+        latent_channels=8,
+        latent_spatial=latent_spatial,
+        hidden_dim=384,
+        num_heads=6,
+        head_dim=64,
+        ffn_ratio=2.667,
+        num_prelude_blocks=1,
+        num_core_layers=3,
+        num_coda_blocks=1,
+        default_iterations=8,
+        max_iterations=16,
+        fourier_num_blocks=6,
+        sparsity_threshold=0.01,
+        recurrence_dim=192,
+        manhattan_window=12,
+        text_dim=768,
+        max_text_tokens=77,
+        patch_size=2,
+    )
+    return IRIS(config)
+
+
+def create_iris_base(latent_spatial: int = 32) -> IRIS:
+    """Create IRIS-Base: ~150M generator params, quality-focused."""
+    config = IRISConfig(
+        latent_channels=16,
+        latent_spatial=latent_spatial,
+        hidden_dim=768,
+        num_heads=12,
+        head_dim=64,
+        ffn_ratio=2.667,
+        num_prelude_blocks=2,
+        num_core_layers=6,
+        num_coda_blocks=2,
+        default_iterations=8,
+        max_iterations=16,
+        fourier_num_blocks=12,
+        sparsity_threshold=0.01,
+        recurrence_dim=384,
+        manhattan_window=16,
+        text_dim=768,
+        max_text_tokens=77,
+        patch_size=2,
+    )
+    return IRIS(config)
+
+
+if __name__ == "__main__":
+    print("=" * 70)
+    print("IRIS: Iterative Recurrent Image Synthesis")
+    print("=" * 70)
+    
+    # Create model variants
+    for name, create_fn in [("IRIS-Tiny", create_iris_tiny), 
+                             ("IRIS-Small", create_iris_small),
+                             ("IRIS-Base", create_iris_base)]:
+        print(f"\n{'─' * 50}")
+        print(f"  {name}")
+        print(f"{'─' * 50}")
+        model = create_fn()
+        counts = count_parameters(model)
+        mem_fp16 = estimate_memory_mb(model, torch.float16)
+        mem_fp32 = estimate_memory_mb(model, torch.float32)
+        
+        print(f"  Total params:     {counts['total']:>12,}")
+        print(f"  Trainable params: {counts['trainable']:>12,}")
+        print(f"  Memory (fp16):    {mem_fp16:>10.1f} MB")
+        print(f"  Memory (fp32):    {mem_fp32:>10.1f} MB")
+        print(f"  Components:")
+        for k, v in counts.items():
+            if k not in ('total', 'trainable'):
+                print(f"    {k:20s}: {v:>10,} ({v/counts['total']*100:.1f}%)")