V2: Add VAE, fix datasets, streaming, precaching

LiquidDiffusion_Training.ipynb

@@ -1,678 +1,33 @@
 {
-  "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"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {"provenance": [], "gpuType": "T4"},
+  "kernelspec": {"name": "python3", "display_name": "Python 3"},
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {"cell_type": "markdown", "metadata": {}, "source": ["# \ud83c\udf0a LiquidDiffusion: Attention-Free Image Generation with Liquid Neural Networks\n", "\n", "A **novel** image generation model combining:\n", "- **Liquid Neural Networks** (CfC \u2014 Closed-form Continuous-depth) for adaptive processing\n", "- **Rectified Flow** for simple, stable training (MSE velocity prediction)\n", "- **Pretrained SD-VAE** for efficient latent-space training (4ch, 8\u00d7 downscale)\n", "- **Zero attention** \u2014 fully convolutional + multi-scale spatial mixing\n", "- **Fully parallelizable** \u2014 no ODE loops, no recurrence\n", "\n", "### Key Innovation\n", "Diffusion timestep = liquid time constant. CfC gate `\u03c3(-f\u00b7t)` adapts behavior to noise level.\n", "\n", "### References\n", "- [CfC Networks (Nature MI 2022)](https://arxiv.org/abs/2106.13898)\n", "- [LiquidTAD (2024)](https://arxiv.org/abs/2604.18274) | [USM (CVPR 2025)](https://arxiv.org/abs/2504.13499)\n", "- [Rectified Flow (ICLR 2023)](https://arxiv.org/abs/2209.03003)\n", "- **Repo**: [krystv/liquid-diffusion](https://huggingface.co/krystv/liquid-diffusion)"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \u2699\ufe0f Configuration"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["#@title \u2699\ufe0f Configuration { display-mode: \"form\" }\n", "\n", "#@markdown ### Model\n", "MODEL_SIZE = 'tiny' #@param ['tiny', 'small', 'custom']\n", "#@markdown > `tiny`=23M (256px, ~6GB) | `small`=69M (256px, ~10GB)\n", "CUSTOM_CHANNELS = [48, 96, 192]\n", "CUSTOM_BLOCKS = [1, 2, 3]\n", "CUSTOM_T_DIM = 192\n", "\n", "#@markdown ### Resolution\n", "IMAGE_SIZE = 256 #@param [128, 256, 512] {type:\"raw\"}\n", "\n", "#@markdown ### VAE (Latent Space)\n", "USE_VAE = True #@param {type:\"boolean\"}\n", "#@markdown > Pretrained SD-VAE encodes images to 4ch latents (8\u00d7 smaller). **Highly recommended.**\n", "VAE_MODEL = 'stabilityai/sd-vae-ft-mse' #@param ['stabilityai/sd-vae-ft-mse', 'madebyollin/sdxl-vae-fp16-fix']\n", "PRECACHE_LATENTS = True #@param {type:\"boolean\"}\n", "#@markdown > Pre-encode all images once. Frees ~160MB VAE VRAM during training.\n", "\n", "#@markdown ### Dataset\n", "DATASET = 'nielsr/CelebA-faces' #@param ['nielsr/CelebA-faces', 'huggan/flowers-102-categories', 'reach-vb/pokemon-blip-captions', 'huggan/anime-faces', 'huggan/AFHQv2', 'Norod78/cartoon-blip-captions']\n", "#@markdown > All verified \u2713 | CelebA=202K faces | flowers=8K | pokemon=833 | anime=21K | AFHQ=16K animals | cartoon=2K\n", "IMAGE_COLUMN = 'image'\n", "MAX_SAMPLES = None  # e.g. 5000 for quick test, None=full\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", "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\n", "TIME_SAMPLING = 'logit_normal' #@param ['uniform', 'logit_normal']\n", "USE_AMP = True #@param {type:\"boolean\"}\n", "AMP_DTYPE = 'float16' #@param ['float16', 'bfloat16']\n", "\n", "#@markdown ### Sampling & Logging\n", "SAMPLE_EVERY = 500 #@param {type:\"integer\"}\n", "NUM_SAMPLE_IMAGES = 8 #@param {type:\"integer\"}\n", "NUM_EULER_STEPS = 50 #@param {type:\"integer\"}\n", "SAVE_EVERY = 2000 #@param {type:\"integer\"}\n", "OUTPUT_DIR = './outputs'\n", "RESUME_FROM = None\n", "LOG_EVERY = 50\n", "\n", "LATENT_SIZE = IMAGE_SIZE // 8 if USE_VAE else IMAGE_SIZE\n", "IN_CHANNELS = 4 if USE_VAE else 3\n", "print(f\"Config: {MODEL_SIZE} | {IMAGE_SIZE}px {'(latent '+str(LATENT_SIZE)+'px)' if USE_VAE else '(pixel)'} | {DATASET}\")\n", "print(f\"Training: bs={BATCH_SIZE}, lr={LEARNING_RATE}, epochs={NUM_EPOCHS}, AMP={USE_AMP}\")"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\udce6 Install & Check GPU"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["!pip install -q datasets diffusers accelerate huggingface_hub Pillow matplotlib transformers\n", "import torch\n", "print(f\"PyTorch: {torch.__version__}, CUDA: {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\")\n", "else:\n", "    print(\"\u26a0\ufe0f No GPU! Enable via Runtime \u2192 Change runtime type.\")"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83c\udfd7\ufe0f Model Architecture"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import math, copy, os, time\nfrom glob import glob\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\nfrom torchvision.utils import save_image, make_grid\n\nclass SinusoidalTimeEmbedding(nn.Module):\n    def __init__(self, dim, max_period=10000):\n        super().__init__()\n        self.dim, self.mp = dim, max_period\n        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.SiLU(), nn.Linear(dim*4, dim))\n    def forward(self, t):\n        h = self.dim // 2\n        f = torch.exp(-math.log(self.mp)*torch.arange(h, device=t.device, dtype=t.dtype)/h)\n        e = torch.cat([torch.cos(t[:,None]*f[None]), torch.sin(t[:,None]*f[None])], -1)\n        if self.dim%2: e = F.pad(e,(0,1))\n        return self.mlp(e)\n\nclass AdaLN(nn.Module):\n    def __init__(self, dim, cd):\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(cd, dim*2))\n    def forward(self, x, te):\n        s, sh = self.proj(te).chunk(2,1)\n        return self.norm(x)*(1+s[:,:,None,None])+sh[:,:,None,None]\n\nclass ParallelCfCBlock(nn.Module):\n    def __init__(self, dim, td, er=2.0, ks=7, dr=0.0):\n        super().__init__()\n        hid = int(dim*er)\n        self.bdw = nn.Conv2d(dim, dim, ks, padding=ks//2, groups=dim)\n        self.bpw = nn.Conv2d(dim, hid, 1)\n        self.ba = nn.SiLU()\n        self.fh = nn.Conv2d(hid, dim, 1)\n        self.gh = nn.Sequential(nn.Conv2d(hid,hid,ks,padding=ks//2,groups=hid),nn.SiLU(),nn.Conv2d(hid,dim,1))\n        self.hh = nn.Sequential(nn.Conv2d(hid,hid,ks,padding=ks//2,groups=hid),nn.SiLU(),nn.Conv2d(hid,dim,1))\n        self.ta, self.tb = nn.Linear(td, dim), nn.Linear(td, dim)\n        self.rho = nn.Parameter(torch.zeros(1,dim,1,1))\n        self.og = nn.Sequential(nn.SiLU(), nn.Linear(td, dim))\n        self.do = nn.Dropout(dr) if dr>0 else nn.Identity()\n    def forward(self, x, te):\n        res = x\n        bb = self.ba(self.bpw(self.bdw(x)))\n        f,g,h = self.fh(bb), self.gh(bb), self.hh(bb)\n        gt = torch.sigmoid(self.ta(te)[:,:,None,None]*f - self.tb(te)[:,:,None,None])\n        co = self.do(gt*g + (1-gt)*h)\n        lam = F.softplus(self.rho)+1e-6\n        al = torch.exp(-lam*te.mean(1,keepdim=True)[:,:,None,None].abs().clamp(min=0.01))\n        return (al*res+(1-al)*co)*torch.sigmoid(self.og(te))[:,:,None,None]\n\nclass MultiScaleSpatialMix(nn.Module):\n    def __init__(self, dim, td):\n        super().__init__()\n        self.d3=nn.Conv2d(dim,dim,3,padding=1,groups=dim)\n        self.d5=nn.Conv2d(dim,dim,5,padding=2,groups=dim)\n        self.d7=nn.Conv2d(dim,dim,7,padding=3,groups=dim)\n        self.gp=nn.AdaptiveAvgPool2d(1); self.gpj=nn.Conv2d(dim,dim,1)\n        self.mg=nn.Conv2d(dim*4,dim,1); self.ac=nn.SiLU(); self.an=AdaLN(dim,td)\n    def forward(self, x, te):\n        xn=self.an(x,te)\n        return x+self.ac(self.mg(torch.cat([self.d3(xn),self.d5(xn),self.d7(xn),self.gpj(self.gp(xn)).expand_as(xn)],1)))\n\nclass LiquidDiffusionBlock(nn.Module):\n    def __init__(self, dim, td, er=2.0, ks=7, dr=0.0):\n        super().__init__()\n        self.a1=AdaLN(dim,td); self.cfc=ParallelCfCBlock(dim,td,er,ks,dr)\n        self.sm=MultiScaleSpatialMix(dim,td); self.a2=AdaLN(dim,td)\n        ff=int(dim*er); self.ff=nn.Sequential(nn.Conv2d(dim,ff,1),nn.SiLU(),nn.Conv2d(ff,dim,1))\n        self.rs=nn.Parameter(torch.ones(1)*0.1)\n    def forward(self, x, te):\n        x=x+self.rs*self.cfc(self.a1(x,te),te); x=self.sm(x,te)\n        return x+self.rs*self.ff(self.a2(x,te))\n\nclass DS(nn.Module):\n    def __init__(self,i,o): super().__init__(); self.c=nn.Conv2d(i,o,3,stride=2,padding=1)\n    def forward(self,x): return self.c(x)\nclass US(nn.Module):\n    def __init__(self,i,o): super().__init__(); self.c=nn.Conv2d(i,o,3,padding=1)\n    def forward(self,x): return self.c(F.interpolate(x,scale_factor=2,mode='nearest'))\nclass SF(nn.Module):\n    def __init__(self,d,td): super().__init__(); self.p=nn.Conv2d(d*2,d,1); self.g=nn.Sequential(nn.SiLU(),nn.Linear(td,d),nn.Sigmoid())\n    def forward(self,x,sk,te): m=self.p(torch.cat([x,sk],1)); g=self.g(te)[:,:,None,None]; return m*g+x*(1-g)\n\nclass LiquidDiffusionUNet(nn.Module):\n    def __init__(self, in_ch=3, chs=None, bps=None, td=256, er=2.0, ks=7, dr=0.0):\n        super().__init__()\n        chs=chs or [64,128,256]; bps=bps or [2,2,4]\n        assert len(chs)==len(bps)\n        self.chs,self.ns=chs,len(chs)\n        self.te=SinusoidalTimeEmbedding(td)\n        self.st=nn.Sequential(nn.Conv2d(in_ch,chs[0],3,padding=1),nn.SiLU(),nn.Conv2d(chs[0],chs[0],3,padding=1))\n        self.enc,self.dn=nn.ModuleList(),nn.ModuleList()\n        for i in range(self.ns):\n            self.enc.append(nn.ModuleList([LiquidDiffusionBlock(chs[i],td,er,ks,dr) for _ in range(bps[i])]))\n            if i<self.ns-1: self.dn.append(DS(chs[i],chs[i+1]))\n        self.bot=nn.ModuleList([LiquidDiffusionBlock(chs[-1],td,er,ks,dr) for _ in range(2)])\n        self.dec,self.up_,self.sf_=nn.ModuleList(),nn.ModuleList(),nn.ModuleList()\n        for i in range(self.ns-1,-1,-1):\n            if i<self.ns-1: self.up_.append(US(chs[i+1],chs[i])); self.sf_.append(SF(chs[i],td))\n            self.dec.append(nn.ModuleList([LiquidDiffusionBlock(chs[i],td,er,ks,dr) for _ in range(bps[i])]))\n        hg=min(32,chs[0])\n        while chs[0]%hg!=0: hg-=1\n        self.hd=nn.Sequential(nn.GroupNorm(hg,chs[0]),nn.SiLU(),nn.Conv2d(chs[0],in_ch,3,padding=1))\n        nn.init.zeros_(self.hd[-1].weight); nn.init.zeros_(self.hd[-1].bias)\n    def forward(self, x, t):\n        te=self.te(t); h=self.st(x); sk=[]\n        for i in range(self.ns):\n            for b in self.enc[i]: h=b(h,te)\n            sk.append(h)\n            if i<self.ns-1: h=self.dn[i](h)\n        for b in self.bot: h=b(h,te)\n        ui=0\n        for di in range(self.ns):\n            si=self.ns-1-di\n            if di>0: h=self.up_[ui](h); h=self.sf_[ui](h,sk[si],te); ui+=1\n            for b in self.dec[di]: h=b(h,te)\n        return self.hd(h)\n    def count_params(self): return sum(p.numel() for p in self.parameters()), sum(p.numel() for p in self.parameters() if p.requires_grad)\n\nprint('\u2705 Model architecture defined.')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\udd27 Build Model + Load VAE"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["device = 'cuda' if torch.cuda.is_available() else 'cpu'\nCFGS = {'tiny': dict(chs=[64,128,256], bps=[2,2,4], td=256), 'small': dict(chs=[96,192,384], bps=[2,3,6], td=384)}\nif MODEL_SIZE=='custom': cfg=dict(chs=CUSTOM_CHANNELS,bps=CUSTOM_BLOCKS,td=CUSTOM_T_DIM)\nelse: cfg=CFGS[MODEL_SIZE]\nmodel = LiquidDiffusionUNet(in_ch=IN_CHANNELS, **cfg).to(device)\ntp,_=model.count_params()\nprint(f'Model: {MODEL_SIZE} | {tp:,} params ({tp/1e6:.1f}M) | in_ch={IN_CHANNELS}')\n\nvae=None; vae_scale=1.0\nif USE_VAE:\n    from diffusers import AutoencoderKL\n    print(f'Loading VAE: {VAE_MODEL}...')\n    vae = AutoencoderKL.from_pretrained(VAE_MODEL, torch_dtype=torch.float16 if device=='cuda' else torch.float32)\n    vae = vae.to(device).eval(); vae.requires_grad_(False)\n    vae_scale = vae.config.scaling_factor\n    print(f'VAE: {sum(p.numel() for p in vae.parameters())/1e6:.1f}M params, latent_ch={vae.config.latent_channels}, scale={vae_scale}')\n    print(f'  {IMAGE_SIZE}px \u2192 {LATENT_SIZE}px latent (8\u00d7 downsample)')\n\nwith torch.no_grad():\n    tx=torch.randn(1,IN_CHANNELS,LATENT_SIZE,LATENT_SIZE,device=device)\n    assert model(tx, torch.tensor([0.5],device=device)).shape==tx.shape\n    print(f'Forward OK: {tx.shape}')\n    del tx\nif device=='cuda': torch.cuda.empty_cache(); print(f'VRAM: {torch.cuda.memory_allocated()/1e9:.2f}GB')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\udcca Load Dataset"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["from PIL import Image\nfrom datasets import load_dataset\nimport matplotlib.pyplot as plt\n\nclass ImageDS(Dataset):\n    def __init__(self, ds, sz, col='image'):\n        self.ds, self.col = ds, col\n        self.tf = transforms.Compose([transforms.Resize(sz, interpolation=transforms.InterpolationMode.LANCZOS),\n            transforms.CenterCrop(sz), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5],[0.5])])\n    def __len__(self): return len(self.ds)\n    def __getitem__(self, i):\n        img = self.ds[i][self.col]\n        if not hasattr(img,'convert'): img=Image.fromarray(img)\n        return self.tf(img.convert('RGB'))\n\nprint(f'Loading: {DATASET}')\nraw = load_dataset(DATASET, split='train')\nif MAX_SAMPLES: raw = raw.select(range(min(MAX_SAMPLES, len(raw))))\nprint(f'  {len(raw):,} images')\ndataset = ImageDS(raw, IMAGE_SIZE, IMAGE_COLUMN)\ndata_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True, drop_last=True, persistent_workers=True)\nprint(f'  {len(data_loader):,} steps/epoch | ~{len(data_loader)*NUM_EPOCHS:,} total steps')\n\nsb=next(iter(data_loader))\nfig,axes=plt.subplots(1,min(8,BATCH_SIZE),figsize=(16,2.5))\nfor i,ax in enumerate(axes if hasattr(axes,'__len__') else [axes]): ax.imshow((sb[i].permute(1,2,0)*0.5+0.5).clamp(0,1)); ax.axis('off')\nplt.suptitle(f'Training samples ({IMAGE_SIZE}px)'); plt.tight_layout(); plt.show()"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\uddc3\ufe0f Pre-cache Latents (if VAE enabled)"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["cached_latents = None\ntrain_loader = data_loader\n\nif USE_VAE and PRECACHE_LATENTS:\n    print(f'Pre-encoding {len(dataset):,} images...')\n    cl = DataLoader(dataset, batch_size=BATCH_SIZE*2, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n    all_z = []\n    vd = torch.float16 if device=='cuda' else torch.float32\n    t0 = time.time()\n    with torch.no_grad():\n        for bi, imgs in enumerate(cl):\n            z = vae.encode(imgs.to(device, dtype=vd)).latent_dist.sample() * vae_scale\n            all_z.append(z.cpu().float())\n            if (bi+1)%50==0: print(f'  {(bi+1)*BATCH_SIZE*2:,}/{len(dataset):,}')\n    cached_latents = torch.cat(all_z)\n    print(f'  Done in {time.time()-t0:.0f}s | Shape: {cached_latents.shape} | {cached_latents.numel()*4/1e9:.2f}GB')\n    vae = vae.cpu()\n    if device=='cuda': torch.cuda.empty_cache(); print(f'  VAE \u2192 CPU. GPU VRAM: {torch.cuda.memory_allocated()/1e9:.2f}GB')\n    class LatDS(Dataset):\n        def __init__(self,z): self.z=z\n        def __len__(self): return len(self.z)\n        def __getitem__(self,i): return self.z[i]\n    train_loader = DataLoader(LatDS(cached_latents), batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True, drop_last=True)\n    print(f'  Latent loader: {len(train_loader)} steps/epoch')\nelif USE_VAE:\n    print('Online VAE encoding (VAE stays on GPU)')\nelse:\n    print('Pixel-space training (no VAE)')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\ude80 Training"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import matplotlib.pyplot as plt\nos.makedirs(f'{OUTPUT_DIR}/samples', exist_ok=True); os.makedirs(f'{OUTPUT_DIR}/checkpoints', exist_ok=True)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY, betas=(0.9,0.999))\ntotal_steps = len(train_loader)*NUM_EPOCHS\nwarmup = min(1000, total_steps//10)\ndef lrl(step):\n    if step<warmup: return step/max(1,warmup)\n    return max(0.0, 0.5*(1+math.cos(math.pi*(step-warmup)/max(1,total_steps-warmup))))\nsched = torch.optim.lr_scheduler.LambdaLR(optimizer, lrl)\n\nema = copy.deepcopy(model).eval()\nfor p in ema.parameters(): p.requires_grad_(False)\nscaler = torch.amp.GradScaler('cuda', enabled=(USE_AMP and device=='cuda'))\namp_dt = getattr(torch, AMP_DTYPE) if USE_AMP and device=='cuda' else torch.float32\n\ndef st(bs):\n    e=1e-5\n    if TIME_SAMPLING=='uniform': return torch.rand(bs,device=device)*(1-2*e)+e\n    return torch.sigmoid(torch.randn(bs,device=device)).clamp(e,1-e)\n\ngstep,start_ep,all_losses,ep_losses=0,0,[],[]\nif RESUME_FROM and os.path.exists(RESUME_FROM):\n    ck=torch.load(RESUME_FROM,map_location=device,weights_only=False)\n    model.load_state_dict(ck['model']); ema.load_state_dict(ck['ema_model']); optimizer.load_state_dict(ck['optimizer'])\n    gstep=ck.get('step',0); start_ep=ck.get('epoch',0); all_losses=ck.get('losses',[])\n    print(f'Resumed from step {gstep}')\n\n@torch.no_grad()\ndef gen_samples(step):\n    ema.eval()\n    z=torch.randn(NUM_SAMPLE_IMAGES,IN_CHANNELS,LATENT_SIZE,LATENT_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_dt,enabled=USE_AMP and device=='cuda'): v=ema(z,t)\n        if USE_AMP and amp_dt==torch.float16: v=v.float()\n        z=z-v*dt\n    z=z.clamp(-3,3)\n    if USE_VAE:\n        _v=vae.to(device); vd=torch.float16 if device=='cuda' else torch.float32\n        imgs=_v.decode(z.to(vd)/vae_scale).sample.float()\n        if PRECACHE_LATENTS: vae.cpu()\n    else: imgs=z\n    imgs=imgs.clamp(-1,1)\n    save_image(make_grid(imgs*0.5+0.5,nrow=int(math.ceil(math.sqrt(NUM_SAMPLE_IMAGES))),padding=2),f'{OUTPUT_DIR}/samples/step_{step:06d}.png')\n    return imgs\n\nprint(f'\\n{\"=\"*60}\\nTraining: {NUM_EPOCHS} epochs, {total_steps:,} steps\\n{\"=\"*60}\\n')\nt_start=time.time(); online_vae=USE_VAE and not PRECACHE_LATENTS; vd=torch.float16 if device=='cuda' else torch.float32\n\nfor epoch in range(start_ep, NUM_EPOCHS):\n    model.train(); el=0\n    for batch in train_loader:\n        if online_vae:\n            with torch.no_grad(): x0=vae.encode(batch.to(device,dtype=vd)).latent_dist.sample()*vae_scale; x0=x0.float()\n        else: x0=batch.to(device)\n        x1=torch.randn_like(x0); t=st(x0.shape[0]); te=t[:,None,None,None]\n        xt=(1-te)*x0+te*x1; vt=x1-x0\n        with torch.amp.autocast(device,dtype=amp_dt,enabled=USE_AMP and device=='cuda'):\n            vp=model(xt,t); loss=F.mse_loss(vp,vt)\n        optimizer.zero_grad(set_to_none=True); scaler.scale(loss).backward()\n        if GRAD_CLIP>0: scaler.unscale_(optimizer); torch.nn.utils.clip_grad_norm_(model.parameters(),GRAD_CLIP)\n        scaler.step(optimizer); scaler.update(); sched.step()\n        with torch.no_grad():\n            for ep,mp in zip(ema.parameters(),model.parameters()): ep.data.mul_(EMA_DECAY).add_(mp.data,alpha=1-EMA_DECAY)\n        gstep+=1; lv=loss.item(); all_losses.append(lv); el+=lv\n        if gstep%LOG_EVERY==0:\n            avg=sum(all_losses[-LOG_EVERY:])/LOG_EVERY; lr=sched.get_last_lr()[0]\n            sps=gstep/(time.time()-t_start); eta=(total_steps-gstep)/max(sps,1e-8)\n            vm=f' | VRAM:{torch.cuda.max_memory_allocated()/1e9:.1f}GB' if device=='cuda' else ''\n            print(f'Step {gstep:6d}/{total_steps} | Loss:{avg:.4f} | LR:{lr:.2e} | {sps:.1f}it/s | ETA:{eta/60:.0f}m{vm}')\n        if gstep%SAMPLE_EVERY==0:\n            print('  \\U0001f4f8 Generating...'); samps=gen_samples(gstep)\n            fig,axes=plt.subplots(1,min(8,NUM_SAMPLE_IMAGES),figsize=(16,2.5))\n            if not hasattr(axes,'__len__'): axes=[axes]\n            for i,ax in enumerate(axes):\n                if i<samps.shape[0]: ax.imshow((samps[i].cpu().permute(1,2,0)*0.5+0.5).clamp(0,1))\n                ax.axis('off')\n            plt.suptitle(f'Step {gstep} | Loss:{lv:.4f}'); plt.tight_layout(); plt.show()\n        if gstep%SAVE_EVERY==0:\n            cp=f'{OUTPUT_DIR}/checkpoints/step_{gstep:06d}.pt'\n            torch.save({'model':model.state_dict(),'ema_model':ema.state_dict(),'optimizer':optimizer.state_dict(),'step':gstep,'epoch':epoch,'losses':all_losses[-2000:],'config':cfg},cp)\n            print(f'  \\U0001f4be Saved: {cp}')\n    ep_losses.append(el/len(train_loader))\n    print(f'  Epoch {epoch+1}/{NUM_EPOCHS} | Avg loss:{ep_losses[-1]:.4f}')\n\nfp=f'{OUTPUT_DIR}/checkpoints/final.pt'\ntorch.save({'model':model.state_dict(),'ema_model':ema.state_dict(),'step':gstep,'config':cfg,'losses':all_losses[-2000:]},fp)\nprint(f'\\n\\u2705 Done! {fp} | {(time.time()-t_start)/3600:.1f}h')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\udcc8 Training Curves"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import numpy as np\nfig,(a1,a2)=plt.subplots(1,2,figsize=(14,5))\na1.plot(all_losses,alpha=0.3,color='blue',lw=0.5)\nw=min(200,len(all_losses)//5)\nif w>1:\n    sm=np.convolve(all_losses,np.ones(w)/w,mode='valid')\n    a1.plot(range(w-1,len(all_losses)),sm,color='red',lw=2,label=f'Smooth(w={w})')\na1.set_xlabel('Step');a1.set_ylabel('Loss');a1.set_title('Training Loss');a1.legend();a1.grid(True,alpha=0.3)\nif ep_losses: a2.plot(range(1,len(ep_losses)+1),ep_losses,'o-',color='green'); a2.set_xlabel('Epoch');a2.set_ylabel('Loss');a2.set_title('Per Epoch');a2.grid(True,alpha=0.3)\nplt.tight_layout();plt.show()"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83c\udfa8 Generate Images"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["N_GEN=16; STEPS=50\nprint(f'Generating {N_GEN} images ({STEPS} steps)...')\nema.eval()\nif USE_VAE: vae=vae.to(device)\nwith torch.no_grad():\n    z=torch.randn(N_GEN,IN_CHANNELS,LATENT_SIZE,LATENT_SIZE,device=device)\n    dt=1.0/STEPS\n    for i in range(STEPS,0,-1):\n        t=torch.full((N_GEN,),i/STEPS,device=device)\n        with torch.amp.autocast(device,dtype=amp_dt,enabled=USE_AMP and device=='cuda'): v=ema(z,t)\n        if USE_AMP and amp_dt==torch.float16: v=v.float()\n        z=z-v*dt\n    if USE_VAE: vdd=torch.float16 if device=='cuda' else torch.float32; gen=vae.decode(z.clamp(-3,3).to(vdd)/vae_scale).sample.float()\n    else: gen=z\n    gen=gen.clamp(-1,1)\nnr=int(math.ceil(math.sqrt(N_GEN)))\nfig,axes=plt.subplots(nr,nr,figsize=(2.5*nr,2.5*nr))\naxes=axes.flatten() if hasattr(axes,'flatten') else [axes]\nfor i,ax in enumerate(axes):\n    if i<N_GEN: ax.imshow((gen[i].cpu().permute(1,2,0)*0.5+0.5).clamp(0,1))\n    ax.axis('off')\nplt.suptitle(f'LiquidDiffusion ({IMAGE_SIZE}px, {STEPS} steps)',fontsize=14);plt.tight_layout();plt.show()\nsave_image(make_grid(gen*0.5+0.5,nrow=nr,padding=2),f'{OUTPUT_DIR}/final_samples.png')\nprint(f'Saved: {OUTPUT_DIR}/final_samples.png')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["## \ud83d\udcbe Push to Hub"]},
+  {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["PUSH=False #@param {type:\"boolean\"}\nHUB_ID='your-username/liquid-diffusion-256' #@param {type:\"string\"}\nif PUSH:\n    from huggingface_hub import HfApi\n    api=HfApi(); api.create_repo(HUB_ID,exist_ok=True)\n    api.upload_file(path_or_fileobj=fp,path_in_repo='model.pt',repo_id=HUB_ID)\n    print(f'Pushed: https://huggingface.co/{HUB_ID}')"]},
+  {"cell_type": "markdown", "metadata": {}, "source": ["---\n", "## \ud83d\udcd6 Architecture Reference\n", "\n", "### CfC Time-Gating\n", "```\n", "gate = \u03c3(time_a(t) \u00b7 f(features) - time_b(t))\n", "out = gate \u00b7 g + (1-gate) \u00b7 h\n", "```\n", "### Liquid Relaxation\n", "```\n", "\u03b1 = exp(-\u03bb\u00b7|t|), out = \u03b1\u00b7input + (1-\u03b1)\u00b7CfC_out\n", "```\n", "High noise \u2192 \u03b1\u22480 \u2192 heavy processing. Low noise \u2192 \u03b1\u22481 \u2192 preserve.\n", "\n", "### VAE: `stabilityai/sd-vae-ft-mse`\n", "83M params, 4ch latents, 8\u00d7 downscale. 256px\u219232\u00d732\u00d74 latent.\n", "\n", "### Verified Datasets\n", "| Dataset | Size | Content |\n", "|---------|------|---------|\n", "| `nielsr/CelebA-faces` | 202K | Celebrity faces |\n", "| `huggan/flowers-102-categories` | 8K | Flowers |\n", "| `reach-vb/pokemon-blip-captions` | 833 | Pokemon art |\n", "| `huggan/anime-faces` | 21K | Anime faces |\n", "| `huggan/AFHQv2` | 16K | Cat/dog/wild |\n", "| `Norod78/cartoon-blip-captions` | 2K | Cartoon characters |"]}
+ ]
 }