Fix deprecated torch.cuda.amp API in training cell — use torch.amp instead

LiquidFlow_Training.ipynb

@@ -1,747 +1 @@
-{
- "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
-}
+{"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","### 💻 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":["!nvidia-smi || echo 'No GPU — CPU only'\nimport torch\nprint(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')\nif 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":["!pip install -q torch torchvision einops pillow matplotlib tqdm"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import os\nif not os.path.exists('liquidflow'):\n    !git clone https://huggingface.co/krystv/LiquidFlow liquidflow_repo\n    !cp -r liquidflow_repo/liquidflow .\nelse:\n    print('liquidflow/ already exists — updating...')\n    !cd liquidflow_repo && git pull && cp -r liquidflow/* ../liquidflow/\n\nfrom liquidflow.model import liquidflow_tiny, liquidflow_small, liquidflow_base, liquidflow_512\nfrom liquidflow.losses import PhysicsInformedFlowLoss, EMAModel\nfrom liquidflow.sampling import euler_sample, heun_sample, generate_grid, make_grid_image\nprint('✅ LiquidFlow imported successfully!')"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 1. ⚙️ Configuration"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["#@title 🎛️ Training Configuration { display-mode: \"form\" }\n\nDATASET = 'cifar10'  #@param ['cifar10', 'flowers', 'celeba', 'folder', 'fashion_mnist', 'afhq', 'lsun_churches']\nCUSTOM_DATA_DIR = '/content/my_images'  #@param {type:\"string\"}\nMODEL_SIZE = 'tiny'  #@param ['tiny', 'small', 'base', '512']\nIMG_SIZE = 128  #@param [32, 64, 128, 256, 512] {type:\"integer\"}\nEPOCHS = 100  #@param {type:\"integer\"}\nBATCH_SIZE = 32  #@param [4, 8, 16, 32, 64, 128] {type:\"integer\"}\nLEARNING_RATE = 3e-4  #@param {type:\"number\"}\nGRAD_ACCUM = 1  #@param [1, 2, 4, 8] {type:\"integer\"}\nUSE_AMP = True  #@param {type:\"boolean\"}\nLAMBDA_SMOOTH = 0.01  #@param {type:\"number\"}\nLAMBDA_TV = 0.001  #@param {type:\"number\"}\nSAMPLE_EVERY = 5  #@param {type:\"integer\"}\nSAMPLE_STEPS = 50  #@param [10, 25, 50, 100] {type:\"integer\"}\nLOG_EVERY = 50  #@param {type:\"integer\"}\nSAVE_EVERY = 10  #@param {type:\"integer\"}\nOUTPUT_DIR = './outputs'\nDATA_DIR = './data'\n\nimport torch\nif 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    recommended = {(32,'tiny'):128,(64,'tiny'):64,(128,'tiny'):32,(32,'small'):64,(64,'small'):32,(128,'small'):16,(256,'base'):8,(512,'512'):4}\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\nelse:\n    print('⚠️ No GPU detected'); USE_AMP = False\n\nprint(f'\\n📋 Config: {MODEL_SIZE}-{IMG_SIZE}, {DATASET}, bs={BATCH_SIZE}, lr={LEARNING_RATE}, epochs={EPOCHS}')"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 2. 📦 Load Dataset"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import torchvision\nimport torchvision.transforms as transforms\nfrom torch.utils.data import DataLoader, Dataset, ConcatDataset\nfrom pathlib import Path\nfrom PIL import Image\nimport os, matplotlib.pyplot as plt, numpy as np\n\ndef get_transform(img_size):\n    return transforms.Compose([transforms.Resize(img_size+img_size//8), transforms.CenterCrop(img_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5]*3,[0.5]*3)])\n\nclass ImageFolderFlat(Dataset):\n    def __init__(self, root, transform):\n        self.transform = transform\n        self.files = []\n        for ext in ['*.png','*.jpg','*.jpeg','*.webp','*.bmp']: self.files.extend(Path(root).rglob(ext))\n        self.files = sorted(self.files); print(f'Found {len(self.files)} images in {root}')\n    def __len__(self): return len(self.files)\n    def __getitem__(self, idx): return self.transform(Image.open(self.files[idx]).convert('RGB'))\n\nclass GrayscaleToRGB:\n    def __call__(self, x): return x.repeat(3,1,1) if x.shape[0]==1 else x\n\ntfm = get_transform(IMG_SIZE)\n\nif DATASET == 'cifar10':\n    dataset = torchvision.datasets.CIFAR10(root=DATA_DIR, train=True, download=True, transform=tfm)\nelif DATASET == 'flowers':\n    dataset = ConcatDataset([torchvision.datasets.Flowers102(root=DATA_DIR, split=s, download=True, transform=tfm) for s in ['train','val','test']])\nelif DATASET == 'celeba':\n    dataset = torchvision.datasets.CelebA(root=DATA_DIR, split='train', download=True, transform=tfm)\nelif DATASET == 'fashion_mnist':\n    fm_tfm = transforms.Compose([transforms.Resize(IMG_SIZE), transforms.ToTensor(), transforms.Normalize([0.5],[0.5]), GrayscaleToRGB()])\n    dataset = torchvision.datasets.FashionMNIST(root=DATA_DIR, train=True, download=True, transform=fm_tfm)\nelif DATASET == 'folder':\n    dataset = ImageFolderFlat(CUSTOM_DATA_DIR, tfm)\nelse:\n    raise ValueError(f'Unknown dataset: {DATASET}')\n\nprint(f'✅ {DATASET}: {len(dataset)} images')\n\nfig, axes = plt.subplots(1, 8, figsize=(16, 2))\nfor i, ax in enumerate(axes):\n    sample = dataset[i]; sample = sample[0] if isinstance(sample,(list,tuple)) else sample\n    ax.imshow((sample*0.5+0.5).permute(1,2,0).clamp(0,1).numpy()); ax.axis('off')\nplt.suptitle(f'{DATASET} samples ({IMG_SIZE}×{IMG_SIZE})'); plt.tight_layout(); plt.show()"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 3. 🏗️ Build Model"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import torch\nfrom liquidflow.model import liquidflow_tiny, liquidflow_small, liquidflow_base, liquidflow_512\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel = {'tiny':liquidflow_tiny,'small':liquidflow_small,'base':liquidflow_base,'512':liquidflow_512}[MODEL_SIZE](img_size=IMG_SIZE).to(device)\n\nnum_params = model.count_params()\nprint(f'🏗️ LiquidFlow-{MODEL_SIZE}: {num_params:,} ({num_params/1e6:.1f}M) params')\nprint(f'   {IMG_SIZE}×{IMG_SIZE}, patch={model.patch_size}, patches={model.num_patches}, dim={model.d_model}, depth={model.depth}')\n\nwith torch.no_grad():\n    v = model(torch.randn(1,3,IMG_SIZE,IMG_SIZE,device=device), torch.tensor([0.5],device=device))\n    print(f'   ✅ Forward pass OK')"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 4. 🚀 Train"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import math, time, json\nimport torch.nn as nn\nfrom liquidflow.losses import PhysicsInformedFlowLoss, EMAModel\nfrom liquidflow.sampling import euler_sample, make_grid_image\nimport matplotlib.pyplot as plt\n\nos.makedirs(f'{OUTPUT_DIR}/samples', exist_ok=True)\nos.makedirs(f'{OUTPUT_DIR}/checkpoints', exist_ok=True)\n\ndataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True, drop_last=True)\noptimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas=(0.9,0.999), weight_decay=0.01)\n\ntotal_steps = EPOCHS * len(dataloader) // GRAD_ACCUM\nwarmup_steps = min(500, total_steps // 10)\ndef cosine_lr(step):\n    if step < warmup_steps: return step / max(1, warmup_steps)\n    p = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * p))\n\nscheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, cosine_lr)\ncriterion = PhysicsInformedFlowLoss(lambda_smooth=LAMBDA_SMOOTH, lambda_tv=LAMBDA_TV).to(device)\nema = EMAModel(model, decay=0.9999)\n\n# Use modern AMP API (no deprecation warnings)\namp_device = 'cuda' if device.type == 'cuda' else 'cpu'\nscaler = torch.amp.GradScaler(amp_device, enabled=USE_AMP)\n\nall_losses = []\nglobal_step = 0\nprint(f'🚀 Training {EPOCHS} epochs, {total_steps} steps')\nprint(f'   Effective batch: {BATCH_SIZE} × {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}')\nprint(f'   LR: {LEARNING_RATE} → warmup {warmup_steps} steps → cosine decay\\n')\n\nt_start = time.time()\nfor epoch in range(EPOCHS):\n    model.train()\n    epoch_loss = epoch_flow = 0.0; n_batches = 0\n\n    for batch_idx, batch_data in enumerate(dataloader):\n        x1 = (batch_data[0] if isinstance(batch_data,(list,tuple)) else batch_data).to(device)\n        B = x1.shape[0]\n        x0 = torch.randn_like(x1)\n        t = torch.rand(B, device=device)\n        x_t = t.view(B,1,1,1) * x1 + (1 - t.view(B,1,1,1)) * x0\n\n        with torch.amp.autocast(amp_device, 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); scaler.update(); optimizer.zero_grad()\n            scheduler.step(); ema.update(model); global_step += 1\n\n            epoch_loss += ld['total'].item(); epoch_flow += ld['flow'].item(); n_batches += 1\n\n            if global_step % LOG_EVERY == 0:\n                avg = epoch_loss/n_batches; avg_f = epoch_flow/n_batches\n                lr_now = scheduler.get_last_lr()[0]; it_s = global_step/(time.time()-t_start)\n                all_losses.append({'step':global_step,'loss':avg,'flow':avg_f,'lr':lr_now,'epoch':epoch})\n                print(f'  E{epoch+1} step {global_step}/{total_steps} | loss={avg:.4f} flow={avg_f:.4f} lr={lr_now:.2e} gn={gn:.2f} [{it_s:.1f} it/s]')\n\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    if (epoch+1) % SAMPLE_EVERY == 0 or epoch == 0:\n        model.eval(); ema.apply_shadow(model)\n        with torch.no_grad():\n            imgs = euler_sample(model, (min(16,BATCH_SIZE),3,IMG_SIZE,IMG_SIZE), 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            fig, ax = plt.subplots(1,1,figsize=(8,8)); ax.imshow(grid)\n            ax.set_title(f'Epoch {epoch+1} — {MODEL_SIZE}-{IMG_SIZE} on {DATASET}'); ax.axis('off'); plt.tight_layout(); plt.show()\n        ema.restore(model); model.train()\n\n    if (epoch+1) % SAVE_EVERY == 0:\n        ckpt = {'model':model.state_dict(),'optimizer':optimizer.state_dict(),'scheduler':scheduler.state_dict(),'ema':ema.state_dict(),'epoch':epoch,'global_step':global_step}\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\nema.apply_shadow(model)\ntorch.save({'model':model.state_dict(),'config':{'model_size':MODEL_SIZE,'img_size':IMG_SIZE,'dataset':DATASET,'num_params':num_params,'epochs':EPOCHS}}, f'{OUTPUT_DIR}/liquidflow_final.pt')\nema.restore(model)\nprint(f'\\n✅ Training complete! {(time.time()-t_start)/60:.1f} min total')"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 5. 📈 Training Curves"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import matplotlib.pyplot as plt\nif all_losses:\n    steps=[d['step'] for d in all_losses]; losses=[d['loss'] for d in all_losses]; flows=[d['flow'] for d in all_losses]; lrs=[d['lr'] for d in all_losses]\n    fig,(ax1,ax2) = plt.subplots(1,2,figsize=(14,5))\n    ax1.plot(steps,losses,label='Total',alpha=0.8); ax1.plot(steps,flows,label='Flow',alpha=0.8)\n    ax1.set_xlabel('Step'); ax1.set_ylabel('Loss'); ax1.set_title('Training Loss'); ax1.legend(); ax1.grid(True,alpha=0.3)\n    ax2.plot(steps,lrs,color='orange'); ax2.set_xlabel('Step'); ax2.set_ylabel('LR'); ax2.set_title('LR Schedule'); ax2.grid(True,alpha=0.3)\n    plt.tight_layout(); plt.savefig(f'{OUTPUT_DIR}/training_curves.png',dpi=150); plt.show()\nelse: 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\" }\nNUM_IMAGES = 16  #@param {type:\"integer\"}\nGEN_STEPS = 50  #@param [10, 25, 50, 100, 200] {type:\"integer\"}\nSAMPLER = 'euler'  #@param ['euler', 'heun']\nSEED = 42  #@param {type:\"integer\"}\n\nimport torch\nfrom liquidflow.sampling import euler_sample, heun_sample, make_grid_image\nimport matplotlib.pyplot as plt\n\nckpt_path = f'{OUTPUT_DIR}/liquidflow_final.pt'\nif os.path.exists(ckpt_path):\n    model.load_state_dict(torch.load(ckpt_path, map_location=device, weights_only=False)['model'])\n    print(f'Loaded: {ckpt_path}')\nmodel.eval(); torch.manual_seed(SEED)\nshape = (NUM_IMAGES, 3, IMG_SIZE, IMG_SIZE)\nwith torch.no_grad():\n    images = (euler_sample if SAMPLER=='euler' else heun_sample)(model, shape, num_steps=GEN_STEPS, device=device)\nimages = images.clamp(-1,1)*0.5+0.5\ngrid = make_grid_image(images, nrow=int(NUM_IMAGES**0.5))\ngrid.save(f'{OUTPUT_DIR}/generated_final.png')\nplt.figure(figsize=(10,10)); plt.imshow(grid)\nplt.title(f'LiquidFlow-{MODEL_SIZE} | {DATASET} {IMG_SIZE}×{IMG_SIZE} | {GEN_STEPS} steps ({SAMPLER})'); plt.axis('off'); plt.show()"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 7. 📱 Export for Mobile"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["model.eval()\nexample_x = torch.randn(1,3,IMG_SIZE,IMG_SIZE,device=device)\nexample_t = torch.tensor([0.5],device=device)\ntry:\n    traced = torch.jit.trace(model, (example_x, example_t))\n    ts_path = f'{OUTPUT_DIR}/liquidflow_mobile.pt'; traced.save(ts_path)\n    print(f'✅ TorchScript: {ts_path} ({os.path.getsize(ts_path)/1e6:.1f} MB)')\nexcept Exception as e: print(f'⚠️ TorchScript failed: {e}')\ntry:\n    onnx_path = f'{OUTPUT_DIR}/liquidflow.onnx'\n    torch.onnx.export(model.cpu(), (example_x.cpu(), example_t.cpu()), onnx_path, opset_version=14,\n        input_names=['image','timestep'], output_names=['velocity'],\n        dynamic_axes={'image':{0:'batch'},'timestep':{0:'batch'},'velocity':{0:'batch'}})\n    print(f'✅ ONNX: {onnx_path} ({os.path.getsize(onnx_path)/1e6:.1f} MB)'); model.to(device)\nexcept Exception as e: print(f'⚠️ ONNX failed: {e}'); model.to(device)"]},{"cell_type":"markdown","metadata":{},"source":["---\n","## 8. 🧪 Recommended Experiments\n","\n","| Goal | Dataset | Model | Size | Epochs | Time (T4) |\n","|------|---------|-------|------|--------|----------|\n","| Sanity check | CIFAR-10 | tiny | 32 | 20 | ~5 min |\n","| Baseline | CIFAR-10 | tiny | 128 | 100 | ~2 hrs |\n","| Quality | Flowers-102 | small | 128 | 200 | ~4 hrs |\n","| Faces | CelebA | small | 128 | 50 | ~6 hrs |\n","| High-res | CelebA | 512 | 512 | 100 | ~12 hrs |"]}],"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}