Upload liqmamba/model.py

liqmamba/model.py

@@ -0,0 +1,280 @@
+"""
+LiqMamba: Liquid-Mamba Image Generator
+
+Complete architecture that combines:
+1. SDXL VAE for encoding/decoding (pretrained, frozen)
+2. CfC-gated Mamba-2 SSD backbone with multi-directional 2D scans
+3. Flow matching objective for stable training
+4. Lipshitz regularization (physics-informed) to prevent collapse
+
+Configurations:
+- LiqMamba-Tiny:  ~8M params  (extreme lightweight)
+- LiqMamba-Small: ~25M params (Colab/Kaggle free tier target)
+- LiqMamba-Base:  ~85M params (higher quality)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+import math
+
+from .mamba2_ssd import MultiDirectionalScan, Mamba2SSDBlock
+from .cfc import CfCLayer
+
+
+class PatchEmbed(nn.Module):
+    """Convert image latents to patch tokens."""
+    def __init__(self, in_channels=4, dim=256, patch_size=1):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels, dim, patch_size, patch_size)
+    
+    def forward(self, x):
+        # x: (B, C, H, W) -> (B, dim, H, W) -> (B, H*W, dim)
+        x = self.proj(x)
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        return x, H, W
+
+
+class Unpatchify(nn.Module):
+    """Convert patch tokens back to image latents."""
+    def __init__(self, dim=256, out_channels=4, patch_size=1):
+        super().__init__()
+        self.proj = nn.Conv2d(dim, out_channels, patch_size, patch_size)
+    
+    def forward(self, x, H, W):
+        B, L, D = x.shape
+        x = x.transpose(1, 2).view(B, D, H, W)
+        return self.proj(x)
+
+
+class AdaLNModulation(nn.Module):
+    """
+    Adaptive Layer Norm modulation (from DiT).
+    Injects timestep and optional class conditioning.
+    """
+    def __init__(self, dim, cond_dim=256):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.scale_shift = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(cond_dim, dim * 6)  # scale, shift, gate x 2
+        )
+        
+    def forward(self, x, c):
+        # x: (B, L, D), c: (B, cond_dim)
+        params = self.scale_shift(c)  # (B, D*6)
+        scale1, shift1, gate1, scale2, shift2, gate2 = params.chunk(6, dim=-1)
+        
+        # Modulate
+        x = self.norm(x) * (1 + scale1.unsqueeze(1)) + shift1.unsqueeze(1)
+        x = x * gate1.unsqueeze(1)
+        return x
+
+
+class TimestepEmbedding(nn.Module):
+    """Sinusoidal timestep embedding."""
+    def __init__(self, dim, max_period=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_period = max_period
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.SiLU(),
+            nn.Linear(dim * 4, dim),
+        )
+    
+    def forward(self, t):
+        # t: (B,) float timesteps in [0,1]
+        half = self.dim // 2
+        freqs = torch.exp(-math.log(self.max_period) * 
+                          torch.arange(0, half, device=t.device).float() / half)
+        args = t.unsqueeze(-1) * freqs.unsqueeze(0)
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.dim % 2:
+            embedding = F.pad(embedding, (0, 1))
+        return self.mlp(embedding)
+
+
+class LiqMambaBlock(nn.Module):
+    """
+    Core LiqMamba block combining:
+    - AdaLN-Zero modulation
+    - Multi-directional Mamba-2 SSD scan
+    - CfC liquid state modulation
+    - Feed-forward with CfC gating
+    """
+    def __init__(self, dim, cond_dim=256, d_state=16, expand=2,
+                 scan_pattern="row_fwd", use_ffn=True):
+        super().__init__()
+        self.dim = dim
+        self.scan_pattern = scan_pattern
+        
+        # AdaLN
+        self.adaLN = AdaLNModulation(dim, cond_dim)
+        
+        # Multi-directional scan
+        self.scan = MultiDirectionalScan(dim, pattern=scan_pattern, 
+                                         d_state=d_state, expand=expand)
+        
+        # CfC liquid layer (replaces FFN for adaptive computation)
+        if use_ffn:
+            self.cfc_ffn = CfCLayer(dim, expansion_factor=2)
+        else:
+            self.cfc_ffn = nn.Identity()
+            
+        # AdaLN for FFN
+        self.adaLN_ffn = AdaLNModulation(dim, cond_dim) if use_ffn else None
+    
+    def forward(self, x, c, H, W):
+        # x: (B, H*W, dim), c: (B, cond_dim)
+        
+        # Scan with conditioning
+        x_mod = self.adaLN(x, c)
+        x = x + self.scan(x_mod, H, W)
+        
+        # CfC FFN with conditioning
+        if self.adaLN_ffn is not None:
+            x_mod2 = self.adaLN_ffn(x, c)
+            x = x + self.cfc_ffn(x_mod2)
+        
+        return x
+
+
+class LiqMamba(nn.Module):
+    """
+    LiqMamba Image Generator — Liquid Neural Network + Mamba-2 SSD
+    
+    Architecture:
+    1. Patch embed: latent (4, H, W) → tokens (H*W, dim)
+    2. Timestep + condition embedding
+    3. N stacked LiqMambaBlocks with alternating scan directions
+    4. Unpatchify: tokens → latent (4, H, W)
+    
+    Config presets:
+    - Tiny:  dim=128, depth=4  → ~8M params
+    - Small: dim=256, depth=8  → ~25M params  
+    - Base:  dim=512, depth=12 → ~85M params
+    """
+    
+    def __init__(
+        self,
+        in_channels: int = 4,          # VAE latent channels
+        out_channels: int = 4,
+        dim: int = 256,                # Hidden dimension
+        depth: int = 8,                # Number of blocks
+        cond_dim: int = 256,           # Conditioning dimension
+        d_state: int = 16,             # SSM state dimension
+        expand: int = 2,               # SSD expansion factor
+        patch_size: int = 1,
+        scan_patterns: list[str] | None = None,
+    ):
+        super().__init__()
+        
+        self.dim = dim
+        self.depth = depth
+        
+        # Scan pattern rotation (matching DiM's 4-pattern cycle)
+        if scan_patterns is None:
+            scan_patterns = ["row_fwd", "row_rev", "col_fwd", "col_rev"]
+        self.scan_patterns = scan_patterns
+        
+        # Patch embedding
+        self.patch_embed = PatchEmbed(in_channels, dim, patch_size)
+        self.unpatchify = Unpatchify(dim, out_channels, patch_size)
+        
+        # Timestep embedding
+        self.time_embed = TimestepEmbedding(cond_dim)
+        
+        # Optional class embedding
+        self.class_embed = nn.Embedding(1000, cond_dim)
+        
+        # Initial CfC layer for liquid state initialization
+        self.cfc_init = CfCLayer(dim, expansion_factor=2)
+        
+        # LiqMamba blocks
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            pattern = scan_patterns[i % len(scan_patterns)]
+            use_ffn = (i % 2 == 0)  # FFN every other block for efficiency
+            self.blocks.append(
+                LiqMambaBlock(
+                    dim=dim,
+                    cond_dim=cond_dim,
+                    d_state=d_state,
+                    expand=expand,
+                    scan_pattern=pattern,
+                    use_ffn=use_ffn,
+                )
+            )
+        
+        # Final CfC refinement layer
+        self.cfc_final = CfCLayer(dim, expansion_factor=2)
+        
+        # Final projection
+        self.final_norm = nn.LayerNorm(dim)
+        
+        # Initialize weights
+        self._init_weights()
+        
+    def _init_weights(self):
+        """Initialize with small values for stable training."""
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+    
+    def forward(self, x, t, class_labels=None, return_dict=False):
+        """
+        Args:
+            x: (B, C, H, W) latent images
+            t: (B,) float timesteps in [0, 1]
+            class_labels: (B,) optional class indices
+        Returns:
+            velocity field v(x, t) used for flow matching
+        """
+        B, C, H, W = x.shape
+        
+        # Patch embed
+        x, H_p, W_p = self.patch_embed(x)  # (B, H*W, dim)
+        
+        # Timestep conditioning
+        c = self.time_embed(t)  # (B, cond_dim)
+        if class_labels is not None:
+            c = c + self.class_embed(class_labels)
+        
+        # Initial liquid state
+        x = self.cfc_init(x)
+        
+        # LiqMamba blocks
+        for block in self.blocks:
+            x = block(x, c, H_p, W_p)
+        
+        # Final refinement
+        x = self.final_norm(x)
+        x = self.cfc_final(x)
+        
+        # Unpatchify
+        x = self.unpatchify(x, H_p, W_p)
+        
+        if return_dict:
+            return {"velocity": x}
+        return x
+    
+    def get_num_params(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+def liqmamba_tiny(**kwargs):
+    """Tiny variant: ~8M params, extreme lightweight."""
+    return LiqMamba(dim=128, depth=4, d_state=8, expand=2, **kwargs)
+
+def liqmamba_small(**kwargs):
+    """Small variant: ~25M params, Colab/Kaggle free tier target."""
+    return LiqMamba(dim=256, depth=8, d_state=16, expand=2, **kwargs)
+
+def liqmamba_base(**kwargs):
+    """Base variant: ~85M params, higher quality."""
+    return LiqMamba(dim=512, depth=12, d_state=16, expand=2, **kwargs)