lira/model.py

"""
LiRA Model: Full Architecture

Architecture Overview (Denoising Network):
==========================================

Input: z_t (noisy latent, B x C x H x W) + t (timestep) + text_features
                    |
                    v
        ┌─────────────────────────┐
        │   Patch Embedding       │  Conv2d(C_lat, D, 1x1) - patchify
        │   + Freq Decomposition  │  Optional: Haar wavelet split
        └────────────┬────────────┘
                     │
                     v
        ┌─────────────────────────┐
        │  Latent Reasoning Loop  │  2-8 adaptive steps (learned)
        │  (generates reasoning   │  → produces reasoning conditioning
        │   conditioning vector)  │  Only ~128 dims, very cheap
        └────────────┬────────────┘
                     │ reasoning_cond + timestep_embed + text_pooled
                     │ → combined conditioning vector
                     v
        ┌─────────────────────────┐
        │  N x LiRA Blocks        │  Each block:
        │  (with HyperConnections)│    1. AdaLN conditioning
        │                         │    2. Bidirectional SSM (4-dir scan)
        │  Every K blocks:        │    3. Mix-FFN (DWConv + GLU)
        │  → GatedCrossStateFusion│    4. Hyper-connection routing
        └────────────┬────────────┘
                     │
                     v
        ┌─────────────────────────┐
        │   Final Norm + Proj     │  LayerNorm → Linear(D, C_lat)
        │   → velocity prediction │  Predicts v = ε - x_0
        └─────────────────────────┘

Model Sizes:
- LiRA-Tiny:   D=384,  N=12,  ~50M params   (for testing)
- LiRA-Small:  D=512,  N=20,  ~120M params  (mobile-optimized)
- LiRA-Base:   D=768,  N=28,  ~300M params  (quality-optimized)
- LiRA-Large:  D=1024, N=36,  ~600M params  (maximum quality)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Dict, Tuple
from einops import rearrange

from .core_modules import (
    LiRABlock,
    GatedCrossStateFusion,
    LatentReasoningLoop,
    TimestepEmbedding,
    TextProjection,
    HyperConnection,
)


# ============================================================================
# Patch Embedding for Latent Space
# ============================================================================

class LatentPatchEmbed(nn.Module):
    """
    Embeds latent space patches into model dimension.
    
    For DC-AE f32: latent is 32x32 for 1024px image, with 32 channels
    For SD3/FLUX f8: latent is 128x128 for 1024px, with 16 channels
    
    We use simple 1x1 conv (no spatial patchify) since the VAE already
    provides heavy spatial compression. Additional patching would lose
    spatial resolution in the latent space.
    
    However, for f8 VAEs (128x128 = 16384 tokens), we optionally use
    2x2 patches to reduce to 64x64 = 4096 tokens.
    """
    
    def __init__(self, in_channels: int, d_model: int, patch_size: int = 1):
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Conv2d(
            in_channels, d_model, 
            kernel_size=patch_size, stride=patch_size
        )
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, int, int]:
        """
        x: (B, C, H, W) latent features
        Returns: (B, H'*W', D), H', W'
        """
        x = self.proj(x)  # (B, D, H', W')
        B, D, H, W = x.shape
        x = rearrange(x, 'b d h w -> b (h w) d')
        x = self.norm(x)
        return x, H, W


class LatentUnpatch(nn.Module):
    """Reverse of LatentPatchEmbed: project back and reshape"""
    
    def __init__(self, d_model: int, out_channels: int, patch_size: int = 1):
        super().__init__()
        self.patch_size = patch_size
        self.out_channels = out_channels
        self.norm = nn.LayerNorm(d_model)
        
        if patch_size > 1:
            # Use pixel shuffle for upsampling
            self.proj = nn.Linear(d_model, out_channels * patch_size * patch_size)
        else:
            self.proj = nn.Linear(d_model, out_channels)
    
    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        """
        x: (B, H'*W', D)
        Returns: (B, C, H_orig, W_orig)
        """
        x = self.norm(x)
        x = self.proj(x)  # (B, H'*W', C*p*p)
        
        x = rearrange(x, 'b (h w) d -> b d h w', h=H, w=W)
        
        if self.patch_size > 1:
            x = F.pixel_shuffle(x, self.patch_size)
        
        return x


# ============================================================================
# LiRA Denoising Network
# ============================================================================

class LiRAModel(nn.Module):
    """
    LiRA: Liquid Reasoning Artisan - Main Denoising Network
    
    Novel architecture combining:
    1. State-space backbone (O(N) complexity)
    2. Latent reasoning loop (adaptive compute)
    3. Hyper-connections (dynamic layer arrangement)
    4. Gated cross-state text fusion (efficient cross-modal)
    5. Mix-FFN (local feature enhancement)
    
    Designed for mobile deployment:
    - No quadratic attention anywhere
    - All operations are O(N) in sequence length
    - Compact parameter count (<400M for Base)
    - Native 1024px via f32 VAE (32x32 = 1024 tokens)
    """
    
    # Predefined configurations
    CONFIGS = {
        'tiny': {
            'd_model': 384, 'n_blocks': 12, 'd_state': 8,
            'd_reason': 96, 'max_reason_steps': 4,
            'ffn_expand': 2.0, 'cross_every': 4,
            'hc_expansion': 2, 'num_heads': 6,
        },
        'small': {
            'd_model': 512, 'n_blocks': 20, 'd_state': 16,
            'd_reason': 128, 'max_reason_steps': 6,
            'ffn_expand': 2.5, 'cross_every': 4,
            'hc_expansion': 2, 'num_heads': 8,
        },
        'base': {
            'd_model': 768, 'n_blocks': 28, 'd_state': 16,
            'd_reason': 192, 'max_reason_steps': 8,
            'ffn_expand': 2.5, 'cross_every': 4,
            'hc_expansion': 2, 'num_heads': 12,
        },
        'large': {
            'd_model': 1024, 'n_blocks': 36, 'd_state': 16,
            'd_reason': 256, 'max_reason_steps': 8,
            'ffn_expand': 3.0, 'cross_every': 4,
            'hc_expansion': 2, 'num_heads': 16,
        },
    }
    
    def __init__(
        self,
        config_name: str = 'small',
        in_channels: int = 32,  # DC-AE f32c32 latent channels
        d_text: int = 768,      # Text encoder dimension (CLIP or small LLM)
        patch_size: int = 1,    # Patch size for latent tokens
        **kwargs
    ):
        super().__init__()
        
        # Get config
        if config_name in self.CONFIGS:
            config = {**self.CONFIGS[config_name], **kwargs}
        else:
            config = kwargs
        
        self.d_model = config['d_model']
        self.n_blocks = config['n_blocks']
        self.d_state = config['d_state']
        self.d_reason = config['d_reason']
        self.cross_every = config['cross_every']
        self.in_channels = in_channels
        
        d_cond = self.d_model  # Conditioning dimension
        
        # ====== Input Processing ======
        self.patch_embed = LatentPatchEmbed(in_channels, self.d_model, patch_size)
        self.unpatch = LatentUnpatch(self.d_model, in_channels, patch_size)
        
        # ====== Conditioning ======
        self.time_embed = TimestepEmbedding(self.d_model)
        self.text_proj = TextProjection(d_text, self.d_model)
        
        # Combine timestep + text pooled + reasoning into single conditioning vector
        self.cond_combine = nn.Sequential(
            nn.Linear(self.d_model * 3, self.d_model * 2),
            nn.SiLU(),
            nn.Linear(self.d_model * 2, self.d_model)
        )
        
        # ====== Latent Reasoning Loop ======
        self.reasoning = LatentReasoningLoop(
            self.d_model, config['d_reason'], config['max_reason_steps']
        )
        
        # ====== Main Backbone: LiRA Blocks ======
        self.blocks = nn.ModuleList()
        self.cross_fusions = nn.ModuleDict()
        
        for i in range(self.n_blocks):
            self.blocks.append(LiRABlock(
                d_model=self.d_model,
                d_cond=d_cond,
                d_state=self.d_state,
                ffn_expand=config['ffn_expand'],
                hc_expansion=config['hc_expansion'],
            ))
            
            # Add cross-modal fusion every K blocks
            if (i + 1) % self.cross_every == 0:
                self.cross_fusions[str(i)] = GatedCrossStateFusion(
                    self.d_model, self.d_model, self.d_state, config['num_heads']
                )
        
        # ====== Long Skip Connection (from U-ViT / DiM) ======
        # Connect block i with block (n_blocks - 1 - i) via learned projection
        self.n_skip = self.n_blocks // 2
        self.skip_projs = nn.ModuleList([
            nn.Linear(self.d_model * 2, self.d_model)
            for _ in range(self.n_skip)
        ])
        
        # ====== Output ======
        self.final_norm = nn.LayerNorm(self.d_model)
        self.final_adaln = nn.Sequential(
            nn.SiLU(),
            nn.Linear(d_cond, 2 * self.d_model)
        )
        nn.init.zeros_(self.final_adaln[1].weight)
        nn.init.zeros_(self.final_adaln[1].bias)
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """Careful weight initialization for training stability"""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def forward(
        self,
        z_t: torch.Tensor,           # (B, C, H, W) noisy latent
        t: torch.Tensor,              # (B,) timestep in [0, 1]
        text_features: torch.Tensor,  # (B, M, D_text) text encoder output
        text_mask: Optional[torch.Tensor] = None,  # (B, M) mask
    ) -> Tuple[torch.Tensor, Dict]:
        """
        Forward pass: predicts velocity v_t = ε - x_0
        
        Returns:
            v_pred: (B, C, H, W) predicted velocity
            info: dict with reasoning stats
        """
        B = z_t.shape[0]
        
        # ====== Embed inputs ======
        x, H, W = self.patch_embed(z_t)  # (B, N, D)
        t_emb = self.time_embed(t)        # (B, D)
        text_tokens, text_pooled = self.text_proj(text_features, text_mask)  # (B, M, D), (B, D)
        
        # ====== Latent Reasoning ======
        reason_cond, reason_info = self.reasoning(x)  # (B, D)
        
        # ====== Combine conditioning ======
        cond = self.cond_combine(torch.cat([t_emb, text_pooled, reason_cond], dim=-1))  # (B, D)
        
        # ====== Main backbone with long skip connections ======
        skip_features = []
        
        for i, block in enumerate(self.blocks):
            # Store features for skip connections (first half)
            if i < self.n_skip:
                skip_features.append(x)
            
            # Apply LiRA block
            x = block(x, cond, H, W)
            
            # Apply cross-modal fusion
            if str(i) in self.cross_fusions:
                x = self.cross_fusions[str(i)](x, text_tokens)
            
            # Apply skip connections (second half)
            if i >= self.n_skip:
                skip_idx = self.n_blocks - 1 - i
                if skip_idx < len(skip_features):
                    x = self.skip_projs[skip_idx](
                        torch.cat([x, skip_features[skip_idx]], dim=-1)
                    )
        
        # ====== Output projection ======
        shift, scale = self.final_adaln(cond).unsqueeze(1).chunk(2, dim=-1)
        x = self.final_norm(x) * (1 + scale) + shift
        
        v_pred = self.unpatch(x, H, W)  # (B, C, H_orig, W_orig)
        
        return v_pred, reason_info
    
    @torch.no_grad()
    def count_parameters(self) -> Dict[str, int]:
        """Count parameters by component"""
        counts = {}
        counts['patch_embed'] = sum(p.numel() for p in self.patch_embed.parameters())
        counts['unpatch'] = sum(p.numel() for p in self.unpatch.parameters())
        counts['time_embed'] = sum(p.numel() for p in self.time_embed.parameters())
        counts['text_proj'] = sum(p.numel() for p in self.text_proj.parameters())
        counts['reasoning'] = sum(p.numel() for p in self.reasoning.parameters())
        counts['blocks'] = sum(p.numel() for p in self.blocks.parameters())
        counts['cross_fusions'] = sum(p.numel() for p in self.cross_fusions.parameters())
        counts['skip_projs'] = sum(p.numel() for p in self.skip_projs.parameters())
        counts['conditioning'] = sum(p.numel() for p in self.cond_combine.parameters())
        counts['output'] = (
            sum(p.numel() for p in self.final_norm.parameters()) +
            sum(p.numel() for p in self.final_adaln.parameters())
        )
        counts['total'] = sum(p.numel() for p in self.parameters())
        return counts


# ============================================================================
# Tiny VAE Decoder for Mobile Deployment
# ============================================================================

class TinyVAEDecoder(nn.Module):
    """
    Ultra-lightweight VAE decoder inspired by SnapGen's tiny decoder.
    
    Key optimizations:
    1. NO attention layers (saves massive memory)
    2. Depthwise separable convolutions instead of full convolutions
    3. Minimal GroupNorm (only where needed to prevent color shift)
    4. PixelShuffle for upsampling (more efficient than transposed conv)
    
    For f32 VAE: 32x32 latent → 1024x1024 image (5 upsampling stages)
    For f8 VAE: 128x128 latent → 1024x1024 image (3 upsampling stages)
    
    Target: ~1.5M parameters, <5MB on disk
    """
    
    def __init__(
        self, 
        in_channels: int = 32,
        out_channels: int = 3,
        spatial_compression: int = 32,  # 32 for f32, 8 for f8
        base_channels: int = 64,
    ):
        super().__init__()
        
        num_upsample = int(math.log2(spatial_compression))  # 5 for f32, 3 for f8
        
        layers = []
        
        # Initial projection
        layers.append(nn.Conv2d(in_channels, base_channels, 3, padding=1))
        layers.append(nn.SiLU())
        
        # Upsampling stages - track channels carefully
        current_ch = base_channels
        for i in range(num_upsample):
            # Gradually reduce channels in later (higher-res) stages
            target_ch = max(base_channels // (2 ** max(0, i)), 16)
            
            # Depthwise separable residual block
            layers.append(SepConvBlock(current_ch, target_ch))
            current_ch = target_ch
            
            # PixelShuffle upsample (2x): needs ch*4 input, outputs ch
            layers.append(nn.Conv2d(current_ch, current_ch * 4, 3, padding=1))
            layers.append(nn.PixelShuffle(2))  # ch*4 → ch, spatial 2x
            layers.append(nn.SiLU())
            # After PixelShuffle, channels stay at current_ch
            
        # Final output
        layers.append(nn.Conv2d(current_ch, out_channels, 3, padding=1))
        layers.append(nn.Tanh())  # Output in [-1, 1]
        
        self.decoder = nn.Sequential(*layers)
    
    def forward(self, z: torch.Tensor) -> torch.Tensor:
        """
        z: (B, C_lat, H_lat, W_lat) latent
        Returns: (B, 3, H_img, W_img) decoded image
        """
        return self.decoder(z)


class SepConvBlock(nn.Module):
    """Depthwise separable convolution block"""
    
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.dwconv = nn.Conv2d(in_ch, in_ch, 3, padding=1, groups=in_ch)
        self.pwconv = nn.Conv2d(in_ch, out_ch, 1)
        self.norm = nn.GroupNorm(min(8, out_ch), out_ch)
        self.act = nn.SiLU()
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
    
    def forward(self, x):
        residual = self.skip(x)
        x = self.dwconv(x)
        x = self.pwconv(x)
        x = self.norm(x)
        x = self.act(x)
        return x + residual


# ============================================================================
# Complete LiRA Pipeline
# ============================================================================

class LiRAPipeline(nn.Module):
    """
    Complete LiRA pipeline combining:
    1. Pretrained VAE encoder (frozen) - for encoding images to latent space
    2. LiRA denoising network - the novel architecture
    3. Tiny VAE decoder - for mobile deployment
    
    During training: 
        image → VAE_encoder → z_0 → add_noise(z_0, t) → z_t → LiRA → v_pred
    
    During inference:
        noise → iterative_denoise(LiRA) → z_0 → TinyVAEDecoder → image
    """
    
    def __init__(
        self,
        config_name: str = 'small',
        latent_channels: int = 32,
        spatial_compression: int = 32,
        d_text: int = 768,
        patch_size: int = 1,
    ):
        super().__init__()
        
        self.spatial_compression = spatial_compression
        self.latent_channels = latent_channels
        
        # Denoising network
        self.denoiser = LiRAModel(
            config_name=config_name,
            in_channels=latent_channels,
            d_text=d_text,
            patch_size=patch_size,
        )
        
        # Tiny decoder for mobile inference
        self.tiny_decoder = TinyVAEDecoder(
            in_channels=latent_channels,
            spatial_compression=spatial_compression,
        )
    
    def forward(self, *args, **kwargs):
        return self.denoiser(*args, **kwargs)
    
    def count_parameters(self):
        counts = self.denoiser.count_parameters()
        counts['tiny_decoder'] = sum(p.numel() for p in self.tiny_decoder.parameters())
        counts['total_with_decoder'] = counts['total'] + counts['tiny_decoder']
        return counts


# ============================================================================
# Helper: Estimate memory usage
# ============================================================================

def estimate_memory_mb(model: nn.Module, batch_size: int = 1, 
                        img_size: int = 1024, spatial_compression: int = 32,
                        latent_channels: int = 32, dtype_bytes: int = 2):
    """Estimate inference memory usage in MB"""
    # Model parameters
    param_bytes = sum(p.numel() * dtype_bytes for p in model.parameters())
    param_mb = param_bytes / (1024 ** 2)
    
    # Latent size
    lat_h = img_size // spatial_compression
    lat_w = img_size // spatial_compression
    latent_bytes = batch_size * latent_channels * lat_h * lat_w * dtype_bytes
    
    # Intermediate activations (rough estimate: 3x latent)
    activation_bytes = latent_bytes * 3
    
    total_mb = param_mb + (latent_bytes + activation_bytes) / (1024 ** 2)
    
    return {
        'params_mb': param_mb,
        'latent_mb': latent_bytes / (1024 ** 2),
        'activation_mb': activation_bytes / (1024 ** 2),
        'total_inference_mb': total_mb,
    }