Add Colab/Kaggle training notebook with all dataset options

LiquidFlow_Training.ipynb

@@ -0,0 +1,747 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🌊 LiquidFlow — Liquid-SSM Flow Matching Image Generator\n",
+    "\n",
+    "A **novel architecture** combining:\n",
+    "- **Liquid Time-Constant Networks** (CfC closed-form) — adaptive ODE dynamics, bounded by construction\n",
+    "- **Selective State Space Models** (Mamba-style) — linear-time long-range context, parallelizable\n",
+    "- **Zigzag Scanning** — 2D spatial awareness for image patches\n",
+    "- **Physics-Informed Regularization** — smoothness + total variation constraints\n",
+    "- **Rectified Flow Matching** — ODE-based generation (no noise schedule tuning)\n",
+    "\n",
+    "### 📋 What this notebook does\n",
+    "1. **Install & clone** the LiquidFlow codebase\n",
+    "2. **Choose a dataset** (CIFAR-10, Flowers-102, CelebA, or custom folder)\n",
+    "3. **Choose a model size** (tiny ~6M, small ~14M, base ~38M)\n",
+    "4. **Train** with one click — all Colab/Kaggle optimized\n",
+    "5. **Generate images** and visualize progress\n",
+    "6. **Export** trained model for mobile deployment\n",
+    "\n",
+    "### 💻 Hardware Requirements\n",
+    "| Config | GPU VRAM | Best For |\n",
+    "|--------|----------|----------|\n",
+    "| tiny-128 (bs=32) | ~4 GB | Colab free T4, Kaggle |\n",
+    "| small-128 (bs=16) | ~8 GB | Colab free T4, Kaggle |\n",
+    "| base-256 (bs=8) | ~12 GB | Colab Pro, Kaggle |\n",
+    "| 512 (bs=4) | ~14 GB | Colab Pro, A100 |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 0. Setup & Install"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Check GPU\n",
+    "!nvidia-smi || echo 'No GPU — CPU training only (very slow)'\n",
+    "import torch\n",
+    "print(f'PyTorch {torch.__version__}, CUDA available: {torch.cuda.is_available()}')\n",
+    "if torch.cuda.is_available():\n",
+    "    print(f'GPU: {torch.cuda.get_device_name(0)}, VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install dependencies\n",
+    "!pip install -q torch torchvision einops pillow matplotlib tqdm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Clone the repo (or just copy the files if already have them)\n",
+    "import os\n",
+    "if not os.path.exists('liquidflow'):\n",
+    "    !git clone https://huggingface.co/krystv/LiquidFlow liquidflow_repo\n",
+    "    !cp -r liquidflow_repo/liquidflow .\n",
+    "else:\n",
+    "    print('liquidflow/ already exists')\n",
+    "\n",
+    "# Verify\n",
+    "from liquidflow.model import liquidflow_tiny, liquidflow_small, liquidflow_base, liquidflow_512\n",
+    "from liquidflow.losses import PhysicsInformedFlowLoss, EMAModel\n",
+    "from liquidflow.sampling import euler_sample, heun_sample, generate_grid, make_grid_image\n",
+    "print('✅ LiquidFlow imported successfully!')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 1. ⚙️ Configuration — EDIT THIS CELL\n",
+    "\n",
+    "Choose your dataset, model size, and training hyperparameters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title 🎛️ Training Configuration { display-mode: \"form\" }\n",
+    "\n",
+    "# ============== DATASET ==============\n",
+    "#@markdown ### Dataset\n",
+    "DATASET = 'cifar10'  #@param ['cifar10', 'flowers', 'celeba', 'folder', 'fashion_mnist', 'afhq', 'lsun_churches']\n",
+    "CUSTOM_DATA_DIR = '/content/my_images'  #@param {type:\"string\"}\n",
+    "#@markdown > For 'folder': put images in CUSTOM_DATA_DIR. Supports .png/.jpg/.webp\n",
+    "\n",
+    "# ============== MODEL ==============\n",
+    "#@markdown ### Model\n",
+    "MODEL_SIZE = 'tiny'  #@param ['tiny', 'small', 'base', '512']\n",
+    "IMG_SIZE = 128  #@param [32, 64, 128, 256, 512] {type:\"integer\"}\n",
+    "\n",
+    "# ============== TRAINING ==============\n",
+    "#@markdown ### Training\n",
+    "EPOCHS = 100  #@param {type:\"integer\"}\n",
+    "BATCH_SIZE = 32  #@param [4, 8, 16, 32, 64, 128] {type:\"integer\"}\n",
+    "LEARNING_RATE = 3e-4  #@param {type:\"number\"}\n",
+    "GRAD_ACCUM = 1  #@param [1, 2, 4, 8] {type:\"integer\"}\n",
+    "USE_AMP = True  #@param {type:\"boolean\"}\n",
+    "\n",
+    "# ============== PHYSICS LOSS ==============\n",
+    "#@markdown ### Physics-Informed Regularization\n",
+    "LAMBDA_SMOOTH = 0.01  #@param {type:\"number\"}\n",
+    "LAMBDA_TV = 0.001  #@param {type:\"number\"}\n",
+    "\n",
+    "# ============== SAMPLING ==============\n",
+    "#@markdown ### Sampling & Logging\n",
+    "SAMPLE_EVERY = 5  #@param {type:\"integer\"}\n",
+    "SAMPLE_STEPS = 50  #@param [10, 25, 50, 100] {type:\"integer\"}\n",
+    "LOG_EVERY = 50  #@param {type:\"integer\"}\n",
+    "SAVE_EVERY = 10  #@param {type:\"integer\"}\n",
+    "\n",
+    "# ============== PATHS ==============\n",
+    "OUTPUT_DIR = './outputs'\n",
+    "DATA_DIR = './data'\n",
+    "\n",
+    "# ============== AUTO-CONFIG ==============\n",
+    "# Smart batch size based on GPU memory\n",
+    "import torch\n",
+    "if torch.cuda.is_available():\n",
+    "    vram_gb = torch.cuda.get_device_properties(0).total_mem / 1e9\n",
+    "    print(f'GPU VRAM: {vram_gb:.1f} GB')\n",
+    "    \n",
+    "    # Auto-adjust batch size if needed\n",
+    "    recommended = {\n",
+    "        (32, 'tiny'): 128, (64, 'tiny'): 64, (128, 'tiny'): 32,\n",
+    "        (32, 'small'): 64, (64, 'small'): 32, (128, 'small'): 16,\n",
+    "        (256, 'base'): 8, (512, '512'): 4,\n",
+    "    }\n",
+    "    key = (IMG_SIZE, MODEL_SIZE)\n",
+    "    if key in recommended and vram_gb < 16:\n",
+    "        rec_bs = recommended[key]\n",
+    "        if BATCH_SIZE > rec_bs:\n",
+    "            print(f'⚠️ Reducing batch size {BATCH_SIZE} → {rec_bs} for {vram_gb:.0f}GB VRAM')\n",
+    "            BATCH_SIZE = rec_bs\n",
+    "else:\n",
+    "    print('⚠️ No GPU detected — training will be very slow!')\n",
+    "    USE_AMP = False\n",
+    "\n",
+    "print(f'\\n📋 Config: {MODEL_SIZE}-{IMG_SIZE}, {DATASET}, bs={BATCH_SIZE}, lr={LEARNING_RATE}, epochs={EPOCHS}')\n",
+    "print(f'   Physics: λ_smooth={LAMBDA_SMOOTH}, λ_tv={LAMBDA_TV}')\n",
+    "print(f'   AMP: {USE_AMP}, GradAccum: {GRAD_ACCUM}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 2. 📦 Load Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torchvision\n",
+    "import torchvision.transforms as transforms\n",
+    "from torch.utils.data import DataLoader, Dataset, ConcatDataset\n",
+    "from pathlib import Path\n",
+    "from PIL import Image\n",
+    "import os\n",
+    "\n",
+    "# Standard transform\n",
+    "def get_transform(img_size):\n",
+    "    return transforms.Compose([\n",
+    "        transforms.Resize(img_size + img_size // 8),\n",
+    "        transforms.CenterCrop(img_size),\n",
+    "        transforms.RandomHorizontalFlip(),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize([0.5]*3, [0.5]*3),\n",
+    "    ])\n",
+    "\n",
+    "class ImageFolderFlat(Dataset):\n",
+    "    \"\"\"Load all images from a folder (recursively).\"\"\"\n",
+    "    def __init__(self, root, transform):\n",
+    "        self.transform = transform\n",
+    "        self.files = []\n",
+    "        for ext in ['*.png', '*.jpg', '*.jpeg', '*.webp', '*.bmp']:\n",
+    "            self.files.extend(Path(root).rglob(ext))\n",
+    "        self.files = sorted(self.files)\n",
+    "        print(f'Found {len(self.files)} images in {root}')\n",
+    "    def __len__(self): return len(self.files)\n",
+    "    def __getitem__(self, idx):\n",
+    "        return self.transform(Image.open(self.files[idx]).convert('RGB'))\n",
+    "\n",
+    "class GrayscaleToRGB:\n",
+    "    \"\"\"Convert 1-channel grayscale to 3-channel RGB.\"\"\"\n",
+    "    def __call__(self, x):\n",
+    "        if x.shape[0] == 1:\n",
+    "            x = x.repeat(3, 1, 1)\n",
+    "        return x\n",
+    "\n",
+    "tfm = get_transform(IMG_SIZE)\n",
+    "\n",
+    "if DATASET == 'cifar10':\n",
+    "    dataset = torchvision.datasets.CIFAR10(root=DATA_DIR, train=True, download=True, transform=tfm)\n",
+    "    print(f'✅ CIFAR-10: {len(dataset)} images')\n",
+    "\n",
+    "elif DATASET == 'flowers':\n",
+    "    ds_train = torchvision.datasets.Flowers102(root=DATA_DIR, split='train', download=True, transform=tfm)\n",
+    "    ds_val = torchvision.datasets.Flowers102(root=DATA_DIR, split='val', download=True, transform=tfm)\n",
+    "    ds_test = torchvision.datasets.Flowers102(root=DATA_DIR, split='test', download=True, transform=tfm)\n",
+    "    dataset = ConcatDataset([ds_train, ds_val, ds_test])  # Use all splits for generation\n",
+    "    print(f'✅ Flowers-102: {len(dataset)} images (all splits)')\n",
+    "\n",
+    "elif DATASET == 'celeba':\n",
+    "    dataset = torchvision.datasets.CelebA(root=DATA_DIR, split='train', download=True, transform=tfm)\n",
+    "    print(f'✅ CelebA: {len(dataset)} images')\n",
+    "\n",
+    "elif DATASET == 'fashion_mnist':\n",
+    "    fm_tfm = transforms.Compose([\n",
+    "        transforms.Resize(IMG_SIZE),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize([0.5], [0.5]),\n",
+    "        GrayscaleToRGB(),\n",
+    "    ])\n",
+    "    dataset = torchvision.datasets.FashionMNIST(root=DATA_DIR, train=True, download=True, transform=fm_tfm)\n",
+    "    print(f'✅ Fashion-MNIST: {len(dataset)} images (converted to RGB)')\n",
+    "\n",
+    "elif DATASET == 'afhq':\n",
+    "    # Download AFHQ from Kaggle or manual\n",
+    "    afhq_dir = os.path.join(DATA_DIR, 'afhq', 'train')\n",
+    "    if not os.path.exists(afhq_dir):\n",
+    "        print('⬇️ Downloading AFHQ...')\n",
+    "        !pip install -q gdown\n",
+    "        !gdown 1Gof5BaELXlmSJIlvKMYCe9ONYPebkNsf -O {DATA_DIR}/afhq.zip\n",
+    "        !unzip -q {DATA_DIR}/afhq.zip -d {DATA_DIR}/afhq\n",
+    "    dataset = ImageFolderFlat(afhq_dir, tfm)\n",
+    "    print(f'✅ AFHQ: {len(dataset)} images')\n",
+    "\n",
+    "elif DATASET == 'lsun_churches':\n",
+    "    # LSUN requires manual download — point to extracted folder\n",
+    "    lsun_dir = os.path.join(DATA_DIR, 'lsun_churches')\n",
+    "    if not os.path.exists(lsun_dir):\n",
+    "        print('❌ LSUN churches not found. Please download and extract to', lsun_dir)\n",
+    "        print('   See: https://github.com/fyu/lsun')\n",
+    "        raise FileNotFoundError(lsun_dir)\n",
+    "    dataset = ImageFolderFlat(lsun_dir, tfm)\n",
+    "    print(f'✅ LSUN Churches: {len(dataset)} images')\n",
+    "\n",
+    "elif DATASET == 'folder':\n",
+    "    dataset = ImageFolderFlat(CUSTOM_DATA_DIR, tfm)\n",
+    "    print(f'✅ Custom folder: {len(dataset)} images from {CUSTOM_DATA_DIR}')\n",
+    "\n",
+    "else:\n",
+    "    raise ValueError(f'Unknown dataset: {DATASET}')\n",
+    "\n",
+    "# Show a few samples\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 8, figsize=(16, 2))\n",
+    "for i, ax in enumerate(axes):\n",
+    "    sample = dataset[i]\n",
+    "    if isinstance(sample, (list, tuple)):\n",
+    "        sample = sample[0]\n",
+    "    img = sample * 0.5 + 0.5  # denormalize\n",
+    "    ax.imshow(img.permute(1, 2, 0).clamp(0, 1).numpy())\n",
+    "    ax.axis('off')\n",
+    "plt.suptitle(f'{DATASET} samples ({IMG_SIZE}×{IMG_SIZE})', fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 3. 🏗️ Build Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from liquidflow.model import liquidflow_tiny, liquidflow_small, liquidflow_base, liquidflow_512\n",
+    "\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "\n",
+    "model_factories = {\n",
+    "    'tiny':  liquidflow_tiny,\n",
+    "    'small': liquidflow_small,\n",
+    "    'base':  liquidflow_base,\n",
+    "    '512':   liquidflow_512,\n",
+    "}\n",
+    "\n",
+    "model = model_factories[MODEL_SIZE](img_size=IMG_SIZE).to(device)\n",
+    "\n",
+    "num_params = model.count_params()\n",
+    "print(f'🏗️ LiquidFlow-{MODEL_SIZE}')\n",
+    "print(f'   Parameters: {num_params:,} ({num_params/1e6:.1f}M)')\n",
+    "print(f'   Image size: {IMG_SIZE}×{IMG_SIZE}')\n",
+    "print(f'   Patch size: {model.patch_size}')\n",
+    "print(f'   Num patches: {model.num_patches}')\n",
+    "print(f'   Model dim: {model.d_model}')\n",
+    "print(f'   Depth: {model.depth}')\n",
+    "print(f'   Device: {device}')\n",
+    "\n",
+    "# Quick forward pass test\n",
+    "with torch.no_grad():\n",
+    "    test_x = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device=device)\n",
+    "    test_t = torch.tensor([0.5], device=device)\n",
+    "    test_v = model(test_x, test_t)\n",
+    "    assert test_v.shape == test_x.shape\n",
+    "    print(f'   ✅ Forward pass OK: {test_x.shape} → {test_v.shape}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 4. 🚀 Train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "import time\n",
+    "import json\n",
+    "import torch.nn as nn\n",
+    "from torch.cuda.amp import autocast, GradScaler\n",
+    "from liquidflow.losses import PhysicsInformedFlowLoss, EMAModel\n",
+    "from liquidflow.sampling import euler_sample, make_grid_image\n",
+    "from IPython.display import display, clear_output\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Prepare\n",
+    "os.makedirs(f'{OUTPUT_DIR}/samples', exist_ok=True)\n",
+    "os.makedirs(f'{OUTPUT_DIR}/checkpoints', exist_ok=True)\n",
+    "\n",
+    "dataloader = DataLoader(\n",
+    "    dataset, batch_size=BATCH_SIZE, shuffle=True,\n",
+    "    num_workers=2, pin_memory=True, drop_last=True\n",
+    ")\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE,\n",
+    "                               betas=(0.9, 0.999), weight_decay=0.01)\n",
+    "\n",
+    "total_steps = EPOCHS * len(dataloader) // GRAD_ACCUM\n",
+    "warmup_steps = min(500, total_steps // 10)\n",
+    "\n",
+    "def cosine_lr(step):\n",
+    "    if step < warmup_steps:\n",
+    "        return step / max(1, warmup_steps)\n",
+    "    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n",
+    "    return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))\n",
+    "\n",
+    "scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, cosine_lr)\n",
+    "criterion = PhysicsInformedFlowLoss(\n",
+    "    lambda_smooth=LAMBDA_SMOOTH, lambda_tv=LAMBDA_TV\n",
+    ").to(device)\n",
+    "ema = EMAModel(model, decay=0.9999)\n",
+    "scaler = GradScaler(enabled=USE_AMP)\n",
+    "\n",
+    "# Training log\n",
+    "all_losses = []\n",
+    "global_step = 0\n",
+    "\n",
+    "print(f'🚀 Training {EPOCHS} epochs, {total_steps} steps')\n",
+    "print(f'   Effective batch: {BATCH_SIZE} × {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}')\n",
+    "print(f'   LR: {LEARNING_RATE} → warmup {warmup_steps} steps → cosine decay')\n",
+    "print()\n",
+    "\n",
+    "t_start = time.time()\n",
+    "\n",
+    "for epoch in range(EPOCHS):\n",
+    "    model.train()\n",
+    "    epoch_loss = 0.0\n",
+    "    epoch_flow = 0.0\n",
+    "    n_batches = 0\n",
+    "\n",
+    "    for batch_idx, batch_data in enumerate(dataloader):\n",
+    "        if isinstance(batch_data, (list, tuple)):\n",
+    "            x1 = batch_data[0].to(device)\n",
+    "        else:\n",
+    "            x1 = batch_data.to(device)\n",
+    "\n",
+    "        B = x1.shape[0]\n",
+    "        x0 = torch.randn_like(x1)\n",
+    "        t = torch.rand(B, device=device)\n",
+    "        t_e = t.view(B, 1, 1, 1)\n",
+    "        x_t = t_e * x1 + (1 - t_e) * x0\n",
+    "\n",
+    "        with autocast(enabled=USE_AMP):\n",
+    "            v_pred = model(x_t, t)\n",
+    "            loss, ld = criterion(v_pred, x0, x1, t, step=global_step)\n",
+    "            loss = loss / GRAD_ACCUM\n",
+    "\n",
+    "        scaler.scale(loss).backward()\n",
+    "\n",
+    "        if (batch_idx + 1) % GRAD_ACCUM == 0:\n",
+    "            scaler.unscale_(optimizer)\n",
+    "            gn = nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "            scaler.step(optimizer)\n",
+    "            scaler.update()\n",
+    "            optimizer.zero_grad()\n",
+    "            scheduler.step()\n",
+    "            ema.update(model)\n",
+    "            global_step += 1\n",
+    "\n",
+    "            epoch_loss += ld['total'].item()\n",
+    "            epoch_flow += ld['flow'].item()\n",
+    "            n_batches += 1\n",
+    "\n",
+    "            if global_step % LOG_EVERY == 0:\n",
+    "                avg = epoch_loss / n_batches\n",
+    "                avg_f = epoch_flow / n_batches\n",
+    "                lr_now = scheduler.get_last_lr()[0]\n",
+    "                elapsed = time.time() - t_start\n",
+    "                it_s = global_step / elapsed\n",
+    "                all_losses.append({'step': global_step, 'loss': avg, 'flow': avg_f,\n",
+    "                                   'lr': lr_now, 'epoch': epoch})\n",
+    "                print(f'  E{epoch+1} step {global_step}/{total_steps} | '\n",
+    "                      f'loss={avg:.4f} flow={avg_f:.4f} lr={lr_now:.2e} '\n",
+    "                      f'gn={gn:.2f} [{it_s:.1f} it/s]')\n",
+    "\n",
+    "    # End of epoch\n",
+    "    avg_epoch = epoch_loss / max(1, n_batches)\n",
+    "    print(f'\\n📊 Epoch {epoch+1}/{EPOCHS} — avg loss: {avg_epoch:.4f}\\n')\n",
+    "\n",
+    "    # Sample\n",
+    "    if (epoch + 1) % SAMPLE_EVERY == 0 or epoch == 0:\n",
+    "        model.eval()\n",
+    "        ema.apply_shadow(model)\n",
+    "        with torch.no_grad():\n",
+    "            n_samples = min(16, BATCH_SIZE)\n",
+    "            imgs = euler_sample(model, (n_samples, 3, IMG_SIZE, IMG_SIZE),\n",
+    "                                num_steps=SAMPLE_STEPS, device=device)\n",
+    "            imgs = imgs.clamp(-1, 1) * 0.5 + 0.5\n",
+    "            grid = make_grid_image(imgs, nrow=4)\n",
+    "            grid.save(f'{OUTPUT_DIR}/samples/epoch_{epoch+1:04d}.png')\n",
+    "\n",
+    "            # Display inline\n",
+    "            fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "            ax.imshow(grid)\n",
+    "            ax.set_title(f'Epoch {epoch+1} — {MODEL_SIZE}-{IMG_SIZE} on {DATASET}')\n",
+    "            ax.axis('off')\n",
+    "            plt.tight_layout()\n",
+    "            plt.show()\n",
+    "\n",
+    "        ema.restore(model)\n",
+    "        model.train()\n",
+    "\n",
+    "    # Checkpoint\n",
+    "    if (epoch + 1) % SAVE_EVERY == 0:\n",
+    "        ckpt = {\n",
+    "            'model': model.state_dict(),\n",
+    "            'optimizer': optimizer.state_dict(),\n",
+    "            'scheduler': scheduler.state_dict(),\n",
+    "            'ema': ema.state_dict(),\n",
+    "            'epoch': epoch,\n",
+    "            'global_step': global_step,\n",
+    "        }\n",
+    "        torch.save(ckpt, f'{OUTPUT_DIR}/checkpoints/epoch_{epoch+1:04d}.pt')\n",
+    "        torch.save(ckpt, f'{OUTPUT_DIR}/checkpoints/latest.pt')\n",
+    "        print(f'💾 Checkpoint saved: epoch {epoch+1}')\n",
+    "\n",
+    "# Save final\n",
+    "ema.apply_shadow(model)\n",
+    "torch.save({'model': model.state_dict(), 'config': {\n",
+    "    'model_size': MODEL_SIZE, 'img_size': IMG_SIZE, 'dataset': DATASET,\n",
+    "    'num_params': num_params, 'epochs': EPOCHS,\n",
+    "}}, f'{OUTPUT_DIR}/liquidflow_final.pt')\n",
+    "ema.restore(model)\n",
+    "\n",
+    "elapsed = time.time() - t_start\n",
+    "print(f'\\n✅ Training complete! {elapsed/60:.1f} min total')\n",
+    "print(f'   Final model: {OUTPUT_DIR}/liquidflow_final.pt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 5. 📈 Training Curves"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "if all_losses:\n",
+    "    steps = [d['step'] for d in all_losses]\n",
+    "    losses = [d['loss'] for d in all_losses]\n",
+    "    flows = [d['flow'] for d in all_losses]\n",
+    "    lrs = [d['lr'] for d in all_losses]\n",
+    "\n",
+    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "    ax1.plot(steps, losses, label='Total Loss', alpha=0.8)\n",
+    "    ax1.plot(steps, flows, label='Flow Loss', alpha=0.8)\n",
+    "    ax1.set_xlabel('Step'); ax1.set_ylabel('Loss')\n",
+    "    ax1.set_title('Training Loss'); ax1.legend(); ax1.grid(True, alpha=0.3)\n",
+    "\n",
+    "    ax2.plot(steps, lrs, color='orange')\n",
+    "    ax2.set_xlabel('Step'); ax2.set_ylabel('LR')\n",
+    "    ax2.set_title('Learning Rate Schedule'); ax2.grid(True, alpha=0.3)\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(f'{OUTPUT_DIR}/training_curves.png', dpi=150)\n",
+    "    plt.show()\n",
+    "else:\n",
+    "    print('No training logs yet.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 6. 🎨 Generate Images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title 🎨 Generation Settings { display-mode: \"form\" }\n",
+    "NUM_IMAGES = 16  #@param {type:\"integer\"}\n",
+    "GEN_STEPS = 50  #@param [10, 25, 50, 100, 200] {type:\"integer\"}\n",
+    "SAMPLER = 'euler'  #@param ['euler', 'heun']\n",
+    "SEED = 42  #@param {type:\"integer\"}\n",
+    "\n",
+    "import torch\n",
+    "from liquidflow.sampling import euler_sample, heun_sample, make_grid_image\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Load best model\n",
+    "ckpt_path = f'{OUTPUT_DIR}/liquidflow_final.pt'\n",
+    "if os.path.exists(ckpt_path):\n",
+    "    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)\n",
+    "    model.load_state_dict(ckpt['model'])\n",
+    "    print(f'Loaded: {ckpt_path}')\n",
+    "else:\n",
+    "    print(f'No checkpoint found, using current model weights')\n",
+    "\n",
+    "model.eval()\n",
+    "torch.manual_seed(SEED)\n",
+    "\n",
+    "shape = (NUM_IMAGES, 3, IMG_SIZE, IMG_SIZE)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    if SAMPLER == 'euler':\n",
+    "        images = euler_sample(model, shape, num_steps=GEN_STEPS, device=device)\n",
+    "    else:\n",
+    "        images = heun_sample(model, shape, num_steps=GEN_STEPS, device=device)\n",
+    "\n",
+    "images = images.clamp(-1, 1) * 0.5 + 0.5\n",
+    "grid = make_grid_image(images, nrow=int(NUM_IMAGES**0.5))\n",
+    "grid.save(f'{OUTPUT_DIR}/generated_final.png')\n",
+    "\n",
+    "plt.figure(figsize=(10, 10))\n",
+    "plt.imshow(grid)\n",
+    "plt.title(f'LiquidFlow-{MODEL_SIZE} | {DATASET} {IMG_SIZE}×{IMG_SIZE} | {GEN_STEPS} steps ({SAMPLER})')\n",
+    "plt.axis('off')\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 7. 📱 Export for Mobile (ONNX + TorchScript)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Export to TorchScript for mobile deployment\n",
+    "model.eval()\n",
+    "\n",
+    "# TorchScript (for PyTorch Mobile / ExecuTorch)\n",
+    "example_x = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device=device)\n",
+    "example_t = torch.tensor([0.5], device=device)\n",
+    "\n",
+    "try:\n",
+    "    traced = torch.jit.trace(model, (example_x, example_t))\n",
+    "    ts_path = f'{OUTPUT_DIR}/liquidflow_mobile.pt'\n",
+    "    traced.save(ts_path)\n",
+    "    ts_size_mb = os.path.getsize(ts_path) / 1e6\n",
+    "    print(f'✅ TorchScript saved: {ts_path} ({ts_size_mb:.1f} MB)')\n",
+    "except Exception as e:\n",
+    "    print(f'⚠️ TorchScript export failed: {e}')\n",
+    "\n",
+    "# ONNX\n",
+    "try:\n",
+    "    onnx_path = f'{OUTPUT_DIR}/liquidflow.onnx'\n",
+    "    torch.onnx.export(\n",
+    "        model.cpu(), (example_x.cpu(), example_t.cpu()),\n",
+    "        onnx_path, opset_version=14,\n",
+    "        input_names=['image', 'timestep'],\n",
+    "        output_names=['velocity'],\n",
+    "        dynamic_axes={'image': {0: 'batch'}, 'timestep': {0: 'batch'}, 'velocity': {0: 'batch'}}\n",
+    "    )\n",
+    "    onnx_size_mb = os.path.getsize(onnx_path) / 1e6\n",
+    "    print(f'✅ ONNX saved: {onnx_path} ({onnx_size_mb:.1f} MB)')\n",
+    "    model.to(device)\n",
+    "except Exception as e:\n",
+    "    print(f'⚠️ ONNX export failed: {e}')\n",
+    "    model.to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 8. 🔬 Architecture Deep Dive\n",
+    "\n",
+    "### How LiquidFlow works\n",
+    "\n",
+    "```\n",
+    "Noise x₀ ~ N(0,I)  ──→  LiquidFlow v_θ(xₜ, t)  ──→  Image x₁\n",
+    "                           │\n",
+    "                    ┌──────┴──────┐\n",
+    "                    │  Patchify   │  (img → non-overlapping patches)\n",
+    "                    │  + PosEmb   │  (2D learnable positions)\n",
+    "                    │  + DepthConv│  (local structure)\n",
+    "                    └──────┬──────┘\n",
+    "                           │\n",
+    "              ┌────────────┼────────────┐\n",
+    "              │    L × LiquidSSM Block   │\n",
+    "              │  ┌──────────────────┐    │\n",
+    "              │  │ AdaLN (t-cond)   │    │\n",
+    "              │  │ Zigzag Scan      │    │  ← rotates scan pattern per layer\n",
+    "              │  │ SelectiveSSM     │    │  ← Mamba-style, input-dependent\n",
+    "              │  │ + LiquidCfC      │    │  ← CfC gating, bounded dynamics\n",
+    "              │  │ + FFN            │    │\n",
+    "              │  │ + Skip Connect   │    │  ← U-Net style long skips\n",
+    "              │  └──────────────────┘    │\n",
+    "              └────────────┼────────────┘\n",
+    "                           │\n",
+    "                    ┌──────┴──────┐\n",
+    "                    │  DepthConv  │\n",
+    "                    │  Unpatchify │  (patches → img)\n",
+    "                    └──────┬──────┘\n",
+    "                           │\n",
+    "                     velocity v_θ\n",
+    "```\n",
+    "\n",
+    "### Key Innovations\n",
+    "\n",
+    "1. **Liquid CfC Cell**: Instead of solving the ODE `dx/dt = f(x,t)` numerically, we use the\n",
+    "   closed-form solution `x(t+Δt) = σ(-f_τ) ⊙ x(t) + (1 - σ(-f_τ)) ⊙ f_x`.\n",
+    "   The sigmoid gating **guarantees bounded dynamics** — no training explosion possible.\n",
+    "\n",
+    "2. **SSM + Liquid dual path**: The SSM branch captures long-range spatial dependencies\n",
+    "   via selective scanning; the Liquid branch adds continuous-time adaptive dynamics.\n",
+    "   A learnable mixing coefficient balances them.\n",
+    "\n",
+    "3. **Physics-informed loss**: Smoothness (Laplacian) and Total Variation regularizers\n",
+    "   act as soft PDE constraints on generated images, improving training stability\n",
+    "   and reducing artifacts without domain-specific physics knowledge.\n",
+    "\n",
+    "4. **Flow Matching = Liquid ODE**: Rectified flow trains `v_θ` to follow straight paths\n",
+    "   from noise to data. This is structurally identical to the LTC ODE, making Liquid\n",
+    "   networks a natural fit as the velocity field parameterization."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 9. 🧪 Recommended Experiments\n",
+    "\n",
+    "| Experiment | Dataset | Model | IMG_SIZE | Epochs | Notes |\n",
+    "|------------|---------|-------|----------|--------|-------|\n",
+    "| Quick sanity check | CIFAR-10 | tiny | 32 | 20 | ~5 min on T4 |\n",
+    "| Baseline 128×128 | CIFAR-10 | tiny | 128 | 100 | ~2 hrs on T4 |\n",
+    "| Quality 128×128 | Flowers-102 | small | 128 | 200 | ~4 hrs on T4 |\n",
+    "| Faces 128×128 | CelebA | small | 128 | 50 | ~6 hrs on T4 |\n",
+    "| High-res 512×512 | CelebA | 512 | 512 | 100 | needs ≥16GB |\n",
+    "| Production | Your data | small | 128 | 300+ | best quality |\n",
+    "\n",
+    "### Tips for best results:\n",
+    "- Start with `tiny` + low epochs to verify everything works\n",
+    "- Use `small` for 128×128 production quality\n",
+    "- Increase `SAMPLE_STEPS` to 100+ for final generation\n",
+    "- `heun` sampler gives better quality at half the steps vs `euler`\n",
+    "- Physics loss warmup is automatic — don't increase λ too much"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}