Add ARCHITECTURE.md

ARCHITECTURE.md

@@ -0,0 +1,364 @@
+# IRIS: Iterative Recurrent Image Synthesis
+## A Novel Architecture for Mobile-First High-Quality Image Generation
+
+### Version 1.0 | Architecture Design Document
+
+---
+
+## 1. Executive Summary
+
+**IRIS** (Iterative Recurrent Image Synthesis) is a novel image generation architecture designed from first principles to achieve high visual quality on mobile devices (< 3-4GB RAM). It combines six key innovations drawn from cutting-edge research across multiple domains:
+
+1. **Wavelet-Frequency Latent Space** — 16× spatial compression via Haar DWT + learned VAE, operating in frequency-aware space
+2. **Recurrent Depth Core** — Shared-weight denoising block iterated N times (inspired by Huginn), achieving deep model behavior from tiny parameter count
+3. **Gated Recurrent Fourier Mixer (GRFM)** — Novel token mixing that combines RG-LRU gated recurrence with Adaptive Fourier Neural Operators, replacing O(N²) attention with O(N log N) global mixing
+4. **Manhattan Spatial Decay** — Learned per-head 2D spatial inductive bias via Manhattan distance exponential decay (from RMT)
+5. **Rectified Flow with Consistency Distillation** — Straight ODE paths for few-step generation (1-4 steps)
+6. **Adaptive Compute Budget** — Same model, variable quality: 4 iterations for mobile, 16 for quality
+
+### Target Specifications
+
+| Metric | Target | Achieved By |
+|--------|--------|------------|
+| Total Parameters | < 250M (generator) | Recurrent depth + efficient blocks |
+| RAM (inference) | < 3GB total | ~600MB model + ~400MB VAE + ~200MB text encoder + buffers |
+| Inference Steps | 1-4 | Rectified flow + consistency distillation |
+| Core Iterations | 4 (fast) / 8-16 (quality) | Recurrent depth, shared weights |
+| Image Quality | Competitive with SDXL at 512px | Frequency-aware latent + proper training |
+| Prompt Adherence | Strong | CLIP-L/14 cross-attention conditioning |
+| Training Cost | < 200 A100 GPU-hours (Stage 1) | Efficient architecture + progressive training |
+
+---
+
+## 2. Theoretical Foundations
+
+### 2.1 Why Current Approaches Fail on Mobile
+
+**Problem 1: Parameter Explosion in Transformers**
+Standard DiT/UNet architectures use independent parameters for each layer. A 24-layer DiT-XL has ~675M params. Each self-attention layer stores O(d²) params for Q,K,V projections × number of layers.
+
+**Problem 2: Quadratic Attention Complexity**
+For 512×512 images with 8× VAE downsampling: 64×64 = 4096 tokens. Self-attention requires 4096² × d operations per layer. At d=768, that's ~12.9 GFLOPS per attention layer.
+
+**Problem 3: Step Count**
+Standard diffusion requires 20-50 neural function evaluations (NFE). Even a small model × 50 steps = impractical.
+
+### 2.2 Our Solution: Mathematical Framework
+
+#### 2.2.1 Recurrent Depth as Implicit Neural ODE
+
+The key insight from Huginn (arXiv:2502.05171): a shared-weight block `R` applied iteratively defines a discrete neural ODE:
+
+```
+s_0 = P(x)                    [Prelude: encode input]
+s_{i+1} = R(s_i, c, t)        [Core: iterate with conditioning c and timestep t]  
+y = C(s_r)                    [Coda: decode output]
+```
+
+This is mathematically equivalent to an Euler discretization of:
+```
+ds/dτ = F_θ(s(τ), c, t)      where τ ∈ [0, 1], discretized into r steps
+```
+
+**Parameter efficiency**: If block R has P parameters, then r iterations give effective depth of r×L layers (where L = layers in R) using only P parameters. A 6-layer block iterated 16 times = 96 effective layers.
+
+**Connection to diffusion**: In standard diffusion, the denoiser f_θ is applied at each noise level t with the SAME parameters — this IS recurrent depth, but over the noise schedule axis. IRIS makes it recurrent over BOTH axes: noise schedule (outer loop, t) and computational depth (inner loop, τ).
+
+#### 2.2.2 Gated Recurrent Fourier Mixer (GRFM) — Novel Contribution
+
+We introduce GRFM, which processes the 2D token sequence through three parallel pathways merged multiplicatively:
+
+**Pathway 1: Fourier Global Mixing (O(N log N))**
+```
+x_fourier = IRFFT2(SoftShrink(BlockMLP(RFFT2(x))))
+```
+From AFNO: captures global structure via frequency-domain mixing. The soft-shrinkage promotes sparsity in Fourier domain (images are naturally sparse in frequency).
+
+**Pathway 2: Gated Linear Recurrence (O(N))**
+```
+a_t = σ(Λ)^(c·σ(W_a · x_t))     [decay gate, per-element]
+i_t = σ(W_x · x_t)              [input gate]
+h_t = a_t ⊙ h_{t-1} + √(1 - a_t²) ⊙ (i_t ⊙ x_t)    [RG-LRU update]
+x_recurrent = W_o · h_T
+```
+From Griffin (arXiv:2402.19427): captures sequential dependencies with O(1) state per token position. Bidirectional (forward + backward scan).
+
+**Pathway 3: Manhattan Spatial Gate**
+```
+D_{nm} = γ_head^(|x_n - x_m| + |y_n - y_m|)    [Manhattan decay matrix]
+gate = σ(W_g · x) ⊙ (D · (W_v · x))
+```
+From RMT (arXiv:2309.11523): per-head learnable spatial decay provides multi-scale locality bias.
+
+**Fusion (Novel)**:
+```
+output = LayerNorm(x_fourier ⊙ σ(gate) + x_recurrent ⊙ (1 - σ(gate)))
+```
+
+The gate adaptively selects between global Fourier features (textures, patterns) and local recurrent features (edges, fine details) based on spatial context. This is NOT a simple concatenation — it's a learned, spatially-varying interpolation.
+
+#### 2.2.3 Wavelet-Frequency Latent Space
+
+Instead of standard VAE operating on pixels, we first apply Haar DWT:
+```
+x ∈ R^{3×H×W} → DWT → y ∈ R^{12×H/2×W/2}
+```
+
+Then a lightweight VAE encoder compresses to:
+```
+z ∈ R^{C×H/8×W/8}   (effective 16× total spatial compression from original)
+```
+
+The VAE operates on wavelet coefficients, preserving frequency structure. The LL (low-low) subband carries global structure; LH, HL, HH carry directional high-frequency details. This means the latent space is inherently frequency-aware.
+
+**Benefit**: The denoiser operates on a latent that already separates structure from detail, making the learning problem easier for a small model.
+
+#### 2.2.4 Rectified Flow + Consistency Distillation
+
+**Training Phase 1 (Rectified Flow)**:
+```
+x_t = (1-t) · x_0 + t · ε           [linear interpolation]
+v_target = ε - x_0                   [velocity field]
+L = w(t) · ||v_θ(x_t, t, c) - v_target||²
+w(t) = t/(1-t+ε)                     [SNR reweighting]
+t ~ LogitNormal(0, 1)                [concentrate on hard timesteps]
+```
+
+**Training Phase 2 (Consistency Distillation)**:
+```
+f_θ(x_t, t) = c_skip(t)·x_t + c_out(t)·F_θ(x_t, t)
+L_CD = d(f_θ(x_{t_{n+1}}, t_{n+1}), f_{θ⁻}(x̂_{t_n}, t_n))
+```
+Where θ⁻ is EMA of θ. This enables 1-4 step generation.
+
+---
+
+## 3. Architecture Details
+
+### 3.1 Overall Pipeline
+
+```
+Text → CLIP-L/14 → c ∈ R^{77×768}
+                                        ┌──────────────────────────┐
+Image → Haar DWT → WaveletVAE Encode → z₀ ∈ R^{C×h×w}           │
+                                        │                          │
+                                        │ Noise schedule (RF):     │
+                                        │ z_t = (1-t)z₀ + t·ε     │
+                                        │                          │
+                                        ▼                          │
+                              ┌─────────────────┐                  │
+                              │    PRELUDE       │                  │
+                              │  (2 blocks)      │                  │
+                              │  PatchEmbed +    │                  │
+                              │  Initial mixing  │                  │
+                              └────────┬─────────┘                  │
+                                       │                           │
+                                       ▼                           │
+                              ┌─────────────────┐                  │
+                              │   CORE (shared)  │◄── Iterate r    │
+                              │   GRFM Block     │    times        │
+                              │   + FFN          │    (4-16)       │
+                              │   + adaLN-Zero   │                  │
+                              └────────┬─────────┘                  │
+                                       │                           │
+                                       ▼                           │
+                              ┌─────────────────┐                  │
+                              │     CODA         │                  │
+                              │   (2 blocks)     │                  │
+                              │   Final refine + │                  │
+                              │   Unpatchify     │                  │
+                              └────────┬─────────┘                  │
+                                       │                           │
+                                       ▼                          │
+                              v̂ = predicted velocity              │
+                              z₀_pred = z_t - t·v̂                 │
+                                        │                          │
+                                        ▼                          │
+                              WaveletVAE Decode → Haar IDWT → Image│
+                                        └──────────────────────────┘
+```
+
+### 3.2 Detailed Block Design
+
+#### Prelude (2 blocks, unique weights)
+```python
+class Prelude:
+    patch_embed: Conv2d(C_latent, D, kernel_size=2, stride=2)  # 2× spatial reduce
+    pos_embed: learned R^{(h/2 × w/2) × D}
+    blocks: [PreludeBlock × 2]
+    
+class PreludeBlock:
+    norm1 → DepthwiseSepConv3x3 → GELU → PointwiseConv → norm2 → FFN
+    # Uses conv instead of attention — cheap local feature extraction
+```
+
+#### Core (shared weights, iterated r times)
+```python
+class CoreBlock:
+    # adaLN-Zero conditioning on (timestep t, iteration i, text_global)
+    adaln_modulation: Linear(D_cond, 6*D)  # scale, shift, gate for norm1, norm2
+    
+    norm1 → GRFM → gate1 → residual
+    norm2 → CrossAttention(q=x, kv=text_tokens) → gate2 → residual  # Only 77 text tokens
+    norm3 → FFN(SiLU) → gate3 → residual
+```
+
+Cross-attention with only 77 text tokens is cheap: O(N × 77 × d) ≈ O(N·d).
+
+#### GRFM (Gated Recurrent Fourier Mixer) — The Core Innovation
+```python
+class GRFM:
+    def forward(x, spatial_shape):
+        B, N, D = x.shape
+        H, W = spatial_shape
+        x_2d = x.reshape(B, H, W, D)
+        
+        # Pathway 1: Fourier Global (O(N log N))
+        x_freq = rfft2(x_2d, dim=(1,2))
+        x_freq = block_mlp(x_freq)  # Block-diagonal MLP in freq domain
+        x_freq = soft_shrink(x_freq, lambd=self.sparsity_threshold)
+        x_fourier = irfft2(x_freq, dim=(1,2))
+        
+        # Pathway 2: Bidirectional Gated Recurrence (O(N))
+        x_flat_fwd = x  # N tokens in raster order
+        x_flat_bwd = x.flip(1)  # Reversed
+        h_fwd = gated_linear_recurrence(x_flat_fwd, self.decay_fwd, self.gate_fwd)
+        h_bwd = gated_linear_recurrence(x_flat_bwd, self.decay_bwd, self.gate_bwd)
+        x_recurrent = linear(concat(h_fwd, h_bwd.flip(1)))
+        
+        # Pathway 3: Manhattan Spatial Gate
+        manhattan_dist = compute_manhattan(H, W)  # Precomputed
+        gamma = sigmoid(self.gamma_param)  # Per-head
+        spatial_decay = gamma.pow(manhattan_dist)  # [heads, N, N] — sparse/windowed
+        x_gated = einsum('hnn,bnd->bnd', spatial_decay[:, :K, :K], value_proj(x))
+        gate = sigmoid(gate_proj(x))
+        
+        # Adaptive Fusion
+        output = x_fourier * gate + x_recurrent * (1 - gate)
+        output = output + 0.1 * x_gated  # Small residual from spatial
+        
+        return output_proj(output)
+```
+
+#### Coda (2 blocks, unique weights)
+```python
+class Coda:
+    blocks: [CodaBlock × 2]
+    unpatchify: ConvTranspose2d(D, C_latent, kernel_size=2, stride=2)
+    final_norm: LayerNorm(D)
+    
+class CodaBlock:
+    norm1 → LocalWindowAttention(window=8) → residual  # Small window, efficient
+    norm2 → FFN → residual
+```
+
+### 3.3 Parameter Budget
+
+| Component | Parameters | Notes |
+|-----------|-----------|-------|
+| WaveletVAE Encoder | ~15M | Lightweight (LiteVAE-style) |
+| WaveletVAE Decoder | ~8M | Tiny decoder (SnapGen-style) |
+| CLIP-L/14 Text Encoder | ~39M | Frozen, not counted for training |
+| Prelude (2 blocks) | ~12M | Conv-based, cheap |
+| Core Block (shared) | ~45M | GRFM + CrossAttn + FFN |
+| Coda (2 blocks) | ~15M | Local attention + FFN |
+| Embeddings/conditioning | ~3M | Time, iteration, position |
+| **Total Generator** | **~75M unique** | Core shared across iterations |
+| **Effective depth** | **75M → behaves like 400M+** | At r=8 iterations |
+| **Total system** | **~137M** | Including VAE + text encoder |
+
+### 3.4 Memory Analysis (Inference at 512×512)
+
+```
+CLIP-L/14 text encoder:     ~156 MB (fp16)
+WaveletVAE Decoder:         ~16 MB (fp16)  
+IRIS Generator:             ~150 MB (fp16)
+Latent tensor:              ~2 MB (32×32×16, fp16)
+KV cache (text cross-attn): ~12 MB
+Intermediate activations:   ~100 MB (single block, not accumulated)
+OS/framework overhead:      ~500 MB
+─────────────────────────────────────────
+Total:                      ~936 MB  ✓ (well under 3GB)
+```
+
+**Key insight**: Because Core block weights are shared, we don't accumulate layer-by-layer activations. Each iteration reuses the same memory buffer.
+
+---
+
+## 4. Training Recipe
+
+### Stage 1: Wavelet VAE Training (Standalone)
+```
+Data: ImageNet (1.2M images) + CC3M (3M images)
+Resolution: 256×256
+Objective: Reconstruction loss + KL + Perceptual (LPIPS) + Wavelet frequency loss
+Batch: 32
+LR: 1e-4, cosine decay
+Duration: ~20 GPU-hours on A100
+```
+
+### Stage 2: Class-Conditional Pretraining
+```
+Data: ImageNet 256×256 (class labels)
+Objective: Rectified Flow velocity matching
+Batch: 256
+LR: 1e-4, warmup 5000 steps, cosine decay
+Core iterations: r=8 (randomly sample r ∈ {4,6,8,10,12} for robustness)
+Duration: ~100 GPU-hours on A100
+```
+
+### Stage 3: Text-Image Alignment
+```
+Data: CC3M + CC12M (15M images with captions, re-captioned by VLM)
+Resolution: 256→512 progressive
+Objective: Rectified Flow + cross-attention on CLIP-L text tokens
+Batch: 128
+LR: 2e-5, constant
+Duration: ~200 GPU-hours on A100
+```
+
+### Stage 4: Aesthetic Fine-tuning
+```
+Data: JourneyDB + high-aesthetic LAION subset (1M images, aesthetic score > 6.0)
+Resolution: 512×512
+Batch: 64
+LR: 5e-6
+Duration: ~50 GPU-hours
+```
+
+### Stage 5: Consistency Distillation
+```
+Teacher: Trained IRIS model from Stage 4
+Student: Same architecture, initialized from teacher
+Objective: Consistency loss (CD) + optional LADD (adversarial)
+Target: 1-4 step generation
+Duration: ~30 GPU-hours
+```
+
+**Total estimated cost: ~400 A100 GPU-hours ≈ $1,600 at cloud prices**
+**Colab/Kaggle feasible**: Stage 1-2 can run on T4/A100 free tier
+
+---
+
+## 5. Novel Contributions Summary
+
+1. **GRFM (Gated Recurrent Fourier Mixer)**: First architecture to fuse Fourier global mixing, gated linear recurrence, and Manhattan spatial decay in a single differentiable block with learned gating
+2. **Recurrent Depth for Image Generation**: First application of the Huginn prelude-core-coda pattern to image generation, enabling budget-adaptive compute
+3. **Wavelet-Frequency Latent Space**: DWT preprocessing before VAE encoding preserves frequency structure in the latent space
+4. **Iteration-Aware Conditioning**: The core block receives both timestep t and iteration index i via adaLN, allowing it to learn different behavior at different depths
+5. **Dual-Axis Recurrence**: Recurrence over both noise schedule (diffusion steps) and computational depth (core iterations) — a new paradigm for efficient generation
+
+---
+
+## 6. Extensions
+
+### 6.1 Image Editing (Inpainting, Super-Resolution)
+The iterative nature of IRIS makes it natural for editing:
+- **Inpainting**: Mask latent tokens, condition core iterations on unmasked context
+- **Super-Resolution**: Encode low-res image via WaveletVAE, condition generation on LL subband
+- **Prompt-based editing**: Encode source image, modify text conditioning, run partial denoising (SDEdit-style)
+
+### 6.2 ControlNet-like Conditioning
+Add lightweight adapter to Prelude that injects spatial control signals (edges, depth, pose) into the latent, then the shared Core naturally propagates this through iterations.
+
+---