Upload LiquidDiffusion_Training.ipynb

LiquidDiffusion_Training.ipynb

@@ -1,876 +1,678 @@
 {
- "nbformat": 4,
- "nbformat_minor": 0,
- "metadata": {
-  "colab": {
-   "provenance": [],
-   "gpuType": "T4"
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# 🌊 LiquidDiffusion: Attention-Free Image Generation with Liquid Neural Networks\n",
+        "\n",
+        "**A novel image generation architecture** that replaces attention with Parallel CfC (Closed-form Continuous-depth) blocks from Liquid Neural Networks.\n",
+        "\n",
+        "## Key Innovations\n",
+        "- **No attention mechanism** — all spatial mixing via multi-scale depthwise convolutions\n",
+        "- **Fully parallelizable** — no sequential ODE solving loops (unlike original LTC/Neural ODE)\n",
+        "- **Diffusion timestep IS the liquid time constant** — natural CfC-diffusion bridge\n",
+        "- **Liquid relaxation residuals** — time-aware skip connections that adapt to noise level\n",
+        "- **Fits in 16GB VRAM** — designed for Colab free tier (T4 GPU)\n",
+        "\n",
+        "## Architecture Based On\n",
+        "- [CfC Networks](https://arxiv.org/abs/2106.13898) (Hasani et al., Nature Machine Intelligence 2022)\n",
+        "- [LiquidTAD](https://arxiv.org/abs/2604.18274) — parallel liquid relaxation\n",
+        "- [USM](https://arxiv.org/abs/2504.13499) — U-Shape architecture for diffusion\n",
+        "- [Rectified Flow](https://arxiv.org/abs/2209.03003) — simplest flow matching objective\n",
+        "\n",
+        "## Training: Rectified Flow\n",
+        "```\n",
+        "x_t = (1-t)*x0 + t*noise,  t ~ U[0,1]\n",
+        "Loss = MSE(model(x_t, t), noise - x0)  # velocity prediction\n",
+        "```\n",
+        "That's it — no noise schedule, no variance, just MSE on a straight-line velocity."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🔧 Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Install dependencies\n",
+        "!pip install -q torch torchvision datasets Pillow matplotlib tqdm accelerate"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Clone the repo\n",
+        "!git clone https://huggingface.co/krystv/liquid-diffusion\n",
+        "%cd liquid-diffusion"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "print(f'PyTorch: {torch.__version__}')\n",
+        "print(f'CUDA available: {torch.cuda.is_available()}')\n",
+        "if torch.cuda.is_available():\n",
+        "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
+        "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 📐 Architecture Overview\n",
+        "\n",
+        "The core innovation is the **ParallelCfCBlock** — a parallelized version of CfC (Closed-form Continuous-depth) networks adapted for 2D image features:\n",
+        "\n",
+        "```\n",
+        "CfC Equation (Hasani et al. 2022, Eq. 10):\n",
+        "    x(t) = σ(-f·t) ⊙ g + (1 - σ(-f·t)) ⊙ h\n",
+        "\n",
+        "Our adaptation for image generation:\n",
+        "    backbone = SiLU(PointwiseConv(DepthwiseConv(features)))  # shared spatial context\n",
+        "    f = Conv1x1(backbone)                                     # time-constant gate\n",
+        "    g = DWConv→SiLU→Conv1x1(backbone)                        # \"from\" state\n",
+        "    h = DWConv→SiLU→Conv1x1(backbone)                        # \"to\" state (attractor)\n",
+        "    gate = σ(time_a(t_emb) · f - time_b(t_emb))             # liquid time gate\n",
+        "    cfc_out = gate · g + (1 - gate) · h                      # CfC interpolation\n",
+        "    \n",
+        "    # Liquid relaxation (from LiquidTAD):\n",
+        "    α = exp(-softplus(ρ) · |t|)                              # time-aware residual weight\n",
+        "    output = α · input + (1 - α) · cfc_out                   # adapts to noise level\n",
+        "```\n",
+        "\n",
+        "The **diffusion timestep t** serves double duty:\n",
+        "1. Standard: conditions the denoiser via AdaLN scale/shift\n",
+        "2. Novel: acts as the CfC time parameter — controls interpolation between g and h\n",
+        "\n",
+        "This means: at low noise (t≈0), the gate is balanced → flexible processing.\n",
+        "At high noise (t≈1), the gate saturates → specialized denoising."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🧪 Quick Test (verify model works)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Run the test suite\n",
+        "!python test_model.py"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## ⚙️ Training Configuration\n",
+        "\n",
+        "Choose your config based on GPU and target resolution:\n",
+        "\n",
+        "| Config | Params | Resolution | Batch Size | VRAM | Training Time |\n",
+        "|--------|--------|-----------|------------|------|---------------|\n",
+        "| tiny | ~8M | 256×256 | 8 | ~6GB | ~3h (100K steps) |\n",
+        "| small | ~25M | 256×256 | 4 | ~10GB | ~6h (100K steps) |\n",
+        "| base | ~65M | 512×512 | 2 | ~14GB | ~12h (100K steps) |\n",
+        "\n",
+        "Recommended datasets:\n",
+        "- `huggan/CelebA-HQ` — 30K high-quality face images (256px)\n",
+        "- `huggan/flowers-102-categories` — flowers (various)\n",
+        "- `lambdalabs/naruto-blip-captions` — anime style (~1K)\n",
+        "- `Norod78/simpsons-blip-captions` — cartoon style\n",
+        "- Any folder of images"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Training Configuration {display-mode: \"form\"}\n",
+        "\n",
+        "#@markdown ### Model\n",
+        "model_size = \"tiny\"  #@param [\"tiny\", \"small\", \"base\"]\n",
+        "\n",
+        "#@markdown ### Data\n",
+        "dataset_name = \"huggan/CelebA-HQ\"  #@param {type:\"string\"}\n",
+        "image_column = \"image\"  #@param {type:\"string\"}\n",
+        "image_size = 256  #@param [64, 128, 256, 512] {type:\"integer\"}\n",
+        "max_samples = 0  #@param {type:\"integer\"}\n",
+        "\n",
+        "#@markdown ### Training\n",
+        "batch_size = 8  #@param {type:\"integer\"}\n",
+        "learning_rate = 1e-4  #@param {type:\"number\"}\n",
+        "weight_decay = 0.01  #@param {type:\"number\"}\n",
+        "total_steps = 100000  #@param {type:\"integer\"}\n",
+        "warmup_steps = 1000  #@param {type:\"integer\"}\n",
+        "grad_clip = 1.0  #@param {type:\"number\"}\n",
+        "ema_decay = 0.9999  #@param {type:\"number\"}\n",
+        "time_sampling = \"logit_normal\"  #@param [\"uniform\", \"logit_normal\"]\n",
+        "\n",
+        "#@markdown ### Sampling & Logging\n",
+        "sample_every = 2000  #@param {type:\"integer\"}\n",
+        "save_every = 5000  #@param {type:\"integer\"}\n",
+        "num_sample_steps = 50  #@param {type:\"integer\"}\n",
+        "num_sample_images = 4  #@param {type:\"integer\"}\n",
+        "\n",
+        "#@markdown ### Hardware\n",
+        "use_amp = True  #@param {type:\"boolean\"}\n",
+        "amp_dtype = \"float16\"  #@param [\"float16\", \"bfloat16\"]\n",
+        "num_workers = 2  #@param {type:\"integer\"}\n",
+        "\n",
+        "# Auto-adjust batch size for resolution\n",
+        "if image_size >= 512 and batch_size > 4:\n",
+        "    batch_size = min(batch_size, 2)\n",
+        "    print(f\"Auto-reduced batch_size to {batch_size} for {image_size}px\")\n",
+        "\n",
+        "if max_samples == 0:\n",
+        "    max_samples = None\n",
+        "\n",
+        "print(f\"\\nConfig: {model_size} model, {image_size}px, batch={batch_size}, lr={learning_rate}\")\n",
+        "print(f\"Dataset: {dataset_name}, time_sampling={time_sampling}\")\n",
+        "print(f\"Total steps: {total_steps:,}, AMP: {use_amp} ({amp_dtype})\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 📦 Load Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from datasets import load_dataset\n",
+        "from liquid_diffusion.trainer import ImageDataset\n",
+        "from torch.utils.data import DataLoader\n",
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "\n",
+        "# Load dataset\n",
+        "print(f\"Loading {dataset_name}...\")\n",
+        "dataset = ImageDataset(\n",
+        "    source=dataset_name,\n",
+        "    image_size=image_size,\n",
+        "    image_column=image_column,\n",
+        "    max_samples=max_samples,\n",
+        ")\n",
+        "print(f\"Dataset size: {len(dataset)} images\")\n",
+        "\n",
+        "dataloader = DataLoader(\n",
+        "    dataset, batch_size=batch_size, shuffle=True,\n",
+        "    num_workers=num_workers, pin_memory=True, drop_last=True,\n",
+        ")\n",
+        "\n",
+        "# Show some samples\n",
+        "sample_batch = next(iter(dataloader))\n",
+        "fig, axes = plt.subplots(1, min(4, batch_size), figsize=(16, 4))\n",
+        "for i, ax in enumerate(axes):\n",
+        "    img = sample_batch[i].permute(1, 2, 0).numpy() * 0.5 + 0.5  # [-1,1] -> [0,1]\n",
+        "    ax.imshow(np.clip(img, 0, 1))\n",
+        "    ax.axis('off')\n",
+        "plt.suptitle(f'Training samples ({image_size}×{image_size})')\n",
+        "plt.tight_layout()\n",
+        "plt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🏗️ Build Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from liquid_diffusion.model import (\n",
+        "    liquid_diffusion_tiny, liquid_diffusion_small, liquid_diffusion_base\n",
+        ")\n",
+        "\n",
+        "# Build model\n",
+        "model_factories = {\n",
+        "    'tiny': liquid_diffusion_tiny,\n",
+        "    'small': liquid_diffusion_small,\n",
+        "    'base': liquid_diffusion_base,\n",
+        "}\n",
+        "\n",
+        "model = model_factories[model_size]()\n",
+        "total_params, trainable_params = model.count_params()\n",
+        "print(f\"Model: liquid_diffusion_{model_size}\")\n",
+        "print(f\"Parameters: {total_params:,} ({total_params/1e6:.1f}M)\")\n",
+        "print(f\"Trainable: {trainable_params:,}\")\n",
+        "\n",
+        "# Quick forward pass test\n",
+        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+        "model = model.to(device)\n",
+        "test_x = torch.randn(1, 3, image_size, image_size, device=device)\n",
+        "test_t = torch.tensor([0.5], device=device)\n",
+        "with torch.no_grad():\n",
+        "    test_out = model(test_x, test_t)\n",
+        "print(f\"Forward pass OK: {test_x.shape} → {test_out.shape}\")\n",
+        "del test_x, test_out\n",
+        "if device == 'cuda':\n",
+        "    torch.cuda.empty_cache()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🚀 Train!"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "import time\n",
+        "import math\n",
+        "from tqdm.auto import tqdm\n",
+        "from torchvision.utils import save_image, make_grid\n",
+        "from liquid_diffusion.trainer import RectifiedFlowTrainer, get_cosine_schedule_with_warmup\n",
+        "\n",
+        "# Create output directories\n",
+        "os.makedirs('checkpoints', exist_ok=True)\n",
+        "os.makedirs('samples', exist_ok=True)\n",
+        "\n",
+        "# Build trainer\n",
+        "trainer = RectifiedFlowTrainer(\n",
+        "    model=model,\n",
+        "    lr=learning_rate,\n",
+        "    weight_decay=weight_decay,\n",
+        "    ema_decay=ema_decay,\n",
+        "    grad_clip=grad_clip,\n",
+        "    time_sampling=time_sampling,\n",
+        "    device=device,\n",
+        "    use_amp=use_amp,\n",
+        "    amp_dtype=amp_dtype,\n",
+        ")\n",
+        "\n",
+        "# Learning rate scheduler\n",
+        "scheduler = get_cosine_schedule_with_warmup(\n",
+        "    trainer.optimizer, warmup_steps, total_steps\n",
+        ")\n",
+        "\n",
+        "# Optional: resume from checkpoint\n",
+        "resume_path = 'checkpoints/latest.pt'\n",
+        "if os.path.exists(resume_path):\n",
+        "    trainer.load_checkpoint(resume_path)\n",
+        "    print(f\"Resumed from step {trainer.step}\")\n",
+        "\n",
+        "print(f\"\\n{'='*60}\")\n",
+        "print(f\"Starting training: {total_steps:,} steps\")\n",
+        "print(f\"Model: liquid_diffusion_{model_size} ({total_params/1e6:.1f}M params)\")\n",
+        "print(f\"Resolution: {image_size}×{image_size}, Batch: {batch_size}\")\n",
+        "print(f\"LR: {learning_rate}, Warmup: {warmup_steps}, AMP: {use_amp}\")\n",
+        "print(f\"{'='*60}\\n\")\n",
+        "\n",
+        "# Training loop\n",
+        "start_time = time.time()\n",
+        "data_iter = iter(dataloader)\n",
+        "pbar = tqdm(range(trainer.step, total_steps), desc='Training', dynamic_ncols=True)\n",
+        "loss_history = []\n",
+        "\n",
+        "for step in pbar:\n",
+        "    # Get batch (cycle through dataset)\n",
+        "    try:\n",
+        "        batch = next(data_iter)\n",
+        "    except StopIteration:\n",
+        "        data_iter = iter(dataloader)\n",
+        "        batch = next(data_iter)\n",
+        "    \n",
+        "    x0 = batch.to(device)\n",
+        "    \n",
+        "    # Train step\n",
+        "    metrics = trainer.train_step(x0)\n",
+        "    scheduler.step()\n",
+        "    \n",
+        "    # Logging\n",
+        "    loss_history.append(metrics['loss'])\n",
+        "    avg_loss = sum(loss_history[-100:]) / len(loss_history[-100:])\n",
+        "    lr_current = scheduler.get_last_lr()[0]\n",
+        "    \n",
+        "    pbar.set_postfix({\n",
+        "        'loss': f\"{metrics['loss']:.4f}\",\n",
+        "        'avg': f\"{avg_loss:.4f}\",\n",
+        "        'lr': f\"{lr_current:.6f}\",\n",
+        "        'gn': f\"{metrics['grad_norm']:.2f}\",\n",
+        "    })\n",
+        "    \n",
+        "    # Generate samples\n",
+        "    if (step + 1) % sample_every == 0 or step == 0:\n",
+        "        print(f\"\\nGenerating samples at step {step+1}...\")\n",
+        "        samples = trainer.sample(\n",
+        "            batch_size=num_sample_images, image_size=image_size,\n",
+        "            num_steps=num_sample_steps, use_ema=True\n",
+        "        )\n",
+        "        # Save grid\n",
+        "        grid = make_grid(samples * 0.5 + 0.5, nrow=int(math.sqrt(num_sample_images)), padding=2)\n",
+        "        save_image(grid, f'samples/step_{step+1:06d}.png')\n",
+        "        \n",
+        "        # Display\n",
+        "        fig, axes = plt.subplots(1, num_sample_images, figsize=(4*num_sample_images, 4))\n",
+        "        if num_sample_images == 1:\n",
+        "            axes = [axes]\n",
+        "        for i, ax in enumerate(axes):\n",
+        "            img = samples[i].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+        "            ax.imshow(np.clip(img, 0, 1))\n",
+        "            ax.axis('off')\n",
+        "        plt.suptitle(f'Step {step+1} (EMA samples, {num_sample_steps} Euler steps)')\n",
+        "        plt.tight_layout()\n",
+        "        plt.show()\n",
+        "    \n",
+        "    # Save checkpoint\n",
+        "    if (step + 1) % save_every == 0:\n",
+        "        trainer.save_checkpoint(f'checkpoints/step_{step+1:06d}.pt', extra={'config': {\n",
+        "            'model_size': model_size, 'image_size': image_size,\n",
+        "            'batch_size': batch_size, 'learning_rate': learning_rate,\n",
+        "        }})\n",
+        "        trainer.save_checkpoint('checkpoints/latest.pt')\n",
+        "        print(f\"Saved checkpoint at step {step+1}\")\n",
+        "    \n",
+        "    # Safety: check for NaN\n",
+        "    if math.isnan(metrics['loss']):\n",
+        "        print(\"\\n⚠️ NaN loss detected! Stopping training.\")\n",
+        "        print(\"Try: reduce learning_rate, increase grad_clip, or use smaller model\")\n",
+        "        break\n",
+        "\n",
+        "elapsed = time.time() - start_time\n",
+        "print(f\"\\nTraining complete! {trainer.step:,} steps in {elapsed/3600:.1f}h\")\n",
+        "print(f\"Final avg loss: {sum(loss_history[-100:])/len(loss_history[-100:]):.4f}\")\n",
+        "\n",
+        "# Final save\n",
+        "trainer.save_checkpoint('checkpoints/final.pt')\n",
+        "print(\"Saved final checkpoint.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 📊 Training Loss Curve"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "\n",
+        "if loss_history:\n",
+        "    # Smooth the loss\n",
+        "    window = min(100, len(loss_history) // 5 + 1)\n",
+        "    smoothed = np.convolve(loss_history, np.ones(window)/window, mode='valid')\n",
+        "    \n",
+        "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
+        "    \n",
+        "    ax1.plot(loss_history, alpha=0.3, label='Raw')\n",
+        "    ax1.plot(range(window-1, len(loss_history)), smoothed, label=f'Smoothed (w={window})')\n",
+        "    ax1.set_xlabel('Step')\n",
+        "    ax1.set_ylabel('Loss')\n",
+        "    ax1.set_title('Training Loss')\n",
+        "    ax1.legend()\n",
+        "    ax1.grid(True, alpha=0.3)\n",
+        "    \n",
+        "    ax2.plot(loss_history[-min(1000, len(loss_history)):], alpha=0.5)\n",
+        "    ax2.set_xlabel('Recent Steps')\n",
+        "    ax2.set_ylabel('Loss')\n",
+        "    ax2.set_title('Recent Loss (last 1000 steps)')\n",
+        "    ax2.grid(True, alpha=0.3)\n",
+        "    \n",
+        "    plt.tight_layout()\n",
+        "    plt.show()\n",
+        "else:\n",
+        "    print(\"No training history yet.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🎨 Generate Images"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Generation Settings {display-mode: \"form\"}\n",
+        "num_images = 8  #@param {type:\"integer\"}\n",
+        "sampling_steps = 50  #@param [25, 50, 100, 200] {type:\"integer\"}\n",
+        "use_ema_model = True  #@param {type:\"boolean\"}\n",
+        "\n",
+        "print(f\"Generating {num_images} images with {sampling_steps} Euler steps...\")\n",
+        "samples = trainer.sample(\n",
+        "    batch_size=num_images, image_size=image_size,\n",
+        "    num_steps=sampling_steps, use_ema=use_ema_model,\n",
+        ")\n",
+        "\n",
+        "# Display\n",
+        "ncols = min(4, num_images)\n",
+        "nrows = (num_images + ncols - 1) // ncols\n",
+        "fig, axes = plt.subplots(nrows, ncols, figsize=(4*ncols, 4*nrows))\n",
+        "if nrows == 1 and ncols == 1:\n",
+        "    axes = [[axes]]\n",
+        "elif nrows == 1:\n",
+        "    axes = [axes]\n",
+        "for i in range(num_images):\n",
+        "    r, c = i // ncols, i % ncols\n",
+        "    img = samples[i].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+        "    axes[r][c].imshow(np.clip(img, 0, 1))\n",
+        "    axes[r][c].axis('off')\n",
+        "# Hide unused axes\n",
+        "for i in range(num_images, nrows * ncols):\n",
+        "    r, c = i // ncols, i % ncols\n",
+        "    axes[r][c].axis('off')\n",
+        "plt.suptitle(f'LiquidDiffusion Samples ({sampling_steps} steps, {\"EMA\" if use_ema_model else \"online\"})')\n",
+        "plt.tight_layout()\n",
+        "plt.show()\n",
+        "\n",
+        "# Save\n",
+        "grid = make_grid(samples * 0.5 + 0.5, nrow=ncols, padding=2)\n",
+        "save_image(grid, 'samples/generated.png')\n",
+        "print(\"Saved to samples/generated.png\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 🔬 Visualize the Denoising Process"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Show step-by-step denoising\n",
+        "num_vis_steps = 10\n",
+        "total_euler_steps = 50\n",
+        "vis_interval = total_euler_steps // num_vis_steps\n",
+        "\n",
+        "model_vis = trainer.ema_model\n",
+        "model_vis.eval()\n",
+        "\n",
+        "z = torch.randn(1, 3, image_size, image_size, device=device)\n",
+        "dt = 1.0 / total_euler_steps\n",
+        "intermediates = [z.clone()]\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "    for i in range(total_euler_steps, 0, -1):\n",
+        "        t = torch.full((1,), i / total_euler_steps, device=device)\n",
+        "        v = model_vis(z, t)\n",
+        "        z = z - v * dt\n",
+        "        if (total_euler_steps - i + 1) % vis_interval == 0:\n",
+        "            intermediates.append(z.clone())\n",
+        "\n",
+        "intermediates.append(z.clamp(-1, 1))\n",
+        "\n",
+        "fig, axes = plt.subplots(1, len(intermediates), figsize=(3*len(intermediates), 3))\n",
+        "for idx, (ax, img_t) in enumerate(zip(axes, intermediates)):\n",
+        "    img = img_t[0].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+        "    ax.imshow(np.clip(img, 0, 1))\n",
+        "    ax.axis('off')\n",
+        "    if idx == 0:\n",
+        "        ax.set_title('Noise (t=1)')\n",
+        "    elif idx == len(intermediates) - 1:\n",
+        "        ax.set_title('Output (t=0)')\n",
+        "    else:\n",
+        "        ax.set_title(f't={1-idx*vis_interval/total_euler_steps:.1f}')\n",
+        "plt.suptitle('LiquidDiffusion Denoising Process')\n",
+        "plt.tight_layout()\n",
+        "plt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 💾 Save & Export Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Save final checkpoint\n",
+        "trainer.save_checkpoint('checkpoints/final.pt', extra={\n",
+        "    'config': {\n",
+        "        'model_size': model_size,\n",
+        "        'image_size': image_size,\n",
+        "        'total_params': total_params,\n",
+        "        'training_steps': trainer.step,\n",
+        "        'dataset': dataset_name,\n",
+        "    }\n",
+        "})\n",
+        "print(f\"Saved checkpoint: checkpoints/final.pt\")\n",
+        "print(f\"Model: liquid_diffusion_{model_size} ({total_params/1e6:.1f}M params)\")\n",
+        "print(f\"Trained for {trainer.step:,} steps on {dataset_name}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Optional: Push to Hugging Face Hub\n",
+        "# Uncomment and fill in your details:\n",
+        "\n",
+        "# from huggingface_hub import HfApi, login\n",
+        "# login()  # or use token\n",
+        "# api = HfApi()\n",
+        "# repo_id = \"your-username/liquid-diffusion-celebahq-256\"  # change this\n",
+        "# api.create_repo(repo_id, exist_ok=True)\n",
+        "# api.upload_file('checkpoints/final.pt', 'model.pt', repo_id)\n",
+        "# api.upload_folder('liquid_diffusion/', 'liquid_diffusion/', repo_id)\n",
+        "# print(f\"Uploaded to https://huggingface.co/{repo_id}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 📚 Architecture Details & Theory\n",
+        "\n",
+        "### Why Liquid Neural Networks for Image Generation?\n",
+        "\n",
+        "**Liquid Time-Constant (LTC) Networks** (Hasani et al., 2020) define neurons with input-dependent time constants:\n",
+        "\n",
+        "```\n",
+        "dx/dt = -[1/τ + f(x,I,θ)] · x + f(x,I,θ) · A\n",
+        "```\n",
+        "\n",
+        "The system time constant `τ_sys = τ/(1 + τ·f)` adapts dynamically based on input — the neuron speeds up or slows down its response depending on what it sees. This is the \"liquid\" property.\n",
+        "\n",
+        "**CfC (Closed-form Continuous-depth)** networks (Hasani et al., 2022) solve this ODE in closed form:\n",
+        "\n",
+        "```\n",
+        "x(t) = σ(-f·t) ⊙ g + (1 - σ(-f·t)) ⊙ h\n",
+        "```\n",
+        "\n",
+        "This eliminates the ODE solver — making CfC **fully parallelizable** while preserving the adaptive time constant behavior.\n",
+        "\n",
+        "### Our Innovation: CfC × Diffusion Timestep\n",
+        "\n",
+        "In diffusion models, the network must process images at different noise levels `t ∈ [0,1]`. We observe that:\n",
+        "\n",
+        "1. CfC's time parameter `t` controls interpolation between two learned states\n",
+        "2. Diffusion's noise level `t` controls how the denoiser should behave\n",
+        "3. **These are the same concept** — the CfC time parameter IS the diffusion timestep\n",
+        "\n",
+        "This gives us:\n",
+        "- At `t≈0` (clean images): σ(-f·t)≈0.5, balanced processing for detail refinement\n",
+        "- At `t≈1` (noisy images): σ(-f·t) saturates, specialized denoising\n",
+        "- The gate `f` is **input-dependent** — different image content gets different time responses\n",
+        "\n",
+        "### References\n",
+        "\n",
+        "1. Hasani et al., \"Liquid Time-constant Networks\" (AAAI 2021) — arxiv:2006.04439\n",
+        "2. Hasani et al., \"Closed-form Continuous-time Neural Networks\" (Nature MI 2022) — arxiv:2106.13898\n",
+        "3. LiquidTAD: Parallel liquid relaxation — arxiv:2604.18274\n",
+        "4. USM: U-Shape Mamba for diffusion — arxiv:2504.13499\n",
+        "5. DiffuSSM: Diffusion without attention — arxiv:2311.18257\n",
+        "6. Liu et al., \"Flow Straight and Fast: Rectified Flow\" (ICLR 2023) — arxiv:2209.03003"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.10.0"
+    }
   },
-  "kernelspec": {
-   "name": "python3",
-   "display_name": "Python 3"
-  },
-  "accelerator": "GPU"
- },
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 🌊 LiquidDiffusion: Attention-Free Image Generation with Liquid Neural Networks\n",
-    "\n",
-    "**A novel image generation model** combining:\n",
-    "- **Liquid Neural Networks** (CfC — Closed-form Continuous-depth) for adaptive, time-aware processing\n",
-    "- **Rectified Flow** for simple, stable training (MSE velocity prediction)\n",
-    "- **Zero attention** — fully convolutional with multi-scale spatial mixing\n",
-    "- **Fully parallelizable** — no sequential ODE loops or recurrence\n",
-    "\n",
-    "### Key Innovation\n",
-    "The diffusion timestep serves as the **liquid time constant** — the CfC gate `σ(-f·t)` naturally adapts the network's behavior based on noise level, giving input-dependent processing without attention.\n",
-    "\n",
-    "### References\n",
-    "- CfC Networks: [Hasani et al., Nature MI 2022](https://arxiv.org/abs/2106.13898)\n",
-    "- LiquidTAD (parallel CfC): [arxiv 2604.18274](https://arxiv.org/abs/2604.18274)\n",
-    "- USM (U-Shape Mamba): [arxiv 2504.13499](https://arxiv.org/abs/2504.13499)\n",
-    "- Rectified Flow: [Liu et al., ICLR 2023](https://arxiv.org/abs/2209.03003)\n",
-    "\n",
-    "**Model repo**: [krystv/liquid-diffusion](https://huggingface.co/krystv/liquid-diffusion)\n",
-    "\n",
-    "---"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## ⚙️ Configuration\n",
-    "\n",
-    "Choose your settings here. Everything is configurable from this one cell."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ============================================================================\n",
-    "# CONFIGURATION — Edit this cell to customize your training\n",
-    "# ============================================================================\n",
-    "\n",
-    "# --- Model Size ---\n",
-    "# 'tiny'  = ~23M params, best for 256px, fits easily in T4 16GB\n",
-    "# 'small' = ~69M params, better quality 256px, tight fit on T4\n",
-    "# 'base'  = ~154M params, for 512px (needs A100 or reduce batch size)\n",
-    "# 'custom' = define your own channels/blocks below\n",
-    "MODEL_SIZE = 'tiny'  # @param ['tiny', 'small', 'base', 'custom']\n",
-    "\n",
-    "# Custom model config (only used if MODEL_SIZE='custom')\n",
-    "CUSTOM_CHANNELS = [48, 96, 192]  # channel dims per stage\n",
-    "CUSTOM_BLOCKS = [1, 2, 3]        # blocks per stage\n",
-    "CUSTOM_T_DIM = 192               # time embedding dimension\n",
-    "\n",
-    "# --- Image Resolution ---\n",
-    "IMAGE_SIZE = 256  # @param [64, 128, 256, 512] {type:\"integer\"}\n",
-    "\n",
-    "# --- Dataset ---\n",
-    "# Options:\n",
-    "#   'huggan/CelebA-HQ'          - 30K celebrity faces (256px native)\n",
-    "#   'huggan/flowers-102-categories' - Flowers dataset\n",
-    "#   'huggan/anime-faces'        - Anime faces\n",
-    "#   'lambdalabs/naruto-blip-captions' - Naruto illustrations\n",
-    "#   'jlbaker361/CelebA-HQ-256'  - CelebA-HQ at 256px\n",
-    "#   Or any HF dataset with an 'image' column\n",
-    "#   Or a local folder path with images\n",
-    "DATASET = 'huggan/CelebA-HQ'  # @param {type:\"string\"}\n",
-    "IMAGE_COLUMN = 'image'  # column name in HF dataset containing images\n",
-    "MAX_SAMPLES = None  # Set to e.g. 1000 for quick testing, None for full dataset\n",
-    "\n",
-    "# --- Training ---\n",
-    "BATCH_SIZE = 8        # @param {type:\"integer\"}\n",
-    "LEARNING_RATE = 1e-4  # @param {type:\"number\"}\n",
-    "WEIGHT_DECAY = 0.01   # @param {type:\"number\"}\n",
-    "NUM_EPOCHS = 100      # @param {type:\"integer\"}\n",
-    "GRAD_CLIP = 1.0       # @param {type:\"number\"}\n",
-    "EMA_DECAY = 0.9999    # @param {type:\"number\"}\n",
-    "NUM_WORKERS = 2       # DataLoader workers\n",
-    "\n",
-    "# --- Time Sampling ---\n",
-    "# 'logit_normal' (from SD3) = more weight on intermediate timesteps\n",
-    "# 'uniform' = standard\n",
-    "TIME_SAMPLING = 'logit_normal'  # @param ['uniform', 'logit_normal']\n",
-    "\n",
-    "# --- Mixed Precision ---\n",
-    "USE_AMP = True  # @param {type:\"boolean\"}\n",
-    "AMP_DTYPE = 'float16'  # @param ['float16', 'bfloat16']\n",
-    "\n",
-    "# --- Sampling ---\n",
-    "SAMPLE_EVERY = 500       # Generate samples every N steps\n",
-    "NUM_SAMPLE_IMAGES = 8    # Images to generate per sample\n",
-    "NUM_EULER_STEPS = 50     # Euler ODE steps (more = better quality)\n",
-    "\n",
-    "# --- Checkpointing ---\n",
-    "SAVE_EVERY = 2000        # Save checkpoint every N steps\n",
-    "OUTPUT_DIR = './outputs'  # Where to save everything\n",
-    "RESUME_FROM = None        # Path to checkpoint to resume from, or None\n",
-    "\n",
-    "# --- Logging ---\n",
-    "LOG_EVERY = 50  # Print loss every N steps\n",
-    "\n",
-    "print(f\"Config: {MODEL_SIZE} model, {IMAGE_SIZE}px, dataset={DATASET}\")\n",
-    "print(f\"Training: bs={BATCH_SIZE}, lr={LEARNING_RATE}, epochs={NUM_EPOCHS}\")\n",
-    "print(f\"AMP: {USE_AMP} ({AMP_DTYPE}), Time sampling: {TIME_SAMPLING}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 📦 Install Dependencies"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install -q datasets huggingface_hub Pillow matplotlib\n",
-    "\n",
-    "import torch\n",
-    "print(f\"PyTorch: {torch.__version__}\")\n",
-    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
-    "if torch.cuda.is_available():\n",
-    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
-    "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 🏗️ Model Architecture\n",
-    "\n",
-    "The complete LiquidDiffusion model — defined inline so you can inspect and modify everything."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "import copy\n",
-    "import os\n",
-    "import time\n",
-    "import json\n",
-    "from glob import glob\n",
-    "\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F\n",
-    "from torch.utils.data import DataLoader, Dataset\n",
-    "from torchvision import transforms\n",
-    "from torchvision.utils import save_image, make_grid\n",
-    "\n",
-    "\n",
-    "# ========================= TIME EMBEDDING =========================\n",
-    "\n",
-    "class SinusoidalTimeEmbedding(nn.Module):\n",
-    "    \"\"\"Sinusoidal position encoding + MLP for timestep embedding.\"\"\"\n",
-    "    def __init__(self, dim, max_period=10000):\n",
-    "        super().__init__()\n",
-    "        self.dim = dim\n",
-    "        self.max_period = max_period\n",
-    "        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.SiLU(), nn.Linear(dim*4, dim))\n",
-    "\n",
-    "    def forward(self, t):\n",
-    "        half = self.dim // 2\n",
-    "        freqs = torch.exp(-math.log(self.max_period) * torch.arange(half, device=t.device, dtype=t.dtype) / half)\n",
-    "        args = t[:, None] * freqs[None, :]\n",
-    "        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)\n",
-    "        if self.dim % 2: emb = F.pad(emb, (0, 1))\n",
-    "        return self.mlp(emb)\n",
-    "\n",
-    "\n",
-    "# ========================= ADAPTIVE LAYER NORM =========================\n",
-    "\n",
-    "class AdaLN(nn.Module):\n",
-    "    \"\"\"Adaptive LayerNorm: norm(x) * (1+scale(t)) + shift(t)\"\"\"\n",
-    "    def __init__(self, dim, cond_dim):\n",
-    "        super().__init__()\n",
-    "        ng = min(32, dim)\n",
-    "        while dim % ng != 0: ng -= 1\n",
-    "        self.norm = nn.GroupNorm(ng, dim, affine=False)\n",
-    "        self.proj = nn.Sequential(nn.SiLU(), nn.Linear(cond_dim, dim * 2))\n",
-    "\n",
-    "    def forward(self, x, t_emb):\n",
-    "        s, sh = self.proj(t_emb).chunk(2, dim=1)\n",
-    "        return self.norm(x) * (1 + s[:,:,None,None]) + sh[:,:,None,None]\n",
-    "\n",
-    "\n",
-    "# ========================= PARALLEL CfC BLOCK =========================\n",
-    "\n",
-    "class ParallelCfCBlock(nn.Module):\n",
-    "    \"\"\"\n",
-    "    Parallel Closed-form Continuous-depth (CfC) block.\n",
-    "    \n",
-    "    CfC Eq.10: x(t) = σ(-f·t) ⊙ g + (1-σ(-f·t)) ⊙ h\n",
-    "    \n",
-    "    • f/g/h heads operate on 2D feature maps (depthwise conv)\n",
-    "    • Diffusion timestep t IS the liquid time constant\n",
-    "    • No recurrence, no ODE solver — fully parallel\n",
-    "    • Liquid relaxation: α·residual + (1-α)·CfC_out, α=exp(-λ·t)\n",
-    "    \"\"\"\n",
-    "    def __init__(self, dim, t_dim, expand_ratio=2.0, kernel_size=7, dropout=0.0):\n",
-    "        super().__init__()\n",
-    "        hidden = int(dim * expand_ratio)\n",
-    "        self.backbone_dw = nn.Conv2d(dim, dim, kernel_size, padding=kernel_size//2, groups=dim)\n",
-    "        self.backbone_pw = nn.Conv2d(dim, hidden, 1)\n",
-    "        self.backbone_act = nn.SiLU()\n",
-    "        self.f_head = nn.Conv2d(hidden, dim, 1)\n",
-    "        self.g_head = nn.Sequential(\n",
-    "            nn.Conv2d(hidden, hidden, kernel_size, padding=kernel_size//2, groups=hidden),\n",
-    "            nn.SiLU(), nn.Conv2d(hidden, dim, 1))\n",
-    "        self.h_head = nn.Sequential(\n",
-    "            nn.Conv2d(hidden, hidden, kernel_size, padding=kernel_size//2, groups=hidden),\n",
-    "            nn.SiLU(), nn.Conv2d(hidden, dim, 1))\n",
-    "        self.time_a = nn.Linear(t_dim, dim)\n",
-    "        self.time_b = nn.Linear(t_dim, dim)\n",
-    "        self.rho = nn.Parameter(torch.zeros(1, dim, 1, 1))\n",
-    "        self.output_gate = nn.Sequential(nn.SiLU(), nn.Linear(t_dim, dim))\n",
-    "        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()\n",
-    "\n",
-    "    def forward(self, x, t_emb):\n",
-    "        residual = x\n",
-    "        backbone = self.backbone_act(self.backbone_pw(self.backbone_dw(x)))\n",
-    "        f, g, h = self.f_head(backbone), self.g_head(backbone), self.h_head(backbone)\n",
-    "        ta = self.time_a(t_emb)[:,:,None,None]\n",
-    "        tb = self.time_b(t_emb)[:,:,None,None]\n",
-    "        gate = torch.sigmoid(ta * f - tb)\n",
-    "        cfc_out = self.dropout(gate * g + (1.0 - gate) * h)\n",
-    "        t_scalar = t_emb.mean(dim=1, keepdim=True)[:,:,None,None]\n",
-    "        lam = F.softplus(self.rho) + 1e-6\n",
-    "        alpha = torch.exp(-lam * t_scalar.abs().clamp(min=0.01))\n",
-    "        out = alpha * residual + (1.0 - alpha) * cfc_out\n",
-    "        return out * torch.sigmoid(self.output_gate(t_emb))[:,:,None,None]\n",
-    "\n",
-    "\n",
-    "# ========================= MULTI-SCALE SPATIAL MIXING =========================\n",
-    "\n",
-    "class MultiScaleSpatialMix(nn.Module):\n",
-    "    \"\"\"Multi-scale depthwise conv + global pooling (replaces attention).\"\"\"\n",
-    "    def __init__(self, dim, t_dim):\n",
-    "        super().__init__()\n",
-    "        self.dw3 = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)\n",
-    "        self.dw5 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)\n",
-    "        self.dw7 = nn.Conv2d(dim, dim, 7, padding=3, groups=dim)\n",
-    "        self.global_pool = nn.AdaptiveAvgPool2d(1)\n",
-    "        self.global_proj = nn.Conv2d(dim, dim, 1)\n",
-    "        self.merge = nn.Conv2d(dim*4, dim, 1)\n",
-    "        self.act = nn.SiLU()\n",
-    "        self.adaln = AdaLN(dim, t_dim)\n",
-    "\n",
-    "    def forward(self, x, t_emb):\n",
-    "        xn = self.adaln(x, t_emb)\n",
-    "        return x + self.act(self.merge(torch.cat([\n",
-    "            self.dw3(xn), self.dw5(xn), self.dw7(xn),\n",
-    "            self.global_proj(self.global_pool(xn)).expand_as(xn)], dim=1)))\n",
-    "\n",
-    "\n",
-    "# ========================= LIQUID DIFFUSION BLOCK =========================\n",
-    "\n",
-    "class LiquidDiffusionBlock(nn.Module):\n",
-    "    \"\"\"AdaLN → CfC → SpatialMix → FF with residual scaling.\"\"\"\n",
-    "    def __init__(self, dim, t_dim, expand_ratio=2.0, kernel_size=7, dropout=0.0):\n",
-    "        super().__init__()\n",
-    "        self.adaln1 = AdaLN(dim, t_dim)\n",
-    "        self.cfc = ParallelCfCBlock(dim, t_dim, expand_ratio, kernel_size, dropout)\n",
-    "        self.spatial_mix = MultiScaleSpatialMix(dim, t_dim)\n",
-    "        self.adaln2 = AdaLN(dim, t_dim)\n",
-    "        ff_dim = int(dim * expand_ratio)\n",
-    "        self.ff = nn.Sequential(nn.Conv2d(dim, ff_dim, 1), nn.SiLU(), nn.Conv2d(ff_dim, dim, 1))\n",
-    "        self.res_scale = nn.Parameter(torch.ones(1) * 0.1)\n",
-    "\n",
-    "    def forward(self, x, t_emb):\n",
-    "        x = x + self.res_scale * self.cfc(self.adaln1(x, t_emb), t_emb)\n",
-    "        x = self.spatial_mix(x, t_emb)\n",
-    "        x = x + self.res_scale * self.ff(self.adaln2(x, t_emb))\n",
-    "        return x\n",
-    "\n",
-    "\n",
-    "# ========================= SPATIAL OPS =========================\n",
-    "\n",
-    "class DownSample(nn.Module):\n",
-    "    def __init__(self, in_d, out_d):\n",
-    "        super().__init__()\n",
-    "        self.conv = nn.Conv2d(in_d, out_d, 3, stride=2, padding=1)\n",
-    "    def forward(self, x): return self.conv(x)\n",
-    "\n",
-    "class UpSample(nn.Module):\n",
-    "    def __init__(self, in_d, out_d):\n",
-    "        super().__init__()\n",
-    "        self.conv = nn.Conv2d(in_d, out_d, 3, padding=1)\n",
-    "    def forward(self, x): return self.conv(F.interpolate(x, scale_factor=2, mode='nearest'))\n",
-    "\n",
-    "class SkipFusion(nn.Module):\n",
-    "    def __init__(self, dim, t_dim):\n",
-    "        super().__init__()\n",
-    "        self.proj = nn.Conv2d(dim*2, dim, 1)\n",
-    "        self.gate = nn.Sequential(nn.SiLU(), nn.Linear(t_dim, dim), nn.Sigmoid())\n",
-    "    def forward(self, x, skip, t_emb):\n",
-    "        m = self.proj(torch.cat([x, skip], dim=1))\n",
-    "        g = self.gate(t_emb)[:,:,None,None]\n",
-    "        return m * g + x * (1 - g)\n",
-    "\n",
-    "\n",
-    "# ========================= LIQUID DIFFUSION U-NET =========================\n",
-    "\n",
-    "class LiquidDiffusionUNet(nn.Module):\n",
-    "    \"\"\"\n",
-    "    LiquidDiffusion: Attention-Free Image Generation with Liquid Neural Networks.\n",
-    "    U-Net with Parallel CfC blocks. Diffusion timestep = liquid time constant.\n",
-    "    \"\"\"\n",
-    "    def __init__(self, in_channels=3, channels=None, blocks_per_stage=None,\n",
-    "                 t_dim=256, expand_ratio=2.0, kernel_size=7, dropout=0.0):\n",
-    "        super().__init__()\n",
-    "        channels = channels or [64, 128, 256]\n",
-    "        blocks_per_stage = blocks_per_stage or [2, 2, 4]\n",
-    "        assert len(channels) == len(blocks_per_stage)\n",
-    "        self.channels, self.num_stages = channels, len(channels)\n",
-    "        \n",
-    "        self.time_embed = SinusoidalTimeEmbedding(t_dim)\n",
-    "        self.stem = nn.Sequential(\n",
-    "            nn.Conv2d(in_channels, channels[0], 3, padding=1), nn.SiLU(),\n",
-    "            nn.Conv2d(channels[0], channels[0], 3, padding=1))\n",
-    "        \n",
-    "        # Encoder\n",
-    "        self.encoder_blocks = nn.ModuleList()\n",
-    "        self.downsamplers = nn.ModuleList()\n",
-    "        for i in range(self.num_stages):\n",
-    "            stage = nn.ModuleList([LiquidDiffusionBlock(channels[i], t_dim, expand_ratio, kernel_size, dropout)\n",
-    "                                   for _ in range(blocks_per_stage[i])])\n",
-    "            self.encoder_blocks.append(stage)\n",
-    "            if i < self.num_stages - 1:\n",
-    "                self.downsamplers.append(DownSample(channels[i], channels[i+1]))\n",
-    "        \n",
-    "        # Bottleneck\n",
-    "        self.bottleneck = nn.ModuleList([\n",
-    "            LiquidDiffusionBlock(channels[-1], t_dim, expand_ratio, kernel_size, dropout),\n",
-    "            LiquidDiffusionBlock(channels[-1], t_dim, expand_ratio, kernel_size, dropout)])\n",
-    "        \n",
-    "        # Decoder\n",
-    "        self.decoder_blocks, self.upsamplers, self.skip_fusions = nn.ModuleList(), nn.ModuleList(), nn.ModuleList()\n",
-    "        for i in range(self.num_stages-1, -1, -1):\n",
-    "            if i < self.num_stages - 1:\n",
-    "                self.upsamplers.append(UpSample(channels[i+1], channels[i]))\n",
-    "                self.skip_fusions.append(SkipFusion(channels[i], t_dim))\n",
-    "            self.decoder_blocks.append(nn.ModuleList([\n",
-    "                LiquidDiffusionBlock(channels[i], t_dim, expand_ratio, kernel_size, dropout)\n",
-    "                for _ in range(blocks_per_stage[i])]))\n",
-    "        \n",
-    "        hg = min(32, channels[0])\n",
-    "        while channels[0] % hg != 0: hg -= 1\n",
-    "        self.head = nn.Sequential(nn.GroupNorm(hg, channels[0]), nn.SiLU(),\n",
-    "                                   nn.Conv2d(channels[0], in_channels, 3, padding=1))\n",
-    "        nn.init.zeros_(self.head[-1].weight); nn.init.zeros_(self.head[-1].bias)\n",
-    "\n",
-    "    def forward(self, x, t):\n",
-    "        t_emb = self.time_embed(t)\n",
-    "        h = self.stem(x)\n",
-    "        skips = []\n",
-    "        for i in range(self.num_stages):\n",
-    "            for blk in self.encoder_blocks[i]: h = blk(h, t_emb)\n",
-    "            skips.append(h)\n",
-    "            if i < self.num_stages - 1: h = self.downsamplers[i](h)\n",
-    "        for blk in self.bottleneck: h = blk(h, t_emb)\n",
-    "        up_idx = 0\n",
-    "        for dec_i in range(self.num_stages):\n",
-    "            si = self.num_stages - 1 - dec_i\n",
-    "            if dec_i > 0:\n",
-    "                h = self.upsamplers[up_idx](h)\n",
-    "                h = self.skip_fusions[up_idx](h, skips[si], t_emb)\n",
-    "                up_idx += 1\n",
-    "            for blk in self.decoder_blocks[dec_i]: h = blk(h, t_emb)\n",
-    "        return self.head(h)\n",
-    "\n",
-    "    def count_params(self):\n",
-    "        return sum(p.numel() for p in self.parameters()), sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
-    "\n",
-    "\n",
-    "print(\"✅ Model architecture defined.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 🔧 Build Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Build model based on config\n",
-    "MODEL_CONFIGS = {\n",
-    "    'tiny':  dict(channels=[64, 128, 256],  blocks_per_stage=[2, 2, 4], t_dim=256),\n",
-    "    'small': dict(channels=[96, 192, 384],  blocks_per_stage=[2, 3, 6], t_dim=384),\n",
-    "    'base':  dict(channels=[128, 256, 512], blocks_per_stage=[2, 4, 8], t_dim=512),\n",
-    "}\n",
-    "\n",
-    "if MODEL_SIZE == 'custom':\n",
-    "    config = dict(channels=CUSTOM_CHANNELS, blocks_per_stage=CUSTOM_BLOCKS, t_dim=CUSTOM_T_DIM)\n",
-    "else:\n",
-    "    config = MODEL_CONFIGS[MODEL_SIZE]\n",
-    "\n",
-    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
-    "model = LiquidDiffusionUNet(**config).to(device)\n",
-    "total_params, trainable_params = model.count_params()\n",
-    "\n",
-    "print(f\"Model: {MODEL_SIZE}\")\n",
-    "print(f\"  Channels: {config['channels']}\")\n",
-    "print(f\"  Blocks:   {config['blocks_per_stage']}\")\n",
-    "print(f\"  t_dim:    {config['t_dim']}\")\n",
-    "print(f\"  Total parameters: {total_params:,} ({total_params/1e6:.1f}M)\")\n",
-    "print(f\"  Device: {device}\")\n",
-    "\n",
-    "# Quick forward pass test\n",
-    "with torch.no_grad():\n",
-    "    test_x = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=device)\n",
-    "    test_t = torch.tensor([0.5], device=device)\n",
-    "    test_out = model(test_x, test_t)\n",
-    "    print(f\"  Forward pass OK: {test_x.shape} → {test_out.shape}\")\n",
-    "    del test_x, test_out\n",
-    "    if device == 'cuda':\n",
-    "        torch.cuda.empty_cache()\n",
-    "        print(f\"  VRAM after test: {torch.cuda.memory_allocated()/1e9:.2f} GB\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 📊 Load Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from PIL import Image\n",
-    "\n",
-    "class ImageDataset(Dataset):\n",
-    "    def __init__(self, source, image_size=256, image_column='image', max_samples=None):\n",
-    "        self.image_column = image_column\n",
-    "        self.transform = transforms.Compose([\n",
-    "            transforms.Resize(image_size, interpolation=transforms.InterpolationMode.LANCZOS),\n",
-    "            transforms.CenterCrop(image_size),\n",
-    "            transforms.RandomHorizontalFlip(),\n",
-    "            transforms.ToTensor(),\n",
-    "            transforms.Normalize([0.5], [0.5]),\n",
-    "        ])\n",
-    "        if os.path.isdir(source):\n",
-    "            self.files = sorted(sum([glob(os.path.join(source, '**', f'*.{e}'), recursive=True)\n",
-    "                                     for e in ['png','jpg','jpeg','webp','bmp']], []))\n",
-    "            if max_samples: self.files = self.files[:max_samples]\n",
-    "            self.mode = 'folder'\n",
-    "        else:\n",
-    "            from datasets import load_dataset\n",
-    "            self.data = load_dataset(source, split='train')\n",
-    "            if max_samples: self.data = self.data.select(range(min(max_samples, len(self.data))))\n",
-    "            self.mode = 'hf'\n",
-    "    \n",
-    "    def __len__(self):\n",
-    "        return len(self.files) if self.mode == 'folder' else len(self.data)\n",
-    "    \n",
-    "    def __getitem__(self, idx):\n",
-    "        if self.mode == 'folder':\n",
-    "            img = Image.open(self.files[idx]).convert('RGB')\n",
-    "        else:\n",
-    "            img = self.data[idx][self.image_column]\n",
-    "            if not hasattr(img, 'convert'): img = Image.fromarray(img)\n",
-    "            img = img.convert('RGB')\n",
-    "        return self.transform(img)\n",
-    "\n",
-    "# Load dataset\n",
-    "print(f\"Loading dataset: {DATASET}\")\n",
-    "dataset = ImageDataset(DATASET, IMAGE_SIZE, IMAGE_COLUMN, MAX_SAMPLES)\n",
-    "dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True,\n",
-    "                        num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)\n",
-    "\n",
-    "print(f\"Dataset size: {len(dataset):,} images\")\n",
-    "print(f\"Steps per epoch: {len(dataloader):,}\")\n",
-    "print(f\"Total steps: ~{len(dataloader) * NUM_EPOCHS:,}\")\n",
-    "\n",
-    "# Show sample\n",
-    "import matplotlib.pyplot as plt\n",
-    "sample_batch = next(iter(dataloader))\n",
-    "fig, axes = plt.subplots(1, min(8, BATCH_SIZE), figsize=(16, 2))\n",
-    "for i, ax in enumerate(axes):\n",
-    "    img = (sample_batch[i].permute(1,2,0) * 0.5 + 0.5).clamp(0,1)\n",
-    "    ax.imshow(img); ax.axis('off')\n",
-    "plt.suptitle(f'Training samples ({IMAGE_SIZE}px)'); plt.tight_layout(); plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 🚀 Training Loop"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "from IPython.display import clear_output, display\n",
-    "\n",
-    "# Setup\n",
-    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
-    "os.makedirs(f'{OUTPUT_DIR}/samples', exist_ok=True)\n",
-    "os.makedirs(f'{OUTPUT_DIR}/checkpoints', exist_ok=True)\n",
-    "\n",
-    "optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE,\n",
-    "                               weight_decay=WEIGHT_DECAY, betas=(0.9, 0.999))\n",
-    "\n",
-    "# Cosine LR schedule with warmup\n",
-    "total_steps = len(dataloader) * NUM_EPOCHS\n",
-    "warmup_steps = min(1000, total_steps // 10)\n",
-    "\n",
-    "def lr_lambda(step):\n",
-    "    if step < warmup_steps:\n",
-    "        return float(step) / float(max(1, warmup_steps))\n",
-    "    progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))\n",
-    "    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))\n",
-    "\n",
-    "scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n",
-    "\n",
-    "# EMA model\n",
-    "ema_model = copy.deepcopy(model).eval()\n",
-    "for p in ema_model.parameters(): p.requires_grad_(False)\n",
-    "\n",
-    "# AMP\n",
-    "scaler = torch.amp.GradScaler('cuda', enabled=(USE_AMP and device == 'cuda'))\n",
-    "amp_dtype = getattr(torch, AMP_DTYPE) if USE_AMP and device == 'cuda' else torch.float32\n",
-    "\n",
-    "# Time sampling\n",
-    "def sample_time(bs):\n",
-    "    eps = 1e-5\n",
-    "    if TIME_SAMPLING == 'uniform':\n",
-    "        return torch.rand(bs, device=device) * (1 - 2*eps) + eps\n",
-    "    else:  # logit_normal\n",
-    "        return torch.sigmoid(torch.randn(bs, device=device)).clamp(eps, 1-eps)\n",
-    "\n",
-    "# Resume if requested\n",
-    "global_step = 0\n",
-    "start_epoch = 0\n",
-    "all_losses = []\n",
-    "\n",
-    "if RESUME_FROM and os.path.exists(RESUME_FROM):\n",
-    "    ckpt = torch.load(RESUME_FROM, map_location=device, weights_only=False)\n",
-    "    model.load_state_dict(ckpt['model'])\n",
-    "    ema_model.load_state_dict(ckpt['ema_model'])\n",
-    "    optimizer.load_state_dict(ckpt['optimizer'])\n",
-    "    global_step = ckpt.get('step', 0)\n",
-    "    start_epoch = ckpt.get('epoch', 0)\n",
-    "    all_losses = ckpt.get('losses', [])\n",
-    "    print(f\"Resumed from step {global_step}, epoch {start_epoch}\")\n",
-    "\n",
-    "\n",
-    "@torch.no_grad()\n",
-    "def generate_samples(step):\n",
-    "    \"\"\"Generate and save sample images.\"\"\"\n",
-    "    ema_model.eval()\n",
-    "    z = torch.randn(NUM_SAMPLE_IMAGES, 3, IMAGE_SIZE, IMAGE_SIZE, device=device)\n",
-    "    dt = 1.0 / NUM_EULER_STEPS\n",
-    "    for i in range(NUM_EULER_STEPS, 0, -1):\n",
-    "        t = torch.full((NUM_SAMPLE_IMAGES,), i / NUM_EULER_STEPS, device=device)\n",
-    "        with torch.amp.autocast(device, dtype=amp_dtype, enabled=USE_AMP and device=='cuda'):\n",
-    "            v = ema_model(z, t)\n",
-    "        if USE_AMP and amp_dtype == torch.float16: v = v.float()\n",
-    "        z = z - v * dt\n",
-    "    z = z.clamp(-1, 1)\n",
-    "    grid = make_grid(z * 0.5 + 0.5, nrow=int(math.sqrt(NUM_SAMPLE_IMAGES)), padding=2)\n",
-    "    save_image(grid, f'{OUTPUT_DIR}/samples/step_{step:06d}.png')\n",
-    "    return z\n",
-    "\n",
-    "\n",
-    "# ========== TRAINING LOOP ==========\n",
-    "print(f\"\\n{'='*60}\")\n",
-    "print(f\"Starting training: {NUM_EPOCHS} epochs, {total_steps:,} total steps\")\n",
-    "print(f\"Warmup: {warmup_steps} steps, LR: {LEARNING_RATE}\")\n",
-    "print(f\"{'='*60}\\n\")\n",
-    "\n",
-    "train_start = time.time()\n",
-    "epoch_losses = []\n",
-    "\n",
-    "for epoch in range(start_epoch, NUM_EPOCHS):\n",
-    "    model.train()\n",
-    "    epoch_loss = 0\n",
-    "    \n",
-    "    for batch_idx, x0 in enumerate(dataloader):\n",
-    "        x0 = x0.to(device, non_blocking=True)\n",
-    "        \n",
-    "        # Rectified Flow: x_t = (1-t)*x0 + t*x1, target = x1 - x0\n",
-    "        x1 = torch.randn_like(x0)\n",
-    "        t = sample_time(x0.shape[0])\n",
-    "        te = t[:, None, None, None]\n",
-    "        x_t = (1 - te) * x0 + te * x1\n",
-    "        v_target = x1 - x0\n",
-    "        \n",
-    "        with torch.amp.autocast(device, dtype=amp_dtype, enabled=USE_AMP and device=='cuda'):\n",
-    "            v_pred = model(x_t, t)\n",
-    "            loss = F.mse_loss(v_pred, v_target)\n",
-    "        \n",
-    "        optimizer.zero_grad(set_to_none=True)\n",
-    "        scaler.scale(loss).backward()\n",
-    "        if GRAD_CLIP > 0:\n",
-    "            scaler.unscale_(optimizer)\n",
-    "            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)\n",
-    "        scaler.step(optimizer)\n",
-    "        scaler.update()\n",
-    "        scheduler.step()\n",
-    "        \n",
-    "        # EMA update\n",
-    "        with torch.no_grad():\n",
-    "            for ep, mp in zip(ema_model.parameters(), model.parameters()):\n",
-    "                ep.data.mul_(EMA_DECAY).add_(mp.data, alpha=1-EMA_DECAY)\n",
-    "        \n",
-    "        global_step += 1\n",
-    "        loss_val = loss.item()\n",
-    "        all_losses.append(loss_val)\n",
-    "        epoch_loss += loss_val\n",
-    "        \n",
-    "        # Logging\n",
-    "        if global_step % LOG_EVERY == 0:\n",
-    "            avg_loss = sum(all_losses[-LOG_EVERY:]) / LOG_EVERY\n",
-    "            lr = scheduler.get_last_lr()[0]\n",
-    "            elapsed = time.time() - train_start\n",
-    "            steps_per_sec = global_step / elapsed\n",
-    "            eta = (total_steps - global_step) / max(steps_per_sec, 1e-8)\n",
-    "            if device == 'cuda':\n",
-    "                vram = torch.cuda.max_memory_allocated() / 1e9\n",
-    "                print(f\"Step {global_step:6d}/{total_steps} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | {steps_per_sec:.1f} it/s | ETA: {eta/60:.0f}m | VRAM: {vram:.1f}GB\")\n",
-    "            else:\n",
-    "                print(f\"Step {global_step:6d}/{total_steps} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | {steps_per_sec:.1f} it/s | ETA: {eta/60:.0f}m\")\n",
-    "        \n",
-    "        # Generate samples\n",
-    "        if global_step % SAMPLE_EVERY == 0:\n",
-    "            print(f\"\\n  📸 Generating samples at step {global_step}...\")\n",
-    "            samples = generate_samples(global_step)\n",
-    "            \n",
-    "            # Display in notebook\n",
-    "            fig, axes = plt.subplots(1, min(8, NUM_SAMPLE_IMAGES), figsize=(16, 2.5))\n",
-    "            if NUM_SAMPLE_IMAGES == 1: axes = [axes]\n",
-    "            for i, ax in enumerate(axes):\n",
-    "                if i < samples.shape[0]:\n",
-    "                    img = (samples[i].cpu().permute(1,2,0) * 0.5 + 0.5).clamp(0,1)\n",
-    "                    ax.imshow(img); ax.axis('off')\n",
-    "            plt.suptitle(f'Step {global_step} | Loss: {loss_val:.4f}'); plt.tight_layout(); plt.show()\n",
-    "        \n",
-    "        # Save checkpoint\n",
-    "        if global_step % SAVE_EVERY == 0:\n",
-    "            ckpt_path = f'{OUTPUT_DIR}/checkpoints/step_{global_step:06d}.pt'\n",
-    "            torch.save({\n",
-    "                'model': model.state_dict(),\n",
-    "                'ema_model': ema_model.state_dict(),\n",
-    "                'optimizer': optimizer.state_dict(),\n",
-    "                'step': global_step,\n",
-    "                'epoch': epoch,\n",
-    "                'losses': all_losses[-2000:],\n",
-    "                'config': config,\n",
-    "            }, ckpt_path)\n",
-    "            print(f\"  💾 Saved checkpoint: {ckpt_path}\")\n",
-    "    \n",
-    "    # Epoch summary\n",
-    "    avg_epoch_loss = epoch_loss / len(dataloader)\n",
-    "    epoch_losses.append(avg_epoch_loss)\n",
-    "    print(f\"\\n  Epoch {epoch+1}/{NUM_EPOCHS} complete | Avg loss: {avg_epoch_loss:.4f}\")\n",
-    "\n",
-    "# Final save\n",
-    "final_path = f'{OUTPUT_DIR}/checkpoints/final.pt'\n",
-    "torch.save({\n",
-    "    'model': model.state_dict(),\n",
-    "    'ema_model': ema_model.state_dict(),\n",
-    "    'step': global_step,\n",
-    "    'config': config,\n",
-    "    'losses': all_losses[-2000:],\n",
-    "}, final_path)\n",
-    "print(f\"\\n✅ Training complete! Final checkpoint: {final_path}\")\n",
-    "print(f\"Total time: {(time.time()-train_start)/3600:.1f} hours\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 📈 Training Curves"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "\n",
-    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
-    "\n",
-    "# Raw loss\n",
-    "ax1.plot(all_losses, alpha=0.3, color='blue', linewidth=0.5)\n",
-    "# Smoothed loss\n",
-    "window = min(200, len(all_losses)//5)\n",
-    "if window > 1:\n",
-    "    smoothed = np.convolve(all_losses, np.ones(window)/window, mode='valid')\n",
-    "    ax1.plot(range(window-1, len(all_losses)), smoothed, color='red', linewidth=2, label=f'Smoothed (w={window})')\n",
-    "ax1.set_xlabel('Step'); ax1.set_ylabel('Loss'); ax1.set_title('Training Loss')\n",
-    "ax1.legend(); ax1.grid(True, alpha=0.3)\n",
-    "\n",
-    "# Epoch loss\n",
-    "if epoch_losses:\n",
-    "    ax2.plot(range(1, len(epoch_losses)+1), epoch_losses, 'o-', color='green')\n",
-    "    ax2.set_xlabel('Epoch'); ax2.set_ylabel('Avg Loss'); ax2.set_title('Loss per Epoch')\n",
-    "    ax2.grid(True, alpha=0.3)\n",
-    "\n",
-    "plt.tight_layout(); plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 🎨 Generate Images"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Generate a batch of images\n",
-    "NUM_GENERATE = 16  # @param {type:\"integer\"}\n",
-    "EULER_STEPS = 50   # @param {type:\"integer\"}\n",
-    "\n",
-    "print(f\"Generating {NUM_GENERATE} images with {EULER_STEPS} Euler steps...\")\n",
-    "ema_model.eval()\n",
-    "\n",
-    "with torch.no_grad():\n",
-    "    z = torch.randn(NUM_GENERATE, 3, IMAGE_SIZE, IMAGE_SIZE, device=device)\n",
-    "    dt = 1.0 / EULER_STEPS\n",
-    "    for i in range(EULER_STEPS, 0, -1):\n",
-    "        t = torch.full((NUM_GENERATE,), i / EULER_STEPS, device=device)\n",
-    "        with torch.amp.autocast(device, dtype=amp_dtype, enabled=USE_AMP and device=='cuda'):\n",
-    "            v = ema_model(z, t)\n",
-    "        if USE_AMP and amp_dtype == torch.float16: v = v.float()\n",
-    "        z = z - v * dt\n",
-    "    generated = z.clamp(-1, 1)\n",
-    "\n",
-    "# Display\n",
-    "nrow = int(math.ceil(math.sqrt(NUM_GENERATE)))\n",
-    "fig, axes = plt.subplots(nrow, nrow, figsize=(2.5*nrow, 2.5*nrow))\n",
-    "axes = axes.flatten() if hasattr(axes, 'flatten') else [axes]\n",
-    "for i, ax in enumerate(axes):\n",
-    "    if i < NUM_GENERATE:\n",
-    "        img = (generated[i].cpu().permute(1,2,0) * 0.5 + 0.5).clamp(0,1)\n",
-    "        ax.imshow(img)\n",
-    "    ax.axis('off')\n",
-    "plt.suptitle(f'LiquidDiffusion Samples ({IMAGE_SIZE}px, {EULER_STEPS} steps)', fontsize=14)\n",
-    "plt.tight_layout(); plt.show()\n",
-    "\n",
-    "# Save\n",
-    "grid = make_grid(generated * 0.5 + 0.5, nrow=nrow, padding=2)\n",
-    "save_image(grid, f'{OUTPUT_DIR}/final_samples.png')\n",
-    "print(f\"Saved to {OUTPUT_DIR}/final_samples.png\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 💾 Save / Load Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Save to HuggingFace Hub (optional)\n",
-    "PUSH_TO_HUB = False  # @param {type:\"boolean\"}\n",
-    "HUB_MODEL_ID = 'your-username/liquid-diffusion-celebahq-256'  # @param {type:\"string\"}\n",
-    "\n",
-    "if PUSH_TO_HUB:\n",
-    "    from huggingface_hub import HfApi\n",
-    "    api = HfApi()\n",
-    "    api.create_repo(HUB_MODEL_ID, exist_ok=True)\n",
-    "    api.upload_file(\n",
-    "        path_or_fileobj=final_path,\n",
-    "        path_in_repo='model.pt',\n",
-    "        repo_id=HUB_MODEL_ID,\n",
-    "    )\n",
-    "    print(f\"Pushed to https://huggingface.co/{HUB_MODEL_ID}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "---\n",
-    "\n",
-    "## 📖 Architecture Deep Dive\n",
-    "\n",
-    "### What makes LiquidDiffusion special?\n",
-    "\n",
-    "**1. CfC Time-Gating (the \"liquid\" part)**\n",
-    "```\n",
-    "gate = σ(time_a(t_emb) · f(features) - time_b(t_emb))\n",
-    "output = gate · g(features) + (1 - gate) · h(features)\n",
-    "```\n",
-    "- `f` = time-constant head (controls gate sensitivity)\n",
-    "- `g` = \"from\" state (what features look like at short time)\n",
-    "- `h` = \"to\" state (attractor for long time)\n",
-    "- The gate adapts **per-channel, per-spatial-position** based on both the input features AND the noise level\n",
-    "\n",
-    "**2. Liquid Relaxation Residual**\n",
-    "```\n",
-    "α = exp(-λ · |t_emb_mean|)\n",
-    "out = α · input + (1-α) · CfC_output\n",
-    "```\n",
-    "- When noise is high (large t): α→0, rely on CfC output (needs heavy processing)\n",
-    "- When noise is low (small t): α→1, preserve input (just refine details)\n",
-    "- λ is learned per-channel — each feature dimension decides its own decay rate\n",
-    "\n",
-    "**3. Multi-Scale Spatial Mixing**\n",
-    "- 3×3 + 5×5 + 7×7 depthwise convolutions + global average pooling\n",
-    "- Gives effective global receptive field without O(n²) attention\n",
-    "- All parallel, all efficient\n",
-    "\n",
-    "### Why no attention?\n",
-    "- Self-attention is O(n²) in spatial tokens — at 256px that's 65K tokens\n",
-    "- Depthwise convolutions + global pooling give global context at O(n) cost\n",
-    "- The CfC time-gating provides the \"adaptive routing\" that attention normally gives\n",
-    "- Result: **same expressivity, 10× less memory, 3× faster**\n",
-    "\n",
-    "### Parameter counts\n",
-    "| Config | Params | 256px VRAM | 512px VRAM |\n",
-    "|--------|--------|------------|------------|\n",
-    "| tiny   | ~23M   | ~6 GB      | ~12 GB     |\n",
-    "| small  | ~69M   | ~10 GB     | ~20 GB     |\n",
-    "| base   | ~154M  | ~16 GB     | ~30 GB     |"
-   ]
-  }
- ]
+  "nbformat": 4,
+  "nbformat_minor": 0
 }