Upload test_syntax.py

test_syntax.py

@@ -0,0 +1,104 @@
+"""
+Quick syntax and import test for LiquidFlow.
+Run: python test_syntax.py
+"""
+import sys
+import os
+
+# Test 1: All files parse correctly
+print("=== Test 1: Syntax Check ===")
+modules = [
+    'liquid_flow/__init__.py',
+    'liquid_flow/cfc_cell.py',
+    'liquid_flow/mamba2_ssd.py',
+    'liquid_flow/liquid_flow_block.py',
+    'liquid_flow/generator.py',
+    'liquid_flow/vae_wrapper.py',
+    'liquid_flow/physics_loss.py',
+    'train.py',
+]
+
+for module in modules:
+    with open(module, 'r') as f:
+        code = f.read()
+    compile(code, module, 'exec')
+    print(f"  ✓ {module}")
+
+# Test 2: Module imports
+print("\n=== Test 2: Import Check ===")
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+import torch.nn as nn
+
+from liquid_flow.cfc_cell import CfCCell, CfCBlock
+print("  ✓ CfC imports")
+
+from liquid_flow.mamba2_ssd import Mamba2SSD, Mamba2Block
+print("  ✓ Mamba-2 imports")
+
+from liquid_flow.liquid_flow_block import LiquidMambaBlock, LiquidFlowBackbone
+print("  ✓ LiquidFlow block imports")
+
+from liquid_flow.generator import LiquidFlowGenerator, create_liquidflow
+print("  ✓ Generator imports")
+
+from liquid_flow.vae_wrapper import TAESDWrapper, SDVAEWrapper
+print("  ✓ VAE wrapper imports")
+
+from liquid_flow.physics_loss import PhysicsRegularizer, DDIMEstimator
+print("  ✓ Physics loss imports")
+
+# Test 3: Forward pass
+print("\n=== Test 3: Forward Pass ===")
+
+# CfCCell
+cell = CfCCell(dim=64)
+x = torch.randn(2, 64)
+h = cell(x)
+assert h.shape == x.shape
+print(f"  ✓ CfCCell: {x.shape} -> {h.shape}")
+
+# Mamba2SSD
+ssd = Mamba2SSD(dim=64, d_state=8, expand=2)
+x_seq = torch.randn(2, 256, 64)
+out = ssd(x_seq)
+assert out.shape == x_seq.shape
+print(f"  ✓ Mamba2SSD: {x_seq.shape} -> {out.shape}")
+
+# LiquidMambaBlock (2D)
+lm = LiquidMambaBlock(dim=64, d_state=8, expand=2)
+x_2d = torch.randn(2, 64, 16, 16)
+out = lm(x_2d)
+assert out.shape == x_2d.shape
+print(f"  ✓ LiquidMambaBlock: {x_2d.shape} -> {out.shape}")
+
+# Full backbone
+backbone = LiquidFlowBackbone(in_channels=4, hidden_dim=64, num_stages=2, blocks_per_stage=2)
+x = torch.randn(2, 4, 32, 32)
+t = torch.tensor([500, 750])
+out = backbone(x, t)
+assert out.shape == x.shape
+print(f"  ✓ Backbone: {x.shape} -> {out.shape}")
+
+# Generator
+model = create_liquidflow(variant='tiny', image_size=128)
+x = torch.randn(2, 4, 16, 16)
+t = torch.tensor([500, 750])
+out = model(x, t)
+assert out.shape == x.shape
+print(f"  ✓ Generator: {x.shape} -> {out.shape}")
+
+# Physics loss
+physics = PhysicsRegularizer()
+x0 = torch.randn(2, 3, 32, 32)
+total, losses = physics(x0)
+print(f"  ✓ Physics Loss: total={total.item():.4f}")
+
+# Count params
+n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+print(f"\n{'='*60}")
+print(f"ALL TESTS PASSED! ✓")
+print(f"Tiny model params: {n_params:,} ({n_params/1e6:.1f}M)")
+print(f"Model compatible with Colab/Kaggle free tier")
+print(f"{'='*60}")