{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LiquidFlow: Liquid Neural Network + Mamba-2 SSD Image Generator\n",
    "\n",
    "**Train on Google Colab Free Tier (T4 GPU) | Export for Mobile Deployment**\n",
    "\n",
    "LiquidFlow combines:\n",
    "- **CfC (Closed-form Continuous-time)** Liquid Neural Networks — adaptive time gates\n",
    "- **Mamba-2 SSD** — linear-time attention replacement, fully parallelizable\n",
    "- **Physics-Informed Regularization** — TV loss, spectral constraints\n",
    "- **TAESD VAE** — Tiny AutoEncoder (< 3M params) for fast encoding\n",
    "\n",
    "---\n",
    "## Quick Start\n",
    "1. Runtime → Change runtime type → GPU (T4)\n",
    "2. Run all cells in order\n",
    "3. Training starts automatically on CIFAR-10\n",
    "4. Check samples in `./outputs/samples/`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 1. Install Dependencies\n",
    "!pip install -q torch torchvision diffusers tqdm pillow numpy accelerate\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\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 2. Clone LiquidFlow Repository\n",
    "import os\n",
    "if not os.path.exists('/content/LiquidFlow'):\n",
    "    !git clone https://huggingface.co/krystv/LiquidFlow-Gen /content/LiquidFlow\n",
    "else:\n",
    "    !cd /content/LiquidFlow && git pull\n",
    "%cd /content/LiquidFlow\n",
    "\n",
    "import sys\n",
    "sys.path.insert(0, '/content/LiquidFlow')\n",
    "\n",
    "from liquid_flow.generator import create_liquidflow\n",
    "from liquid_flow.vae_wrapper import TAESDWrapper\n",
    "print('LiquidFlow imported successfully!')"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 3. Configuration\n",
    "\n",
    "# Model variant: 'tiny' (~3.6M), 'small' (~11M), 'base' (~36M)\n",
    "MODEL_VARIANT = 'small'  # @param ['tiny', 'small', 'base']\n",
    "\n",
    "# Image size (128 recommended for T4 free tier)\n",
    "IMAGE_SIZE = 128  # @param [64, 128, 256, 512] {type:\"integer\"}\n",
    "\n",
    "# Training hyperparameters\n",
    "BATCH_SIZE = 32  # @param [8, 16, 32, 64] {type:\"integer\"}\n",
    "EPOCHS = 50  # @param [10, 25, 50, 100] {type:\"integer\"}\n",
    "LEARNING_RATE = 2e-4  # @param {type:\"number\"}\n",
    "\n",
    "# Dataset\n",
    "DATASET = 'cifar10'  # @param ['cifar10', 'cifar100', 'stl10']\n",
    "\n",
    "# Sampling\n",
    "SAMPLE_EVERY = 5  # @param {type:\"integer\"}\n",
    "SAMPLE_STEPS = 50  # @param {type:\"integer\"}\n",
    "\n",
    "# Ensure integer types (Colab forms can return strings)\n",
    "IMAGE_SIZE = int(IMAGE_SIZE)\n",
    "BATCH_SIZE = int(BATCH_SIZE)\n",
    "EPOCHS = int(EPOCHS)\n",
    "SAMPLE_EVERY = int(SAMPLE_EVERY)\n",
    "SAMPLE_STEPS = int(SAMPLE_STEPS)\n",
    "LEARNING_RATE = float(LEARNING_RATE)\n",
    "\n",
    "print(f\"Config: {MODEL_VARIANT} model, {IMAGE_SIZE}px, batch={BATCH_SIZE}, epochs={EPOCHS}, lr={LEARNING_RATE}\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 4. Load VAE & Create Model\n",
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Load TAESD (Tiny AutoEncoder)\n",
    "print(\"Loading TAESD VAE...\")\n",
    "vae = TAESDWrapper.load(device)\n",
    "latent_size = IMAGE_SIZE // 8\n",
    "print(f\"VAE loaded! Latent: {IMAGE_SIZE}x{IMAGE_SIZE} -> {latent_size}x{latent_size}x4\")\n",
    "\n",
    "# Create LiquidFlow model\n",
    "print(f\"\\nCreating '{MODEL_VARIANT}' LiquidFlow model...\")\n",
    "model = create_liquidflow(variant=MODEL_VARIANT, image_size=IMAGE_SIZE)\n",
    "model = model.to(device)\n",
    "\n",
    "n_params = model.count_parameters()\n",
    "print(f\"Model parameters: {n_params:,} ({n_params/1e6:.1f}M)\")\n",
    "print(f\"\\nReady to train!\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 5. Load Dataset\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n",
    "])\n",
    "\n",
    "if DATASET == 'cifar10':\n",
    "    dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
    "elif DATASET == 'cifar100':\n",
    "    dataset = datasets.CIFAR100(root='./data', train=True, download=True, transform=transform)\n",
    "elif DATASET == 'stl10':\n",
    "    dataset = datasets.STL10(root='./data', split='train', download=True, transform=transform)\n",
    "\n",
    "dataloader = DataLoader(\n",
    "    dataset, batch_size=BATCH_SIZE, shuffle=True,\n",
    "    num_workers=2, pin_memory=True, drop_last=True,\n",
    ")\n",
    "\n",
    "print(f\"Dataset: {DATASET} ({len(dataset):,} images)\")\n",
    "print(f\"Batches per epoch: {len(dataloader)}\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 6. Train!\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "os.makedirs('./outputs/samples', exist_ok=True)\n",
    "os.makedirs('./outputs/checkpoints', exist_ok=True)\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS * len(dataloader))\n",
    "\n",
    "use_amp = device.type == 'cuda'\n",
    "scaler = torch.cuda.amp.GradScaler() if use_amp else None\n",
    "\n",
    "print(f\"Training: {EPOCHS} epochs, AMP={use_amp}\")\n",
    "print('=' * 60)\n",
    "\n",
    "best_loss = float('inf')\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    pbar = tqdm(dataloader, desc=f'Epoch {epoch+1}/{EPOCHS}')\n",
    "    \n",
    "    for images, _ in pbar:\n",
    "        images = images.to(device)\n",
    "        \n",
    "        # Encode to latent space\n",
    "        with torch.no_grad():\n",
    "            latents = TAESDWrapper.encode(vae, images)\n",
    "        \n",
    "        # Training step\n",
    "        loss_dict = model.training_step(latents, optimizer, scaler, use_amp)\n",
    "        scheduler.step()\n",
    "        \n",
    "        epoch_loss += loss_dict['total']\n",
    "        pbar.set_postfix(loss=f\"{loss_dict['total']:.4f}\", diff=f\"{loss_dict['diffusion']:.4f}\")\n",
    "    \n",
    "    avg = epoch_loss / len(dataloader)\n",
    "    print(f'Epoch {epoch+1}: loss={avg:.4f}')\n",
    "    \n",
    "    # Generate samples\n",
    "    if (epoch + 1) % SAMPLE_EVERY == 0 or epoch == EPOCHS - 1:\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            z = model.sample(batch_size=16, steps=SAMPLE_STEPS, ddim=True, progress=False)\n",
    "            imgs = TAESDWrapper.decode(vae, z)\n",
    "        save_image(imgs, f'./outputs/samples/epoch_{epoch+1:03d}.png', nrow=4, normalize=True, value_range=(-1,1))\n",
    "        print(f'  Samples saved: ./outputs/samples/epoch_{epoch+1:03d}.png')\n",
    "    \n",
    "    # Checkpoint\n",
    "    if avg < best_loss:\n",
    "        best_loss = avg\n",
    "        torch.save(model.state_dict(), './outputs/checkpoints/best.pt')\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        torch.save({'epoch': epoch+1, 'model': model.state_dict(), 'opt': optimizer.state_dict()},\n",
    "                   f'./outputs/checkpoints/epoch_{epoch+1:03d}.pt')\n",
    "\n",
    "print('=' * 60)\n",
    "print(f'Done! Best loss: {best_loss:.4f}')"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# @title 7. Display Generated Samples\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import glob\n",
    "\n",
    "sample_files = sorted(glob.glob('./outputs/samples/epoch_*.png'))\n",
    "if sample_files:\n",
    "    img = Image.open(sample_files[-1])\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(img)\n",
    "    plt.title(f'LiquidFlow — {MODEL_VARIANT}, {IMAGE_SIZE}px (latest)')\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "else:\n",
    "    print('No samples yet — run training first!')"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "colab": {"name": "LiquidFlow_Train", "provenance": []},
  "kernelspec": {"display_name": "Python 3", "name": "python3"}
 },
 "nbformat": 4,
 "nbformat_minor": 0
}