Upload validate.py

validate.py

@@ -0,0 +1,336 @@
+#!/usr/bin/env python3
+"""
+LiquidDiffusion — Self-Contained Validation Script
+
+Run this to verify everything works before training:
+    python validate.py
+
+Tests:
+1. Model construction at all scales
+2. Forward pass at multiple resolutions  
+3. Backward pass and gradient flow
+4. 20-step training stability with random data
+5. Sampling with Euler ODE
+6. VRAM estimation
+7. Full trainer pipeline
+"""
+
+import sys
+import math
+import time
+import copy
+
+print("=" * 70)
+print("LiquidDiffusion Validation Suite")
+print("=" * 70)
+
+# Check imports
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    print(f"✓ PyTorch {torch.__version__}")
+except ImportError:
+    print("✗ PyTorch not installed. Run: pip install torch torchvision")
+    sys.exit(1)
+
+try:
+    from torchvision.utils import save_image
+    print("✓ torchvision")
+except ImportError:
+    print("✗ torchvision not installed. Run: pip install torchvision")
+    sys.exit(1)
+
+# Import our modules
+try:
+    from liquid_diffusion.model import (
+        LiquidDiffusionUNet, liquid_diffusion_tiny,
+        liquid_diffusion_small, liquid_diffusion_base,
+        SinusoidalTimeEmbedding, ParallelCfCBlock, AdaLN,
+    )
+    print("✓ liquid_diffusion.model")
+except ImportError as e:
+    print(f"✗ Failed to import model: {e}")
+    print("  Make sure you're in the liquid-diffusion directory")
+    sys.exit(1)
+
+try:
+    from liquid_diffusion.trainer import RectifiedFlowTrainer, get_cosine_schedule_with_warmup
+    print("✓ liquid_diffusion.trainer")
+except ImportError as e:
+    print(f"✗ Failed to import trainer: {e}")
+    sys.exit(1)
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"\nDevice: {device}")
+if device == 'cuda':
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+    print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+
+all_passed = True
+test_num = 0
+
+def test(name):
+    global test_num
+    test_num += 1
+    print(f"\n--- Test {test_num}: {name} ---")
+
+def fail(msg):
+    global all_passed
+    all_passed = False
+    print(f"  ✗ FAIL: {msg}")
+
+def ok(msg):
+    print(f"  ✓ {msg}")
+
+
+# =========================================================================
+test("Model Construction & Parameter Count")
+# =========================================================================
+for name, factory in [("tiny", liquid_diffusion_tiny), ("small", liquid_diffusion_small), ("base", liquid_diffusion_base)]:
+    try:
+        m = factory()
+        total, trainable = m.count_params()
+        ok(f"{name:8s}: {total:>12,} params ({total/1e6:.1f}M)")
+        del m
+    except Exception as e:
+        fail(f"{name}: {e}")
+
+# =========================================================================
+test("Forward Pass (multiple resolutions)")
+# =========================================================================
+model = liquid_diffusion_tiny()
+for res in [32, 64, 128]:
+    try:
+        x = torch.randn(2, 3, res, res)
+        t = torch.rand(2)
+        out = model(x, t)
+        assert out.shape == x.shape, f"Shape mismatch: {out.shape} vs {x.shape}"
+        assert not torch.isnan(out).any(), "NaN in output"
+        assert not torch.isinf(out).any(), "Inf in output"
+        ok(f"{res}x{res}: output shape {out.shape}, range [{out.min():.4f}, {out.max():.4f}]")
+    except Exception as e:
+        fail(f"{res}x{res}: {e}")
+
+# =========================================================================
+test("Backward Pass (gradient flow)")
+# =========================================================================
+model = liquid_diffusion_tiny()
+x = torch.randn(2, 3, 64, 64, requires_grad=False)
+t = torch.rand(2)
+out = model(x, t)
+loss = out.mean()
+loss.backward()
+
+total_params = 0
+params_with_grad = 0
+nan_grads = 0
+zero_grads = 0
+for name_p, p in model.named_parameters():
+    total_params += 1
+    if p.grad is not None:
+        params_with_grad += 1
+        if torch.isnan(p.grad).any():
+            nan_grads += 1
+        if p.grad.abs().max() == 0:
+            zero_grads += 1
+
+if nan_grads > 0:
+    fail(f"NaN gradients in {nan_grads}/{total_params} parameters")
+elif params_with_grad == 0:
+    fail("No parameters received gradients")
+else:
+    ok(f"{params_with_grad}/{total_params} params have gradients, {nan_grads} NaN, {zero_grads} zero-only")
+
+# Check gradient magnitude distribution
+grad_maxes = [p.grad.abs().max().item() for p in model.parameters() if p.grad is not None]
+ok(f"Gradient |max| range: [{min(grad_maxes):.2e}, {max(grad_maxes):.2e}]")
+
+# =========================================================================
+test("Training Stability (20 steps, random data)")
+# =========================================================================
+model = liquid_diffusion_tiny()
+optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
+
+losses = []
+for step in range(20):
+    model.train()
+    x0 = torch.randn(4, 3, 64, 64)
+    x1 = torch.randn_like(x0)
+    t_val = torch.rand(4)
+    t_expand = t_val[:, None, None, None]
+    x_t = (1 - t_expand) * x0 + t_expand * x1
+    v_target = x1 - x0
+    
+    v_pred = model(x_t, t_val)
+    loss = F.mse_loss(v_pred, v_target)
+    
+    optimizer.zero_grad()
+    loss.backward()
+    gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    optimizer.step()
+    
+    losses.append(loss.item())
+    if step % 5 == 0:
+        print(f"    Step {step:3d}: loss={loss.item():.4f}, grad_norm={gn.item():.4f}")
+
+stable = all(not math.isnan(l) and not math.isinf(l) for l in losses)
+not_exploding = max(losses) < 100
+
+if stable:
+    ok(f"No NaN/Inf in any of {len(losses)} steps")
+else:
+    fail("NaN or Inf detected in loss")
+
+if not_exploding:
+    ok(f"Loss range: [{min(losses):.4f}, {max(losses):.4f}]")
+else:
+    fail(f"Loss exploded: max={max(losses):.4f}")
+
+# =========================================================================
+test("Sampling (Euler ODE, 10 steps)")
+# =========================================================================
+model.eval()
+with torch.no_grad():
+    z = torch.randn(2, 3, 64, 64)
+    num_steps = 10
+    dt = 1.0 / num_steps
+    for i in range(num_steps, 0, -1):
+        t_s = torch.full((2,), i / num_steps)
+        v = model(z, t_s)
+        z = z - v * dt
+    z = z.clamp(-1, 1)
+    
+    if torch.isnan(z).any():
+        fail("NaN in generated samples")
+    elif torch.isinf(z).any():
+        fail("Inf in generated samples")
+    else:
+        ok(f"Shape: {z.shape}, range: [{z.min():.3f}, {z.max():.3f}], "
+           f"mean: {z.mean():.4f}, std: {z.std():.4f}")
+
+# =========================================================================
+test("Timestep Sensitivity")
+# =========================================================================
+model.eval()
+x = torch.randn(1, 3, 64, 64)
+outputs = {}
+for t_val in [0.01, 0.25, 0.5, 0.75, 0.99]:
+    with torch.no_grad():
+        out = model(x, torch.tensor([t_val]))
+    outputs[t_val] = out
+    print(f"    t={t_val:.2f}: mean={out.mean():.6f}, std={out.std():.6f}, |max|={out.abs().max():.6f}")
+
+# Check that different timesteps give different outputs
+diff_01_099 = (outputs[0.01] - outputs[0.99]).abs().mean().item()
+if diff_01_099 > 1e-6:
+    ok(f"Timestep affects output (mean diff t=0.01 vs t=0.99: {diff_01_099:.6f})")
+else:
+    fail(f"Timestep has no effect on output (diff={diff_01_099:.10f})")
+
+# =========================================================================
+test("Full Trainer Pipeline (CPU, 5 steps)")
+# =========================================================================
+model = liquid_diffusion_tiny()
+
+trainer = RectifiedFlowTrainer(
+    model=model,
+    lr=1e-4,
+    device='cpu',
+    use_amp=False,  # CPU doesn't support AMP
+    time_sampling='logit_normal',
+)
+
+for step in range(5):
+    x0 = torch.randn(2, 3, 64, 64)
+    metrics = trainer.train_step(x0)
+    if step == 0:
+        print(f"    Step {step}: loss={metrics['loss']:.4f}, grad_norm={metrics['grad_norm']:.4f}")
+
+if math.isnan(metrics['loss']):
+    fail("Trainer produced NaN loss")
+else:
+    ok(f"Trainer works: final loss={metrics['loss']:.4f}, step={trainer.step}")
+
+# Test sampling
+try:
+    samples = trainer.sample(batch_size=1, image_size=64, num_steps=5, use_ema=True)
+    if torch.isnan(samples).any():
+        fail("Trainer sampling produced NaN")
+    else:
+        ok(f"Sampling works: shape={samples.shape}, range=[{samples.min():.3f}, {samples.max():.3f}]")
+except Exception as e:
+    fail(f"Sampling failed: {e}")
+
+# Test checkpoint save/load
+try:
+    import tempfile, os
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = os.path.join(tmpdir, 'test_ckpt.pt')
+        trainer.save_checkpoint(ckpt_path)
+        
+        # Create new trainer and load
+        model2 = liquid_diffusion_tiny()
+        trainer2 = RectifiedFlowTrainer(model2, lr=1e-4, device='cpu', use_amp=False)
+        trainer2.load_checkpoint(ckpt_path)
+        
+        assert trainer2.step == trainer.step, f"Step mismatch: {trainer2.step} vs {trainer.step}"
+        ok(f"Checkpoint save/load works (step={trainer2.step})")
+except Exception as e:
+    fail(f"Checkpoint save/load failed: {e}")
+
+# =========================================================================
+test("Architecture Properties")
+# =========================================================================
+m = liquid_diffusion_tiny()
+total_blocks = (sum(len(s) for s in m.encoder_blocks) + 
+                len(m.bottleneck) + 
+                sum(len(s) for s in m.decoder_blocks))
+
+# Count attention layers (should be 0)
+attention_count = 0
+for name_m, module in m.named_modules():
+    if 'attention' in name_m.lower() or 'attn' in name_m.lower():
+        attention_count += 1
+
+ok(f"Attention layers: {attention_count} (should be 0)")
+ok(f"LiquidCfC blocks: {total_blocks}")
+ok(f"Training: Rectified Flow (MSE velocity)")
+ok(f"Sampling: Euler ODE (configurable steps)")
+
+# =========================================================================
+test("VRAM Estimation for Colab T4 (16GB)")
+# =========================================================================
+for name, factory, res, bs in [
+    ("tiny @256px bs=8", liquid_diffusion_tiny, 256, 8),
+    ("tiny @256px bs=4", liquid_diffusion_tiny, 256, 4),
+    ("small @256px bs=4", liquid_diffusion_small, 256, 4),
+    ("base @512px bs=2", liquid_diffusion_base, 512, 2),
+    ("tiny @512px bs=4", liquid_diffusion_tiny, 512, 4),
+]:
+    m = factory()
+    tp = sum(p.numel() for p in m.parameters())
+    # Conservative VRAM estimate:
+    # params (fp16) + gradients (fp32) + Adam states (2×fp32) + activations
+    param_gb = tp * 2 / 1e9   # fp16
+    grad_gb = tp * 4 / 1e9    # fp32
+    optim_gb = tp * 8 / 1e9   # Adam: 2× fp32  
+    # Activation estimate: ~4 bytes per element, scale with resolution and batch
+    act_gb = bs * res * res * max(m.channels) * 4 * len(m.channels) * 2 / 1e9  
+    total_gb = param_gb + grad_gb + optim_gb + act_gb
+    fits = "✓ fits T4" if total_gb < 15 else "✗ too large"
+    print(f"    {name:25s}: {tp/1e6:5.1f}M params, ~{total_gb:5.1f}GB  {fits}")
+    del m
+
+
+# =========================================================================
+# FINAL SUMMARY
+# =========================================================================
+print("\n" + "=" * 70)
+if all_passed:
+    print("✅ ALL TESTS PASSED")
+    print("\nReady for training! Open the Colab notebook:")
+    print("  LiquidDiffusion_Training.ipynb")
+else:
+    print("❌ SOME TESTS FAILED — check output above")
+print("=" * 70)