Upload LiquidFlow_Colab.ipynb

LiquidFlow_Colab.ipynb

@@ -12,13 +12,7 @@
     "- **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 (< 1M params) for fast encoding\n",
-    "\n",
-    "Based on:\n",
-    "- CfC: Hasani et al., Nature MI 2022\n",
-    "- Mamba-2: Dao & Gu, 2024  \n",
-    "- PINN Diffusion: Bastek & Sun, ICLR 2025\n",
-    "- DiMSUM: NeurIPS 2024\n",
+    "- **TAESD VAE** — Tiny AutoEncoder (< 3M params) for fast encoding\n",
     "\n",
     "---\n",
     "## Quick Start\n",
@@ -33,9 +27,8 @@
    "execution_count": null,
    "metadata": {},
    "source": [
-    "# @title 1. Install Dependencies (~2 min)\n",
-    "!pip install -q torch torchvision diffusers tqdm pillow numpy\n",
-    "!pip install -q git+https://github.com/huggingface/diffusers.git\n",
+    "# @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",
@@ -52,7 +45,11 @@
    "metadata": {},
    "source": [
     "# @title 2. Clone LiquidFlow Repository\n",
-    "!git clone https://huggingface.co/krystv/LiquidFlow-Gen /content/LiquidFlow\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",
@@ -60,8 +57,7 @@
     "\n",
     "from liquid_flow.generator import create_liquidflow\n",
     "from liquid_flow.vae_wrapper import TAESDWrapper\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F"
+    "print('LiquidFlow imported successfully!')"
    ],
    "outputs": []
   },
@@ -70,33 +66,35 @@
    "execution_count": null,
    "metadata": {},
    "source": [
-    "# @title 3. Configuration — Adjust these settings!\n",
+    "# @title 3. Configuration\n",
     "\n",
-    "# Model size: 'tiny' (~2M), 'small' (~8M), 'base' (~30M)\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, 512 needs more VRAM\n",
-    "IMAGE_SIZE = 128  # @param [64, 128, 256, 512]\n",
+    "# Image size (128 recommended for T4 free tier)\n",
+    "IMAGE_SIZE = 128  # @param [64, 128, 256, 512] {type:\"integer\"}\n",
     "\n",
-    "# Training\n",
-    "BATCH_SIZE = 32  # @param [8, 16, 32, 64]\n",
-    "EPOCHS = 50  # @param [10, 25, 50, 100]\n",
-    "LEARNING_RATE = 2e-4  # @param [1e-4, 2e-4, 5e-4, 1e-3]\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 (DDIM steps)\n",
-    "SAMPLE_EVERY = 5  # @param [1, 5, 10]\n",
-    "SAMPLE_STEPS = 50  # @param [20, 50, 100]\n",
+    "# Sampling\n",
+    "SAMPLE_EVERY = 5  # @param {type:\"integer\"}\n",
+    "SAMPLE_STEPS = 50  # @param {type:\"integer\"}\n",
     "\n",
-    "# Physics regularization weights\n",
-    "PHYSICS_TV_WEIGHT = 0.01\n",
-    "PHYSICS_SPEC_WEIGHT = 0.01\n",
-    "PHYSICS_GRAD_WEIGHT = 0.001\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}\")\n",
-    "print(f\"Physics loss: TV={PHYSICS_TV_WEIGHT}, Spec={PHYSICS_SPEC_WEIGHT}, Grad={PHYSICS_GRAD_WEIGHT}\")"
+    "print(f\"Config: {MODEL_VARIANT} model, {IMAGE_SIZE}px, batch={BATCH_SIZE}, epochs={EPOCHS}, lr={LEARNING_RATE}\")"
    ],
    "outputs": []
   },
@@ -117,31 +115,17 @@
     "# Load TAESD (Tiny AutoEncoder)\n",
     "print(\"Loading TAESD VAE...\")\n",
     "vae = TAESDWrapper.load(device)\n",
-    "print(f\"VAE loaded! Latent compression: {IMAGE_SIZE}x{IMAGE_SIZE} → {IMAGE_SIZE//8}x{IMAGE_SIZE//8}\")\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\"Creating {MODEL_VARIANT} LiquidFlow model...\")\n",
-    "model = create_liquidflow(\n",
-    "    variant=MODEL_VARIANT,\n",
-    "    image_size=IMAGE_SIZE,\n",
-    "    physics_weights={\n",
-    "        'tv': PHYSICS_TV_WEIGHT,\n",
-    "        'cons': 0.001,\n",
-    "        'spec': PHYSICS_SPEC_WEIGHT,\n",
-    "        'grad': PHYSICS_GRAD_WEIGHT,\n",
-    "    },\n",
-    ")\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 = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
-    "print(f\"Model: {n_params:,} parameters ({n_params/1e6:.1f}M)\")\n",
-    "\n",
-    "# Memory estimate\n",
-    "latent_size = IMAGE_SIZE // 8\n",
-    "mem_per_sample = latent_size * latent_size * 4 * 4 / 1e6  # MB\n",
-    "print(f\"Memory per sample: {mem_per_sample:.1f} MB\")\n",
-    "print(f\"Estimated batch memory: {mem_per_sample * BATCH_SIZE:.1f} MB\")\n",
-    "print(f\"T4 VRAM: 15 GB — should fit!\" if mem_per_sample * BATCH_SIZE < 10 else \"Watch memory!\")"
+    "n_params = model.count_parameters()\n",
+    "print(f\"Model parameters: {n_params:,} ({n_params/1e6:.1f}M)\")\n",
+    "print(f\"\\nReady to train!\")"
    ],
    "outputs": []
   },
@@ -154,8 +138,9 @@
     "\n",
     "transform = transforms.Compose([\n",
     "    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
     "    transforms.ToTensor(),\n",
-    "    transforms.Normalize([0.5], [0.5]),\n",
+    "    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n",
     "])\n",
     "\n",
     "if DATASET == 'cifar10':\n",
@@ -164,17 +149,14 @@
     "    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",
-    "else:\n",
-    "    raise ValueError(f\"Unknown dataset: {DATASET}\")\n",
     "\n",
     "dataloader = DataLoader(\n",
     "    dataset, batch_size=BATCH_SIZE, shuffle=True,\n",
-    "    num_workers=min(4, os.cpu_count() or 1),\n",
-    "    pin_memory=True, drop_last=True,\n",
+    "    num_workers=2, pin_memory=True, drop_last=True,\n",
     ")\n",
     "\n",
-    "print(f\"Dataset: {DATASET}\")\n",
-    "print(f\"Images: {len(dataset):,}, Batches per epoch: {len(dataloader)}\")"
+    "print(f\"Dataset: {DATASET} ({len(dataset):,} images)\")\n",
+    "print(f\"Batches per epoch: {len(dataloader)}\")"
    ],
    "outputs": []
   },
@@ -183,109 +165,64 @@
    "execution_count": null,
    "metadata": {},
    "source": [
-    "# @title 6. Training Loop\n",
-    "\n",
+    "# @title 6. Train!\n",
     "from torchvision.utils import save_image\n",
-    "import math\n",
     "\n",
     "os.makedirs('./outputs/samples', exist_ok=True)\n",
     "os.makedirs('./outputs/checkpoints', exist_ok=True)\n",
     "\n",
-    "# Optimizer\n",
-    "optimizer = torch.optim.AdamW(\n",
-    "    model.parameters(),\n",
-    "    lr=LEARNING_RATE,\n",
-    "    betas=(0.9, 0.999),\n",
-    "    weight_decay=1e-4,\n",
-    ")\n",
-    "\n",
-    "# Cosine LR scheduler\n",
-    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
-    "    optimizer, T_max=EPOCHS * len(dataloader)\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",
-    "# AMP\n",
     "use_amp = device.type == 'cuda'\n",
     "scaler = torch.cuda.amp.GradScaler() if use_amp else None\n",
     "\n",
-    "print(f\"Training: {EPOCHS} epochs, LR={LEARNING_RATE}, AMP={use_amp}\")\n",
-    "print(\"=\"*60)\n",
+    "print(f\"Training: {EPOCHS} epochs, AMP={use_amp}\")\n",
+    "print('=' * 60)\n",
     "\n",
-    "global_step = 0\n",
     "best_loss = float('inf')\n",
     "\n",
     "for epoch in range(EPOCHS):\n",
     "    model.train()\n",
-    "    epoch_total = 0\n",
-    "    \n",
-    "    pbar = tqdm(dataloader, desc=f\"Epoch {epoch+1}/{EPOCHS}\")\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\n",
+    "        # Encode to latent space\n",
     "        with torch.no_grad():\n",
     "            latents = TAESDWrapper.encode(vae, images)\n",
     "        \n",
-    "        # Training step with physics regularization\n",
+    "        # Training step\n",
     "        loss_dict = model.training_step(latents, optimizer, scaler, use_amp)\n",
-    "        \n",
-    "        # Track\n",
-    "        total_loss = loss_dict['total']\n",
-    "        epoch_total += total_loss\n",
-    "        \n",
-    "        # Update scheduler\n",
     "        scheduler.step()\n",
     "        \n",
-    "        # Progress bar\n",
-    "        pbar.set_postfix({\n",
-    "            'loss': f\"{total_loss:.4f}\",\n",
-    "            'diff': f\"{loss_dict.get('diffusion', 0):.4f}\",\n",
-    "            'phys': f\"{loss_dict.get('physics', 0):.4f}\",\n",
-    "            'lr': f\"{optimizer.param_groups[0]['lr']:.2e}\",\n",
-    "        })\n",
-    "        \n",
-    "        global_step += 1\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_loss = epoch_total / len(dataloader)\n",
-    "    print(f\"Epoch {epoch+1}: avg_loss={avg_loss:.4f}\")\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",
-    "        print(\"  Generating samples...\")\n",
     "        model.eval()\n",
     "        with torch.no_grad():\n",
-    "            latents_gen = model.sample(\n",
-    "                batch_size=16,\n",
-    "                steps=SAMPLE_STEPS,\n",
-    "                ddim=True,\n",
-    "                progress=False,\n",
-    "            )\n",
-    "            images_gen = TAESDWrapper.decode(vae, latents_gen)\n",
-    "        \n",
-    "        save_image(\n",
-    "            images_gen, f'./outputs/samples/epoch_{epoch+1:03d}.png',\n",
-    "            nrow=4, normalize=True, value_range=(-1, 1)\n",
-    "        )\n",
-    "        print(f\"  Saved to ./outputs/samples/epoch_{epoch+1:03d}.png\")\n",
-    "    \n",
-    "    # Save checkpoint\n",
-    "    if (epoch + 1) % 10 == 0 or epoch == EPOCHS - 1:\n",
-    "        torch.save({\n",
-    "            'epoch': epoch + 1,\n",
-    "            'model_state_dict': model.state_dict(),\n",
-    "            'optimizer_state_dict': optimizer.state_dict(),\n",
-    "            'loss': avg_loss,\n",
-    "        }, f'./outputs/checkpoints/epoch_{epoch+1:03d}.pt')\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",
-    "    if avg_loss < best_loss:\n",
-    "        best_loss = avg_loss\n",
-    "        torch.save(model.state_dict(), './outputs/checkpoints/best_model.pt')\n",
-    "\n",
-    "print(\"=\"*60)\n",
-    "print(f\"Training complete! Best loss: {best_loss:.4f}\")\n",
-    "print(f\"Checkpoints saved to ./outputs/checkpoints/\")\n",
-    "print(f\"Samples saved to ./outputs/samples/\")"
+    "    # 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": []
   },
@@ -294,134 +231,28 @@
    "execution_count": null,
    "metadata": {},
    "source": [
-    "# @title 7. Generate & Display Samples\n",
-    "\n",
+    "# @title 7. Display Generated Samples\n",
     "import matplotlib.pyplot as plt\n",
     "from PIL import Image\n",
     "import glob\n",
     "\n",
-    "# Load latest sample\n",
     "sample_files = sorted(glob.glob('./outputs/samples/epoch_*.png'))\n",
     "if sample_files:\n",
-    "    latest = sample_files[-1]\n",
-    "    img = Image.open(latest)\n",
-    "    plt.figure(figsize=(12, 12))\n",
+    "    img = Image.open(sample_files[-1])\n",
+    "    plt.figure(figsize=(10, 10))\n",
     "    plt.imshow(img)\n",
-    "    plt.title(f'LiquidFlow Samples — {MODEL_VARIANT} model, {IMAGE_SIZE}px')\n",
+    "    plt.title(f'LiquidFlow — {MODEL_VARIANT}, {IMAGE_SIZE}px (latest)')\n",
     "    plt.axis('off')\n",
     "    plt.show()\n",
     "else:\n",
-    "    print(\"No samples generated yet. Train for more epochs!\")"
+    "    print('No samples yet — run training first!')"
    ],
    "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "source": [
-    "# @title 8. Export Model for Mobile (ONNX)\n",
-    "\n",
-    "# LiquidFlow can be exported to ONNX for mobile deployment\n",
-    "# since it uses pure PyTorch (no custom CUDA kernels)\n",
-    "\n",
-    "def export_to_onnx(model, output_path='liquidflow_model.onnx', image_size=128):\n",
-    "    \"\"\"Export LiquidFlow to ONNX for mobile deployment.\"\"\"\n",
-    "    model = model.cpu()\n",
-    "    model.eval()\n",
-    "    \n",
-    "    latent_size = image_size // 8\n",
-    "    \n",
-    "    # Dummy inputs\n",
-    "    x = torch.randn(1, 4, latent_size, latent_size)\n",
-    "    t = torch.tensor([500], dtype=torch.long)\n",
-    "    \n",
-    "    # Export\n",
-    "    torch.onnx.export(\n",
-    "        model,\n",
-    "        (x, t),\n",
-    "        output_path,\n",
-    "        input_names=['noisy_latent', 'timestep'],\n",
-    "        output_names=['predicted_noise'],\n",
-    "        dynamic_axes={\n",
-    "            'noisy_latent': {0: 'batch'},\n",
-    "            'predicted_noise': {0: 'batch'},\n",
-    "        },\n",
-    "        opset_version=14,\n",
-    "    )\n",
-    "    \n",
-    "    import os\n",
-    "    size_mb = os.path.getsize(output_path) / 1e6\n",
-    "    print(f\"ONNX model exported to {output_path} ({size_mb:.1f} MB)\")\n",
-    "    return output_path\n",
-    "\n",
-    "# Load best model and export\n",
-    "best_model_path = './outputs/checkpoints/best_model.pt'\n",
-    "if os.path.exists(best_model_path):\n",
-    "    model.load_state_dict(torch.load(best_model_path, map_location='cpu'))\n",
-    "    export_to_onnx(model, 'liquidflow_128.onnx', IMAGE_SIZE)\n",
-    "    print(\"Ready for mobile deployment!\")\n",
-    "else:\n",
-    "    print(\"Train model first before exporting.\")"
-   ],
-   "outputs": []
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Architecture Details\n",
-    "\n",
-    "### LiquidFlow Block Architecture\n",
-    "```\n",
-    "Input → [CfC Gate → Mamba-2 SSD → CfC Gate] → Output\n",
-    "            ↑                        ↑\n",
-    "        Adaptive time gate      Gated output\n",
-    "```\n",
-    "\n",
-    "### CfC (Closed-form Continuous-time) Cell\n",
-    "```\n",
-    "h(t) = σ(-f(x,I;θ_f)·t) ⊙ g(x,I;θ_g) + (1-σ(-f(x,I;θ_f)·t)) ⊙ h(x,I;θ_h)\n",
-    "```\n",
-    "- **No ODE solver needed** — 100x faster than Neural ODEs\n",
-    "- Time-continuous gating adaptively controls information flow\n",
-    "- Closed-form solution → stable gradients\n",
-    "\n",
-    "### Mamba-2 SSD (State Space Duality)\n",
-    "```\n",
-    "h_t = A_t * h_{t-1} + B_t * x_t\n",
-    "y_t = C_t^T * h_t\n",
-    "```\n",
-    "- **O(N) linear complexity** vs Transformers O(N²)\n",
-    "- **Parallelizable** via associative scan (Blelloch)\n",
-    "- **Scalar-A** formulation enables chunk-scan optimization\n",
-    "- Pure PyTorch — no CUDA kernels needed\n",
-    "\n",
-    "### Physics-Informed Regularization\n",
-    "- **Total Variation**: `L_TV = ||∇_x x̂||₁ + ||∇_y x̂||₁`\n",
-    "- **Spectral**: Penalize high-frequency artifacts\n",
-    "- **Gradient**: Sobolev norm for stable training\n",
-    "- Pattern from Bastek & Sun (ICLR 2025): physics loss as training-only regularizer\n",
-    "\n",
-    "### Model Variants\n",
-    "| Variant | Params | Hidden Dim | Stages | Blocks | T4 VRAM |\n",
-    "|---------|--------|------------|--------|--------|---------|\n",
-    "| Tiny    | ~2M    | 128        | 2      | 2      | < 2 GB  |\n",
-    "| Small   | ~8M    | 256        | 4      | 4      | ~4 GB   |\n",
-    "| Base    | ~30M   | 384        | 6      | 6      | ~8 GB   |"
-   ]
   }
  ],
  "metadata": {
-  "colab": {
-   "name": "LiquidFlow: LiquidNN + Mamba-2 SSD Image Generator",
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  }
+  "colab": {"name": "LiquidFlow_Train", "provenance": []},
+  "kernelspec": {"display_name": "Python 3", "name": "python3"}
  },
  "nbformat": 4,
  "nbformat_minor": 0