Fix: remove duplicate forward method in LiquidSSMBlock, clean up dead code"

liquidflow/model.py

@@ -2,11 +2,10 @@
 LiquidFlow: A Novel Liquid-SSM Flow Matching Image Generator
 v0.2.0 — Memory-optimized for Colab T4 (15GB VRAM)
 
-CHANGES from v0.1:
-- SSM scan computes per-step instead of pre-materializing (B,L,D,N) 4D tensors
-- Gradient checkpointing on all blocks (saves ~60% activation memory)
-- Liquid CfC avoids expanding h to full sequence length
-- Fixed deprecated torch.cuda.amp API
+Key fixes from v0.1:
+- SSM scan computes per-step (no 4D tensor materialization → saves ~6GB)
+- Gradient checkpointing on SSM + Liquid branches (saves ~60% activations)
+- Liquid CfC simplified to single fused projection (saves ~2GB)
 """
 
 import math
@@ -22,56 +21,37 @@ from torch.utils.checkpoint import checkpoint
 
 class LiquidCfCCell(nn.Module):
     """
-    Closed-form Continuous-depth Liquid Cell.
+    Closed-form Continuous-depth Liquid Cell (memory-optimized).
     
-    CfC solution (parallel, fast, stable):
-        gate = σ(-f_τ)
-        new_h = gate * h + (1 - gate) * f_x
-    
-    Sigmoid gating guarantees bounded dynamics — no explosion by construction.
-    
-    MEMORY FIX v0.2: Uses a single linear projection instead of two separate
-    networks + avoids expanding hidden state to (B, L, D).
+    Single fused projection instead of two separate MLP networks.
+    gate = σ(-f_τ), out = gate * h + (1 - gate) * f_x
+    Sigmoid gating guarantees bounded dynamics.
     """
     
     def __init__(self, input_dim, hidden_dim):
         super().__init__()
         self.hidden_dim = hidden_dim
-        # Single fused projection: input → (tau, state_update) 
-        # Much more memory efficient than two separate networks
         self.backbone = nn.Linear(input_dim, hidden_dim)
-        self.gate_proj = nn.Linear(hidden_dim, hidden_dim * 2)  # outputs [f_tau, f_x]
+        self.gate_proj = nn.Linear(hidden_dim, hidden_dim * 2)
         self.act = nn.Tanh()
         
     def forward(self, x):
-        """
-        x: (B, L, input_dim)
-        Returns: (B, L, hidden_dim)
-        """
-        # Project input
-        h = self.backbone(x)          # (B, L, hidden_dim)
-        h = self.act(h)
-        proj = self.gate_proj(h)       # (B, L, hidden_dim * 2)
-        f_tau, f_x = proj.chunk(2, dim=-1)
-        
-        # CfC gating: gate ∈ (0,1) by sigmoid → bounded output
+        h = self.act(self.backbone(x))
+        f_tau, f_x = self.gate_proj(h).chunk(2, dim=-1)
         gate = torch.sigmoid(-f_tau)
-        # Mix: gate * input_proj + (1-gate) * state_update
-        out = gate * h + (1.0 - gate) * f_x
-        return out
+        return gate * h + (1.0 - gate) * f_x
 
 
 # ============================================================
-# 2. SELECTIVE STATE SPACE BLOCK (Pure PyTorch Mamba-style)
+# 2. SELECTIVE STATE SPACE BLOCK (Pure PyTorch, memory-lean)
 # ============================================================
 
 class SelectiveSSM(nn.Module):
     """
-    Selective SSM in pure PyTorch — memory-optimized.
+    Selective SSM — memory-optimized scan.
     
-    MEMORY FIX v0.2: The scan loop computes discretized A,B per-step
-    instead of pre-materializing (B, L, d_inner, d_state) 4D tensors.
-    This reduces peak memory from O(B*L*D*N) to O(B*D*N).
+    Per-step discretization inside loop avoids materializing
+    (B, L, d_inner, d_state) 4D tensors. Peak memory: O(B*D*N) not O(B*L*D*N).
     """
     
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
@@ -81,7 +61,6 @@ class SelectiveSSM(nn.Module):
         self.d_inner = int(d_model * expand)
         
         self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
-        
         self.conv1d = nn.Conv1d(
             self.d_inner, self.d_inner, d_conv,
             padding=d_conv - 1, groups=self.d_inner, bias=True,
@@ -99,65 +78,44 @@ class SelectiveSSM(nn.Module):
             dt_init = torch.exp(
                 torch.rand(self.d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
             )
-            inv_dt = dt_init + torch.log(-torch.expm1(-dt_init))
-            self.dt_proj.bias.copy_(inv_dt)
+            self.dt_proj.bias.copy_(dt_init + torch.log(-torch.expm1(-dt_init)))
     
     def forward(self, x):
-        B, L, D = x.shape
-        
+        B, L, _ = x.shape
         xz = self.in_proj(x)
         x_inner, z = xz.chunk(2, dim=-1)
         
-        x_conv = self.conv1d(x_inner.transpose(1, 2))[:, :, :L].transpose(1, 2)
-        x_conv = F.silu(x_conv)
+        x_conv = F.silu(self.conv1d(x_inner.transpose(1, 2))[:, :, :L].transpose(1, 2))
         
         x_ssm = self.x_proj(x_conv)
         B_sel = x_ssm[:, :, :self.d_state]
         C_sel = x_ssm[:, :, self.d_state:2*self.d_state]
-        dt = x_ssm[:, :, -1:]
-        dt = F.softplus(self.dt_proj(dt))
-        
-        A = -torch.exp(self.A_log)  # (d_inner, d_state)
+        dt = F.softplus(self.dt_proj(x_ssm[:, :, -1:]))
         
-        y = self._selective_scan_lean(x_conv, dt, A, B_sel, C_sel)
+        A = -torch.exp(self.A_log)
+        y = self._scan(x_conv, dt, A, B_sel, C_sel)
         
         y = y + x_conv * self.D.unsqueeze(0).unsqueeze(0)
-        y = y * F.silu(z)
-        return self.out_proj(y)
+        return self.out_proj(y * F.silu(z))
     
-    def _selective_scan_lean(self, x, dt, A, B, C):
-        """
-        Memory-lean selective scan.
-        Computes discretization per-step inside the loop to avoid
-        materializing the full (B, L, d_inner, d_state) tensors.
-        
-        Peak memory: O(B * d_inner * d_state) instead of O(B * L * d_inner * d_state).
-        """
+    def _scan(self, x, dt, A, B, C):
+        """Memory-lean sequential scan — no 4D tensor allocation."""
         B_batch, L, d_inner = x.shape
-        d_state = A.shape[1]
         
-        h = torch.zeros(B_batch, d_inner, d_state, device=x.device, dtype=x.dtype)
+        h = torch.zeros(B_batch, d_inner, self.d_state, device=x.device, dtype=x.dtype)
         ys = []
         
         for i in range(L):
-            # Per-step discretization — no 4D tensor allocation
-            dt_i = dt[:, i, :]                    # (B, d_inner)
-            B_i = B[:, i, :]                      # (B, d_state)
-            C_i = C[:, i, :]                      # (B, d_state)
-            x_i = x[:, i, :]                      # (B, d_inner)
-            
-            # dA_i = exp(dt_i * A) — broadcast: (B, d_inner, 1) * (1, d_inner, d_state)
-            dA_i = torch.exp(dt_i.unsqueeze(-1) * A.unsqueeze(0))  # (B, d_inner, d_state)
+            dt_i = dt[:, i]        # (B, d_inner)
+            B_i  = B[:, i]         # (B, d_state)
+            C_i  = C[:, i]         # (B, d_state)
+            x_i  = x[:, i]         # (B, d_inner)
             
-            # dB_i * x_i: (B, d_inner, 1) * (B, 1, d_state) * (B, d_inner, 1)
-            dBx_i = dt_i.unsqueeze(-1) * B_i.unsqueeze(1) * x_i.unsqueeze(-1)  # (B, d_inner, d_state)
+            dA_i  = torch.exp(dt_i.unsqueeze(-1) * A.unsqueeze(0))           # (B, d_inner, d_state)
+            dBx_i = dt_i.unsqueeze(-1) * B_i.unsqueeze(1) * x_i.unsqueeze(-1) # (B, d_inner, d_state)
             
-            # Recurrence
             h = dA_i * h + dBx_i
-            
-            # Output
-            y_i = (h * C_i.unsqueeze(1)).sum(-1)  # (B, d_inner)
-            ys.append(y_i)
+            ys.append((h * C_i.unsqueeze(1)).sum(-1))
         
         return torch.stack(ys, dim=1)
 
@@ -169,22 +127,15 @@ class SelectiveSSM(nn.Module):
 def create_scan_patterns(H, W):
     total = H * W
     indices = torch.arange(total)
-    
-    row_major = indices.clone()
-    row_major_rev = indices.flip(0)
-    
     grid = indices.view(H, W)
-    col_major = grid.t().contiguous().view(-1)
     
-    zigzag = []
-    for i in range(H):
-        row = grid[i]
-        if i % 2 == 1:
-            row = row.flip(0)
-        zigzag.append(row)
-    zigzag = torch.cat(zigzag)
+    patterns = [
+        indices.clone(),                                          # row-major
+        indices.flip(0),                                          # reversed
+        grid.t().contiguous().view(-1),                           # column-major
+        torch.cat([grid[i].flip(0) if i % 2 else grid[i] for i in range(H)]),  # zigzag
+    ]
     
-    patterns = [row_major, row_major_rev, col_major, zigzag]
     inverse_patterns = []
     for p in patterns:
         inv = torch.zeros_like(p)
@@ -195,86 +146,33 @@ def create_scan_patterns(H, W):
 
 
 # ============================================================
-# 4. LIQUID-SSM BLOCK with gradient checkpointing
+# 4. LIQUID-SSM BLOCK (with gradient checkpointing)
 # ============================================================
 
 class LiquidSSMBlock(nn.Module):
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dropout=0.0):
         super().__init__()
-        
         self.norm1 = nn.LayerNorm(d_model)
         self.ssm = SelectiveSSM(d_model, d_state, d_conv, expand)
-        
         self.norm2 = nn.LayerNorm(d_model)
         self.liquid = LiquidCfCCell(d_model, d_model)
-        
         self.norm3 = nn.LayerNorm(d_model)
         self.ff = nn.Sequential(
-            nn.Linear(d_model, d_model * 4),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(d_model * 4, d_model),
-            nn.Dropout(dropout),
+            nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(d_model * 4, d_model), nn.Dropout(dropout),
         )
-        
         self.mix_alpha = nn.Parameter(torch.tensor(0.5))
-        
-    def _inner_forward(self, x, x_scanned):
-        """Inner forward for gradient checkpointing."""
-        ssm_out = self.ssm(self.norm1(x_scanned))
-        liquid_out = self.liquid(self.norm2(x))
-        
-        alpha = torch.sigmoid(self.mix_alpha)
-        mixed = alpha * ssm_out + (1.0 - alpha) * liquid_out
-        return mixed
-        
-    def forward(self, x, scan_idx=None, unscan_idx=None):
-        if scan_idx is not None:
-            x_scanned = x[:, scan_idx]
-        else:
-            x_scanned = x
-        
-        # Gradient checkpointing: recompute forward during backward
-        # to save activation memory
-        if self.training and x.requires_grad:
-            mixed = checkpoint(self._inner_forward, x, x_scanned, use_reentrant=False)
-        else:
-            mixed = self._inner_forward(x, x_scanned)
-        
-        # Unscan the SSM output portion
-        # Note: mixed already contains both SSM (scanned) and Liquid (unscanned)
-        # The SSM part was scanned, so we need to unscan the full mixed output
-        # Actually since we mix before unscanning, and liquid operates on original order,
-        # we need to handle this differently. Let's unscan only the SSM part.
-        # FIXED: unscan happens inside _inner_forward is wrong — we need it outside.
-        # Re-architect: unscan the SSM output before mixing.
-        
-        # Actually the mixing happens inside _inner_forward on the scanned SSM output.
-        # The Liquid branch sees original order. The mix combines them.
-        # For the SSM branch to be correct, we should unscan its output before mixing.
-        # Let me fix this properly:
-        
-        # The above checkpoint call passes x_scanned which is in scan order.
-        # SSM processes it in scan order and outputs in scan order.
-        # We need to unscan before mixing with Liquid (which is in original order).
-        # This is handled by splitting the logic:
-        
-        if unscan_idx is not None:
-            # We need to redo this without checkpoint for correct unscan
-            # Actually let's restructure to handle unscan inside
-            pass
-        
-        x = x + mixed
-        x = x + self.ff(self.norm3(x))
-        return x
-    
+
+    def _ssm_forward(self, x_scanned):
+        return self.ssm(self.norm1(x_scanned))
+
+    def _liquid_forward(self, x):
+        return self.liquid(self.norm2(x))
+
     def forward(self, x, scan_idx=None, unscan_idx=None):
-        """Clean forward with proper scan/unscan and checkpointing."""
-        if scan_idx is not None:
-            x_scanned = x[:, scan_idx]
-        else:
-            x_scanned = x
+        x_scanned = x[:, scan_idx] if scan_idx is not None else x
         
+        # Gradient checkpointing: recompute during backward → saves activation memory
         if self.training and x.requires_grad:
             ssm_out = checkpoint(self._ssm_forward, x_scanned, use_reentrant=False)
             liquid_out = checkpoint(self._liquid_forward, x, use_reentrant=False)
@@ -287,45 +185,34 @@ class LiquidSSMBlock(nn.Module):
             ssm_out = ssm_out[:, unscan_idx]
         
         alpha = torch.sigmoid(self.mix_alpha)
-        mixed = alpha * ssm_out + (1.0 - alpha) * liquid_out
-        
-        x = x + mixed
+        x = x + alpha * ssm_out + (1.0 - alpha) * liquid_out
         x = x + self.ff(self.norm3(x))
         return x
-    
-    def _ssm_forward(self, x_scanned):
-        return self.ssm(self.norm1(x_scanned))
-    
-    def _liquid_forward(self, x):
-        return self.liquid(self.norm2(x))
 
 
 # ============================================================
-# 5. TIMESTEP & CONDITION EMBEDDINGS
+# 5. EMBEDDINGS
 # ============================================================
 
 class SinusoidalPosEmb(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.dim = dim
-    
     def forward(self, t):
-        half_dim = self.dim // 2
-        emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
+        half = self.dim // 2
+        emb = math.log(10000) / (half - 1)
+        emb = torch.exp(torch.arange(half, device=t.device) * -emb)
         emb = t.unsqueeze(-1) * emb.unsqueeze(0)
         return torch.cat([emb.sin(), emb.cos()], dim=-1)
 
-
 class AdaptiveLayerNorm(nn.Module):
     def __init__(self, d_model, cond_dim):
         super().__init__()
         self.norm = nn.LayerNorm(d_model, elementwise_affine=False)
         self.proj = nn.Sequential(nn.SiLU(), nn.Linear(cond_dim, d_model * 2))
-    
     def forward(self, x, cond):
-        scale, shift = self.proj(cond).chunk(2, dim=-1)
-        return self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        s, b = self.proj(cond).chunk(2, dim=-1)
+        return self.norm(x) * (1 + s.unsqueeze(1)) + b.unsqueeze(1)
 
 
 # ============================================================
@@ -333,10 +220,8 @@ class AdaptiveLayerNorm(nn.Module):
 # ============================================================
 
 class LiquidFlowNet(nn.Module):
-    def __init__(
-        self, img_size=128, patch_size=4, in_channels=3, d_model=256,
-        depth=8, d_state=16, d_conv=4, expand=2, dropout=0.0, num_classes=0,
-    ):
+    def __init__(self, img_size=128, patch_size=4, in_channels=3, d_model=256,
+                 depth=8, d_state=16, d_conv=4, expand=2, dropout=0.0, num_classes=0):
         super().__init__()
         self.img_size = img_size
         self.patch_size = patch_size
@@ -350,28 +235,16 @@ class LiquidFlowNet(nn.Module):
         self.num_patches = self.num_patches_h * self.num_patches_w
         self.patch_dim = in_channels * patch_size * patch_size
         
-        self.patch_embed = nn.Sequential(
-            nn.Linear(self.patch_dim, d_model), nn.LayerNorm(d_model),
-        )
+        self.patch_embed = nn.Sequential(nn.Linear(self.patch_dim, d_model), nn.LayerNorm(d_model))
         self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, d_model) * 0.02)
-        
         self.time_embed = nn.Sequential(
-            SinusoidalPosEmb(d_model),
-            nn.Linear(d_model, d_model * 4), nn.GELU(),
-            nn.Linear(d_model * 4, d_model),
+            SinusoidalPosEmb(d_model), nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model),
         )
-        
         self.class_embed = nn.Embedding(num_classes, d_model) if num_classes > 0 else None
         
-        self.blocks = nn.ModuleList([
-            LiquidSSMBlock(d_model, d_state, d_conv, expand, dropout) for _ in range(depth)
-        ])
-        self.adaln_blocks = nn.ModuleList([
-            AdaptiveLayerNorm(d_model, d_model) for _ in range(depth)
-        ])
-        self.skip_projs = nn.ModuleList([
-            nn.Linear(d_model * 2, d_model) for _ in range(depth // 2)
-        ])
+        self.blocks = nn.ModuleList([LiquidSSMBlock(d_model, d_state, d_conv, expand, dropout) for _ in range(depth)])
+        self.adaln_blocks = nn.ModuleList([AdaptiveLayerNorm(d_model, d_model) for _ in range(depth)])
+        self.skip_projs = nn.ModuleList([nn.Linear(d_model * 2, d_model) for _ in range(depth // 2)])
         
         self.final_norm = nn.LayerNorm(d_model)
         self.final_proj = nn.Linear(d_model, self.patch_dim)
@@ -384,45 +257,34 @@ class LiquidFlowNet(nn.Module):
         
         self.pre_conv = nn.Conv2d(d_model, d_model, 3, padding=1, groups=d_model)
         self.post_conv = nn.Conv2d(d_model, d_model, 3, padding=1, groups=d_model)
-        
         self._init_weights()
     
     def _init_weights(self):
         for m in self.modules():
             if isinstance(m, nn.Linear):
                 nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
+                if m.bias is not None: nn.init.zeros_(m.bias)
             elif isinstance(m, (nn.Conv2d, nn.Conv1d)):
                 nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
+                if m.bias is not None: nn.init.zeros_(m.bias)
         nn.init.zeros_(self.final_proj.weight)
         nn.init.zeros_(self.final_proj.bias)
     
     def patchify(self, x):
         B, C, H, W = x.shape
         p = self.patch_size
-        x = x.unfold(2, p, p).unfold(3, p, p)
-        x = x.contiguous().view(B, C, self.num_patches_h, self.num_patches_w, p * p)
-        x = x.permute(0, 2, 3, 1, 4).contiguous().view(B, self.num_patches, self.patch_dim)
-        return x
+        return x.unfold(2,p,p).unfold(3,p,p).contiguous().view(B,C,self.num_patches_h,self.num_patches_w,p*p).permute(0,2,3,1,4).contiguous().view(B,self.num_patches,self.patch_dim)
     
     def unpatchify(self, x):
-        B = x.shape[0]
-        p = self.patch_size
-        x = x.view(B, self.num_patches_h, self.num_patches_w, self.in_channels, p, p)
-        x = x.permute(0, 3, 1, 4, 2, 5).contiguous()
-        return x.view(B, self.in_channels, self.num_patches_h * p, self.num_patches_w * p)
+        B = x.shape[0]; p = self.patch_size
+        return x.view(B,self.num_patches_h,self.num_patches_w,self.in_channels,p,p).permute(0,3,1,4,2,5).contiguous().view(B,self.in_channels,self.num_patches_h*p,self.num_patches_w*p)
     
     def forward(self, x, t, class_label=None):
         B = x.shape[0]
-        
         tokens = self.patch_embed(self.patchify(x)) + self.pos_embed
         
-        # Pre-conv for local structure
-        h2d = tokens.view(B, self.num_patches_h, self.num_patches_w, self.d_model).permute(0, 3, 1, 2)
-        tokens = self.pre_conv(h2d).permute(0, 2, 3, 1).contiguous().view(B, self.num_patches, self.d_model)
+        h2d = tokens.view(B, self.num_patches_h, self.num_patches_w, self.d_model).permute(0,3,1,2)
+        tokens = self.pre_conv(h2d).permute(0,2,3,1).contiguous().view(B, self.num_patches, self.d_model)
         
         t_emb = self.time_embed(t)
         if self.class_embed is not None and class_label is not None:
@@ -432,23 +294,15 @@ class LiquidFlowNet(nn.Module):
         for i, (block, adaln) in enumerate(zip(self.blocks, self.adaln_blocks)):
             tokens = adaln(tokens, t_emb)
             si = i % self.num_scan_patterns
-            scan_idx = getattr(self, f'scan_{si}')
-            unscan_idx = getattr(self, f'unscan_{si}')
-            
-            if i < self.depth // 2:
-                skips.append(tokens)
-            
-            tokens = block(tokens, scan_idx, unscan_idx)
-            
+            if i < self.depth // 2: skips.append(tokens)
+            tokens = block(tokens, getattr(self, f'scan_{si}'), getattr(self, f'unscan_{si}'))
             if i >= self.depth // 2:
                 skip_idx = self.depth - 1 - i
                 if skip_idx < len(skips):
                     tokens = self.skip_projs[skip_idx](torch.cat([tokens, skips[skip_idx]], dim=-1))
         
-        # Post-conv
-        h2d = tokens.view(B, self.num_patches_h, self.num_patches_w, self.d_model).permute(0, 3, 1, 2)
-        tokens = self.post_conv(h2d).permute(0, 2, 3, 1).contiguous().view(B, self.num_patches, self.d_model)
-        
+        h2d = tokens.view(B, self.num_patches_h, self.num_patches_w, self.d_model).permute(0,3,1,2)
+        tokens = self.post_conv(h2d).permute(0,2,3,1).contiguous().view(B, self.num_patches, self.d_model)
         return self.unpatchify(self.final_proj(self.final_norm(tokens)))
     
     def count_params(self):
@@ -460,51 +314,22 @@ class LiquidFlowNet(nn.Module):
 # ============================================================
 
 def liquidflow_tiny(img_size=128, num_classes=0):
-    """~5M params — Colab free tier, mobile deployment"""
-    return LiquidFlowNet(
-        img_size=img_size, patch_size=4, in_channels=3,
-        d_model=192, depth=6, d_state=8, d_conv=4, expand=2,
-        num_classes=num_classes,
-    )
+    return LiquidFlowNet(img_size=img_size, patch_size=4, d_model=192, depth=6, d_state=8, expand=2, num_classes=num_classes)
 
 def liquidflow_small(img_size=128, num_classes=0):
-    """~12M params — production 128×128"""
-    return LiquidFlowNet(
-        img_size=img_size, patch_size=4, in_channels=3,
-        d_model=256, depth=8, d_state=16, d_conv=4, expand=2,
-        num_classes=num_classes,
-    )
+    return LiquidFlowNet(img_size=img_size, patch_size=4, d_model=256, depth=8, d_state=16, expand=2, num_classes=num_classes)
 
 def liquidflow_base(img_size=256, num_classes=0):
-    """~25M params — 256×256"""
-    return LiquidFlowNet(
-        img_size=img_size, patch_size=8, in_channels=3,
-        d_model=384, depth=10, d_state=16, d_conv=4, expand=2,
-        num_classes=num_classes,
-    )
+    return LiquidFlowNet(img_size=img_size, patch_size=8, d_model=384, depth=10, d_state=16, expand=2, num_classes=num_classes)
 
 def liquidflow_512(img_size=512, num_classes=0):
-    """~25M params — 512×512"""
-    return LiquidFlowNet(
-        img_size=img_size, patch_size=16, in_channels=3,
-        d_model=384, depth=10, d_state=16, d_conv=4, expand=2,
-        num_classes=num_classes,
-    )
+    return LiquidFlowNet(img_size=img_size, patch_size=16, d_model=384, depth=10, d_state=16, expand=2, num_classes=num_classes)
 
 
 if __name__ == "__main__":
-    device = torch.device("cpu")
-    for name, factory in [
-        ("tiny-128", lambda: liquidflow_tiny(128)),
-        ("small-128", lambda: liquidflow_small(128)),
-        ("base-256", lambda: liquidflow_base(256)),
-        ("512", lambda: liquidflow_512(512)),
-    ]:
-        model = factory().to(device)
-        print(f"\n{name}: {model.count_params()/1e6:.2f}M params")
-        B = 2
-        x = torch.randn(B, 3, model.img_size, model.img_size)
-        t = torch.rand(B)
-        v = model(x, t)
-        print(f"  {x.shape} → {v.shape} ✓")
-        assert v.shape == x.shape
+    for name, factory in [("tiny-128", lambda: liquidflow_tiny(128)), ("small-128", lambda: liquidflow_small(128))]:
+        m = factory()
+        print(f"{name}: {m.count_params()/1e6:.1f}M params")
+        x = torch.randn(2, 3, m.img_size, m.img_size)
+        v = m(x, torch.rand(2))
+        print(f"  {x.shape} → {v.shape} ✓"); assert v.shape == x.shape