Fix: streaming dataset in notebook (no full download on Colab)

LiquidGen_Colab_Notebook.ipynb

@@ -21,21 +21,17 @@
         "\n",
         "**A novel attention-free diffusion model using CfC Liquid Neural Network dynamics.**\n",
         "\n",
-        "### Key Features:\n",
         "- **No Attention** \u2014 O(n) complexity using liquid time constants\n",
         "- **Fully Parallelizable** \u2014 No sequential ODE solving\n",
-        "- **Flow Matching** \u2014 Modern velocity-prediction training\n",
-        "- **Frozen Flux VAE** \u2014 16-channel latent space\n",
-        "- **Fits 16GB VRAM** \u2014 Designed for Colab free tier\n",
-        "\n",
-        "Based on: Liquid Time-constant Networks (NeurIPS 2020), CfC (Nature MI 2022), ZigMa (ECCV 2024), DiMSUM (NeurIPS 2024)\n"
+        "- **Streaming Dataset** \u2014 No full download, starts training immediately\n",
+        "- **Fits 16GB VRAM** \u2014 Designed for Colab free tier T4\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udce6 Install Dependencies"
+        "## \ud83d\udce6 Step 1: Install"
       ]
     },
     {
@@ -44,14 +40,21 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "!pip install -q torch torchvision diffusers datasets accelerate huggingface_hub Pillow"
+        "!pip install -q torch torchvision diffusers datasets accelerate Pillow"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udd27 Configuration"
+        "## \ud83d\udd27 Step 2: Configuration\n",
+        "\n",
+        "**Dataset options:**\n",
+        "| Dataset | Size | Download | Type |\n",
+        "|---------|------|----------|------|\n",
+        "| `huggan/wikiart` | ~80K | **Streaming** (no download!) | Art, 27 styles |\n",
+        "| `reach-vb/pokemon-blip-captions` | 833 | 95MB (fast) | Pokemon |\n",
+        "| `huggan/flowers-102-categories` | 8K | 330MB | Flowers |\n"
       ]
     },
     {
@@ -60,48 +63,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "MODEL_SIZE = \"small\"  # \"small\" (~55M), \"base\" (~140M), \"large\" (~280M)\n",
-        "IMAGE_SIZE = 256  # 256 or 512\n",
-        "DATASET_NAME = \"huggan/wikiart\"\n",
-        "IMAGE_COLUMN = \"image\"\n",
-        "LABEL_COLUMN = \"style\"  # \"style\" (27), \"genre\" (11), \"\" for unconditional\n",
-        "NUM_CLASSES = 27\n",
-        "BATCH_SIZE = 8\n",
-        "GRADIENT_ACCUMULATION = 4\n",
-        "LEARNING_RATE = 1e-4\n",
-        "WEIGHT_DECAY = 0.01\n",
-        "MAX_GRAD_NORM = 2.0\n",
-        "NUM_EPOCHS = 50\n",
-        "WARMUP_STEPS = 500\n",
-        "EMA_DECAY = 0.9999\n",
-        "NUM_SAMPLE_STEPS = 50\n",
-        "CFG_SCALE = 2.0\n",
-        "OUTPUT_DIR = \"/content/liquidgen_outputs\"\n",
-        "SAVE_EVERY = 2000\n",
-        "SAMPLE_EVERY = 500\n",
-        "LOG_EVERY = 50\n",
-        "PUSH_TO_HUB = False\n",
-        "HUB_MODEL_ID = \"\"\n",
-        "VAE_ID = \"black-forest-labs/FLUX.1-schnell\"\n",
-        "VAE_SUBFOLDER = \"vae\"\n",
-        "\n",
-        "import torch\n",
-        "if torch.cuda.is_available():\n",
-        "    gpu = torch.cuda.get_device_name(0)\n",
-        "    mem = torch.cuda.get_device_properties(0).total_mem / 1024**3\n",
-        "    print(f\"GPU: {gpu} ({mem:.1f} GB)\")\n",
-        "    if mem < 12: print(\"\u26a0\ufe0f Low VRAM! Use small model, 256px, bs=4\")\n",
-        "    elif mem < 20: print(\"\u2705 T4 detected. Good for base model, 256px\")\n",
-        "    else: print(\"\ud83d\ude80 Large GPU! Can run large model, 512px\")\n",
-        "else:\n",
-        "    print(\"\u26a0\ufe0f No GPU! Go to Runtime \u2192 Change runtime type \u2192 GPU\")"
+        "# ============================================================================\n# CONFIGURATION\n# ============================================================================\n\nMODEL_SIZE = \"small\"  # \"small\" (~55M), \"base\" (~140M), \"large\" (~280M)\nIMAGE_SIZE = 256      # 256 or 512\n\n# --- Dataset (Option A: WikiArt streaming \u2014 NO download) ---\nDATASET_NAME = \"huggan/wikiart\"\nIMAGE_COLUMN = \"image\"\nLABEL_COLUMN = \"style\"   # \"style\"(27), \"genre\"(11), \"\" for unconditional\nNUM_CLASSES = 27\nUSE_STREAMING = True      # KEY: no full download!\n\n# --- Dataset (Option B: Pokemon \u2014 small, fast download, good for testing) ---\n# DATASET_NAME = \"reach-vb/pokemon-blip-captions\"\n# IMAGE_COLUMN = \"image\"; LABEL_COLUMN = \"\"; NUM_CLASSES = 0; USE_STREAMING = False\n\n# --- Training ---\nBATCH_SIZE = 8; GRADIENT_ACCUMULATION = 4\nLEARNING_RATE = 1e-4; WEIGHT_DECAY = 0.01; MAX_GRAD_NORM = 2.0\nMAX_STEPS = 20000; WARMUP_STEPS = 500; EMA_DECAY = 0.9999\nNUM_SAMPLE_STEPS = 50; CFG_SCALE = 2.0\n\n# --- Saving ---\nOUTPUT_DIR = \"/content/liquidgen_outputs\"\nSAVE_EVERY = 5000; SAMPLE_EVERY = 500; LOG_EVERY = 50\n\n# --- VAE ---\nVAE_ID = \"black-forest-labs/FLUX.1-schnell\"\nVAE_SUBFOLDER = \"vae\"\nSCALE, SHIFT = 0.3611, 0.1159\n\nimport torch\nif torch.cuda.is_available():\n    gpu = torch.cuda.get_device_name(0)\n    mem = torch.cuda.get_device_properties(0).total_mem / 1024**3\n    print(f\"GPU: {gpu} ({mem:.1f} GB)\")\nelse:\n    print(\"No GPU! Go to Runtime > Change runtime type > GPU\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83c\udfd7\ufe0f Model Architecture"
+        "## \ud83c\udfd7\ufe0f Step 3: Model Architecture"
       ]
     },
     {
@@ -117,7 +86,7 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udd04 Training Utilities"
+        "## \ud83d\udd04 Step 4: Training Utilities"
       ]
     },
     {
@@ -126,77 +95,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import os, json, time, math\n",
-        "import numpy as np\n",
-        "from torch.utils.data import DataLoader, Dataset\n",
-        "from torch.amp import autocast, GradScaler\n",
-        "from torchvision import transforms\n",
-        "from torchvision.utils import save_image\n",
-        "from PIL import Image\n",
-        "\n",
-        "class FlowMatchingScheduler:\n",
-        "    def __init__(self, min_t=0.001, max_t=0.999): self.min_t, self.max_t = min_t, max_t\n",
-        "    def sample_timesteps(self, bs, dev): return torch.rand(bs, device=dev) * (self.max_t - self.min_t) + self.min_t\n",
-        "    def add_noise(self, x0, noise, t): t = t.view(-1,1,1,1); return (1-t)*x0 + t*noise\n",
-        "    def get_velocity_target(self, x0, noise): return noise - x0\n",
-        "    @torch.no_grad()\n",
-        "    def sample(self, model, shape, dev, num_steps=50, labels=None, cfg=1.0):\n",
-        "        model.eval(); x = torch.randn(shape, device=dev)\n",
-        "        dt = 1.0 / num_steps\n",
-        "        for t_val in torch.linspace(1.0, dt, num_steps, device=dev):\n",
-        "            t = torch.full((shape[0],), t_val.item(), device=dev)\n",
-        "            with torch.amp.autocast(\"cuda\"):\n",
-        "                if cfg > 1.0 and labels is not None:\n",
-        "                    vc = model(x,t,labels); vu = model(x,t,torch.zeros_like(labels))\n",
-        "                    v = vu + cfg * (vc - vu)\n",
-        "                else: v = model(x, t, labels)\n",
-        "            x = x - dt * v.float()\n",
-        "        return x\n",
-        "\n",
-        "class EMAModel:\n",
-        "    def __init__(self, model, decay=0.9999):\n",
-        "        self.decay = decay\n",
-        "        self.shadow = {n: p.clone().detach() for n,p in model.named_parameters() if p.requires_grad}\n",
-        "    @torch.no_grad()\n",
-        "    def update(self, model):\n",
-        "        for n,p in model.named_parameters():\n",
-        "            if p.requires_grad and n in self.shadow: self.shadow[n].mul_(self.decay).add_(p.data, alpha=1-self.decay)\n",
-        "    def apply(self, model):\n",
-        "        self.backup = {n: p.data.clone() for n,p in model.named_parameters() if p.requires_grad}\n",
-        "        for n,p in model.named_parameters():\n",
-        "            if p.requires_grad and n in self.shadow: p.data.copy_(self.shadow[n])\n",
-        "    def restore(self, model):\n",
-        "        for n,p in model.named_parameters():\n",
-        "            if p.requires_grad and n in self.backup: p.data.copy_(self.backup[n])\n",
-        "\n",
-        "class ImageDataset(Dataset):\n",
-        "    def __init__(self, ds, tf, img_col, lbl_col=\"\"): self.ds, self.tf, self.ic, self.lc = ds, tf, img_col, lbl_col\n",
-        "    def __len__(self): return len(self.ds)\n",
-        "    def __getitem__(self, i):\n",
-        "        item = self.ds[i]; img = item[self.ic]\n",
-        "        if img.mode != \"RGB\": img = img.convert(\"RGB\")\n",
-        "        label = item[self.lc] if self.lc and self.lc in item else -1\n",
-        "        return self.tf(img), label\n",
-        "\n",
-        "def cosine_sched(opt, warmup, total):\n",
-        "    def lr(s):\n",
-        "        if s < warmup: return s / max(1, warmup)\n",
-        "        return max(0, 0.5*(1+math.cos(math.pi*(s-warmup)/max(1,total-warmup))))\n",
-        "    return torch.optim.lr_scheduler.LambdaLR(opt, lr)\n",
-        "\n",
-        "MODEL_CONFIGS = {\n",
-        "    \"small\": dict(embed_dim=512, depth=12, spatial_kernel=7, scan_kernel=31, expand_ratio=2.0, mlp_ratio=3.0),\n",
-        "    \"base\": dict(embed_dim=640, depth=18, spatial_kernel=7, scan_kernel=31, expand_ratio=2.0, mlp_ratio=4.0),\n",
-        "    \"large\": dict(embed_dim=768, depth=24, spatial_kernel=7, scan_kernel=31, expand_ratio=2.5, mlp_ratio=4.0),\n",
-        "}\n",
-        "print(\"\u2705 Training utilities ready!\")"
+        "import os, time, math\nimport numpy as np\nfrom torch.utils.data import DataLoader, IterableDataset, Dataset\nfrom torch.amp import autocast, GradScaler\nfrom torchvision import transforms\nfrom torchvision.utils import save_image\nfrom PIL import Image\n\nclass StreamingImageDataset(IterableDataset):\n    \"\"\"Streaming dataset \u2014 NO full download. Images load on-the-fly.\"\"\"\n    def __init__(self, name, img_col=\"image\", lbl_col=\"\", img_size=256,\n                 split=\"train\", config=\"\", buffer=1000, seed=42):\n        super().__init__()\n        self.name, self.img_col, self.lbl_col = name, img_col, lbl_col\n        self.split, self.config, self.buffer, self.seed = split, config, buffer, seed\n        self.tf = transforms.Compose([\n            transforms.Resize(img_size, interpolation=transforms.InterpolationMode.LANCZOS),\n            transforms.CenterCrop(img_size), transforms.RandomHorizontalFlip(), transforms.ToTensor()])\n\n    def __iter__(self):\n        from datasets import load_dataset\n        kw = {\"name\": self.config} if self.config else {}\n        ds = load_dataset(self.name, split=self.split, streaming=True, **kw)\n        ds = ds.shuffle(seed=self.seed, buffer_size=self.buffer)\n        for item in ds:\n            try:\n                img = item[self.img_col]\n                if img.mode != \"RGB\": img = img.convert(\"RGB\")\n                lbl = item[self.lbl_col] if self.lbl_col and self.lbl_col in item else -1\n                yield self.tf(img), lbl\n            except: continue\n\nclass MapImageDataset(Dataset):\n    \"\"\"For small datasets (<500MB) \u2014 downloads once.\"\"\"\n    def __init__(self, name, img_col=\"image\", lbl_col=\"\", img_size=256, split=\"train\"):\n        from datasets import load_dataset\n        print(f\"Downloading {name}...\")\n        self.ds = load_dataset(name, split=split)\n        self.img_col, self.lbl_col = img_col, lbl_col\n        self.tf = transforms.Compose([\n            transforms.Resize(img_size, interpolation=transforms.InterpolationMode.LANCZOS),\n            transforms.CenterCrop(img_size), transforms.RandomHorizontalFlip(), transforms.ToTensor()])\n        print(f\"  {len(self.ds)} images\")\n\n    def __len__(self): return len(self.ds)\n    def __getitem__(self, i):\n        item = self.ds[i]; img = item[self.img_col]\n        if img.mode != \"RGB\": img = img.convert(\"RGB\")\n        lbl = item[self.lbl_col] if self.lbl_col and self.lbl_col in item else -1\n        return self.tf(img), lbl\n\nclass FlowMatchingScheduler:\n    def __init__(self, min_t=0.001, max_t=0.999): self.min_t, self.max_t = min_t, max_t\n    def sample_t(self, bs, dev): return torch.rand(bs, device=dev)*(self.max_t-self.min_t)+self.min_t\n    def add_noise(self, x0, noise, t): return (1-t.view(-1,1,1,1))*x0 + t.view(-1,1,1,1)*noise\n    def velocity(self, x0, noise): return noise - x0\n    @torch.no_grad()\n    def sample(self, model, shape, dev, steps=50, labels=None, cfg=1.0):\n        model.eval(); x = torch.randn(shape, device=dev); dt = 1.0/steps\n        for tv in torch.linspace(1.0, dt, steps, device=dev):\n            t = torch.full((shape[0],), tv.item(), device=dev)\n            with torch.amp.autocast(\"cuda\"):\n                if cfg > 1.0 and labels is not None:\n                    vc = model(x,t,labels); vu = model(x,t,torch.zeros_like(labels))\n                    v = vu + cfg*(vc-vu)\n                else: v = model(x,t,labels)\n            x = x - dt * v.float()\n        return x\n\nclass EMAModel:\n    def __init__(self, model, decay=0.9999):\n        self.decay = decay\n        self.shadow = {n:p.clone().detach() for n,p in model.named_parameters() if p.requires_grad}\n    @torch.no_grad()\n    def update(self, m):\n        for n,p in m.named_parameters():\n            if p.requires_grad and n in self.shadow: self.shadow[n].mul_(self.decay).add_(p.data, alpha=1-self.decay)\n    def apply(self, m):\n        self.bk = {n:p.data.clone() for n,p in m.named_parameters() if p.requires_grad}\n        for n,p in m.named_parameters():\n            if p.requires_grad and n in self.shadow: p.data.copy_(self.shadow[n])\n    def restore(self, m):\n        for n,p in m.named_parameters():\n            if p.requires_grad and n in self.bk: p.data.copy_(self.bk[n])\n\ndef cosine_sched(opt, warmup, total):\n    def lr(s):\n        if s < warmup: return s/max(1,warmup)\n        return max(0, 0.5*(1+math.cos(math.pi*(s-warmup)/max(1,total-warmup))))\n    return torch.optim.lr_scheduler.LambdaLR(opt, lr)\n\nMODEL_CONFIGS = {\n    \"small\": dict(embed_dim=512, depth=12, spatial_kernel=7, scan_kernel=31, expand_ratio=2.0, mlp_ratio=3.0),\n    \"base\": dict(embed_dim=640, depth=18, spatial_kernel=7, scan_kernel=31, expand_ratio=2.0, mlp_ratio=4.0),\n    \"large\": dict(embed_dim=768, depth=24, spatial_kernel=7, scan_kernel=31, expand_ratio=2.5, mlp_ratio=4.0),\n}\nprint(\"Training utilities ready!\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udcca Load Dataset & VAE"
+        "## \ud83d\udcca Step 5: Load Dataset & VAE"
       ]
     },
     {
@@ -205,33 +111,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from datasets import load_dataset\n",
-        "from diffusers import AutoencoderKL\n",
-        "\n",
-        "print(f\"Loading dataset: {DATASET_NAME}...\")\n",
-        "dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
-        "print(f\"  {len(dataset)} images\")\n",
-        "\n",
-        "transform = transforms.Compose([\n",
-        "    transforms.Resize(IMAGE_SIZE, interpolation=transforms.InterpolationMode.LANCZOS),\n",
-        "    transforms.CenterCrop(IMAGE_SIZE), transforms.RandomHorizontalFlip(), transforms.ToTensor(),\n",
-        "])\n",
-        "train_ds = ImageDataset(dataset, transform, IMAGE_COLUMN, LABEL_COLUMN)\n",
-        "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True, drop_last=True)\n",
-        "\n",
-        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-        "vae = AutoencoderKL.from_pretrained(VAE_ID, subfolder=VAE_SUBFOLDER, torch_dtype=torch.float16).to(device).eval()\n",
-        "for p in vae.parameters(): p.requires_grad_(False)\n",
-        "print(f\"VAE: {sum(p.numel() for p in vae.parameters())/1e6:.1f}M params (frozen)\")\n",
-        "SCALE, SHIFT = 0.3611, 0.1159\n",
-        "print(\"\u2705 Ready!\")"
+        "from diffusers import AutoencoderKL\n\nif USE_STREAMING:\n    print(f\"Loading {DATASET_NAME} in STREAMING mode (no full download)...\")\n    train_ds = StreamingImageDataset(DATASET_NAME, IMAGE_COLUMN, LABEL_COLUMN, IMAGE_SIZE, buffer=1000)\n    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, num_workers=0, pin_memory=True)\n    print(\"  Streaming ready! Images load on-the-fly.\")\nelse:\n    train_ds = MapImageDataset(DATASET_NAME, IMAGE_COLUMN, LABEL_COLUMN, IMAGE_SIZE)\n    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True, drop_last=True)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Loading VAE to {device}...\")\nvae = AutoencoderKL.from_pretrained(VAE_ID, subfolder=VAE_SUBFOLDER, torch_dtype=torch.float16).to(device).eval()\nfor p in vae.parameters(): p.requires_grad_(False)\nprint(f\"  VAE: {sum(p.numel() for p in vae.parameters())/1e6:.1f}M params (frozen)\")\nprint(\"Ready!\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83c\udfcb\ufe0f Create Model & Train"
+        "## \ud83c\udfcb\ufe0f Step 6: Create Model"
       ]
     },
     {
@@ -240,21 +127,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "cfg = MODEL_CONFIGS[MODEL_SIZE].copy()\n",
-        "cfg[\"num_classes\"] = NUM_CLASSES; cfg[\"class_drop_prob\"] = 0.1; cfg[\"use_zigzag\"] = True\n",
-        "model = LiquidGen(**cfg).to(device)\n",
-        "print(f\"LiquidGen-{MODEL_SIZE}: {model.count_params()/1e6:.1f}M params\")\n",
-        "\n",
-        "optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n",
-        "total_steps = len(train_loader) * NUM_EPOCHS // GRADIENT_ACCUMULATION\n",
-        "scheduler = cosine_sched(optimizer, WARMUP_STEPS, total_steps)\n",
-        "ema = EMAModel(model, EMA_DECAY)\n",
-        "scaler = GradScaler(\"cuda\")\n",
-        "fm = FlowMatchingScheduler()\n",
-        "os.makedirs(f\"{OUTPUT_DIR}/samples\", exist_ok=True)\n",
-        "os.makedirs(f\"{OUTPUT_DIR}/checkpoints\", exist_ok=True)\n",
-        "print(f\"Total steps: {total_steps}, Effective batch: {BATCH_SIZE*GRADIENT_ACCUMULATION}\")\n",
-        "if torch.cuda.is_available(): print(f\"VRAM: {torch.cuda.memory_allocated()/1024**3:.2f} GB used\")"
+        "cfg = MODEL_CONFIGS[MODEL_SIZE].copy()\ncfg[\"num_classes\"] = NUM_CLASSES; cfg[\"class_drop_prob\"] = 0.1; cfg[\"use_zigzag\"] = True\nmodel = LiquidGen(**cfg).to(device)\nprint(f\"LiquidGen-{MODEL_SIZE}: {model.count_params()/1e6:.1f}M params\")\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\nscheduler = cosine_sched(optimizer, WARMUP_STEPS, MAX_STEPS)\nema = EMAModel(model, EMA_DECAY)\nscaler = GradScaler(\"cuda\")\nfm = FlowMatchingScheduler()\nos.makedirs(f\"{OUTPUT_DIR}/samples\", exist_ok=True)\nos.makedirs(f\"{OUTPUT_DIR}/checkpoints\", exist_ok=True)\nprint(f\"Training: {MAX_STEPS} steps, effective batch {BATCH_SIZE*GRADIENT_ACCUMULATION}\")\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## \ud83d\ude80 Step 7: Train!"
       ]
     },
     {
@@ -263,57 +143,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "global_step = 0; loss_accum = 0.0; log_losses = []\n",
-        "print(\"\ud83d\ude80 Training!\n\")\n",
-        "t0 = time.time()\n",
-        "for epoch in range(NUM_EPOCHS):\n",
-        "    model.train(); ep_loss = 0; ep_steps = 0; ep_t = time.time()\n",
-        "    for bi, (imgs, lbls) in enumerate(train_loader):\n",
-        "        imgs = imgs.to(device)\n",
-        "        lbls = lbls.to(device) if NUM_CLASSES > 0 else None\n",
-        "        with torch.no_grad():\n",
-        "            lats = vae.encode(imgs.half()*2-1).latent_dist.sample()\n",
-        "            lats = ((lats - SHIFT) * SCALE).float()\n",
-        "        t = fm.sample_timesteps(lats.shape[0], device)\n",
-        "        noise = torch.randn_like(lats)\n",
-        "        xt = fm.add_noise(lats, noise, t)\n",
-        "        vtgt = fm.get_velocity_target(lats, noise)\n",
-        "        with autocast(\"cuda\"): loss = F.mse_loss(model(xt, t, lbls), vtgt) / GRADIENT_ACCUMULATION\n",
-        "        scaler.scale(loss).backward()\n",
-        "        loss_accum += loss.item()\n",
-        "        if (bi+1) % GRADIENT_ACCUMULATION == 0:\n",
-        "            scaler.unscale_(optimizer)\n",
-        "            gn = torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)\n",
-        "            scaler.step(optimizer); scaler.update(); optimizer.zero_grad(); scheduler.step()\n",
-        "            ema.update(model); global_step += 1\n",
-        "            if global_step % LOG_EVERY == 0:\n",
-        "                al = loss_accum / LOG_EVERY; lr = optimizer.param_groups[0][\"lr\"]\n",
-        "                vram = torch.cuda.memory_allocated()/1024**3 if torch.cuda.is_available() else 0\n",
-        "                print(f\"step={global_step:>6d} | ep={epoch} | loss={al:.4f} | gn={gn:.2f} | lr={lr:.2e} | vram={vram:.1f}G\")\n",
-        "                log_losses.append(al); loss_accum = 0\n",
-        "                if math.isnan(al) or al > 50: print(\"\ud83d\udca5 Diverged!\"); break\n",
-        "            if global_step % SAMPLE_EVERY == 0:\n",
-        "                ema.apply(model); model.eval()\n",
-        "                ls = IMAGE_SIZE // 8\n",
-        "                sl = torch.randint(0, max(1,NUM_CLASSES), (4,), device=device) if NUM_CLASSES > 0 else None\n",
-        "                samp = fm.sample(model, (4,16,ls,ls), device, NUM_SAMPLE_STEPS, sl, CFG_SCALE)\n",
-        "                with torch.no_grad(): si = ((vae.decode(samp.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n",
-        "                save_image(si, f\"{OUTPUT_DIR}/samples/step_{global_step:07d}.png\", nrow=2)\n",
-        "                ema.restore(model); model.train()\n",
-        "            if global_step % SAVE_EVERY == 0:\n",
-        "                torch.save({\"model\":model.state_dict(),\"ema\":ema.shadow,\"step\":global_step,\"cfg\":cfg},\n",
-        "                    f\"{OUTPUT_DIR}/checkpoints/step_{global_step:07d}.pt\")\n",
-        "        ep_loss += loss.item()*GRADIENT_ACCUMULATION; ep_steps += 1\n",
-        "    print(f\"Epoch {epoch} | loss={ep_loss/max(ep_steps,1):.4f} | {time.time()-ep_t:.0f}s\")\n",
-        "torch.save({\"model\":model.state_dict(),\"ema\":ema.shadow,\"cfg\":cfg,\"step\":global_step},f\"{OUTPUT_DIR}/checkpoints/final.pt\")\n",
-        "print(f\"\ud83c\udf89 Done! {global_step} steps in {(time.time()-t0)/60:.1f} min\")"
+        "global_step = 0; loss_accum = 0.0; log_losses = []; accum_count = 0\nprint(\"Training started!\\n\")\nt0 = time.time(); model.train()\n\nwhile global_step < MAX_STEPS:\n    for imgs, lbls in train_loader:\n        if global_step >= MAX_STEPS: break\n        imgs = imgs.to(device)\n        lbls = lbls.to(device) if NUM_CLASSES > 0 else None\n\n        with torch.no_grad():\n            lats = vae.encode(imgs.half()*2-1).latent_dist.sample()\n            lats = ((lats - SHIFT) * SCALE).float()\n\n        t = fm.sample_t(lats.shape[0], device)\n        noise = torch.randn_like(lats)\n        xt = fm.add_noise(lats, noise, t)\n        vtgt = fm.velocity(lats, noise)\n\n        with autocast(\"cuda\"):\n            loss = F.mse_loss(model(xt, t, lbls), vtgt) / GRADIENT_ACCUMULATION\n        scaler.scale(loss).backward()\n        loss_accum += loss.item()\n        accum_count += 1\n\n        if accum_count % GRADIENT_ACCUMULATION == 0:\n            scaler.unscale_(optimizer)\n            gn = torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)\n            scaler.step(optimizer); scaler.update(); optimizer.zero_grad()\n            scheduler.step(); ema.update(model); global_step += 1\n\n            if global_step % LOG_EVERY == 0:\n                al = loss_accum / LOG_EVERY; lr = optimizer.param_groups[0][\"lr\"]\n                vram = torch.cuda.memory_allocated()/1024**3 if torch.cuda.is_available() else 0\n                print(f\"step={global_step:>6d} | loss={al:.4f} | gn={gn:.2f} | lr={lr:.2e} | vram={vram:.1f}G | {time.time()-t0:.0f}s\")\n                log_losses.append(al); loss_accum = 0\n                if math.isnan(al) or al > 50: print(\"Diverged!\"); break\n\n            if global_step % SAMPLE_EVERY == 0:\n                ema.apply(model); model.eval()\n                ls = IMAGE_SIZE // 8\n                sl = torch.randint(0, max(1,NUM_CLASSES), (4,), device=device) if NUM_CLASSES > 0 else None\n                samp = fm.sample(model, (4,16,ls,ls), device, NUM_SAMPLE_STEPS, sl, CFG_SCALE)\n                with torch.no_grad():\n                    si = ((vae.decode(samp.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n                save_image(si, f\"{OUTPUT_DIR}/samples/step_{global_step:07d}.png\", nrow=2)\n                print(f\"  Saved samples\")\n                ema.restore(model); model.train()\n\n            if global_step % SAVE_EVERY == 0:\n                torch.save({\"model\":model.state_dict(),\"ema\":ema.shadow,\"step\":global_step,\"cfg\":cfg},\n                    f\"{OUTPUT_DIR}/checkpoints/step_{global_step:07d}.pt\")\n                print(f\"  Checkpoint saved\")\n\ntorch.save({\"model\":model.state_dict(),\"ema\":ema.shadow,\"cfg\":cfg,\"step\":global_step},\n    f\"{OUTPUT_DIR}/checkpoints/final.pt\")\nprint(f\"\\nDone! {global_step} steps in {(time.time()-t0)/60:.1f} min\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udcc8 Training Loss"
+        "## \ud83d\udcc8 Step 8: Loss Curve"
       ]
     },
     {
@@ -322,18 +159,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import matplotlib.pyplot as plt\n",
-        "if log_losses:\n",
-        "    plt.figure(figsize=(10,4)); plt.plot(log_losses); plt.xlabel(f\"Steps (\u00d7{LOG_EVERY})\"); plt.ylabel(\"Loss\")\n",
-        "    plt.title(\"Training Loss\"); plt.grid(True, alpha=0.3); plt.savefig(f\"{OUTPUT_DIR}/loss.png\", dpi=150); plt.show()\n",
-        "    print(f\"Min loss: {min(log_losses):.4f}\")"
+        "import matplotlib.pyplot as plt\nif log_losses:\n    plt.figure(figsize=(10,4)); plt.plot(log_losses)\n    plt.xlabel(f\"Steps (x{LOG_EVERY})\"); plt.ylabel(\"Loss\")\n    plt.title(\"Training Loss\"); plt.grid(True, alpha=0.3)\n    plt.savefig(f\"{OUTPUT_DIR}/loss.png\", dpi=150); plt.show()\n    print(f\"Min loss: {min(log_losses):.4f}\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83c\udfa8 Generate Images"
+        "## \ud83c\udfa8 Step 9: Generate"
       ]
     },
     {
@@ -342,32 +175,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "ema.apply(model); model.eval()\n",
-        "N, STEPS, CFG = 8, 50, 2.5\n",
-        "ls = IMAGE_SIZE // 8\n",
-        "STYLES = [\"Abstract Expressionism\",\"Baroque\",\"Cubism\",\"Expressionism\",\"Impressionism\",\n",
-        "          \"Pop Art\",\"Realism\",\"Romanticism\",\"Symbolism\",\"Ukiyo-e\"]\n",
-        "if NUM_CLASSES > 0:\n",
-        "    for ci in range(min(NUM_CLASSES, 8)):\n",
-        "        l = torch.full((N,), ci, device=device, dtype=torch.long)\n",
-        "        s = fm.sample(model, (N,16,ls,ls), device, STEPS, l, CFG)\n",
-        "        with torch.no_grad(): i = ((vae.decode(s.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n",
-        "        nm = STYLES[ci] if ci < len(STYLES) else f\"Class_{ci}\"\n",
-        "        save_image(i, f\"{OUTPUT_DIR}/gen_{nm.replace(chr(32),chr(95))}.png\", nrow=4)\n",
-        "        print(f\"Generated: {nm}\")\n",
-        "else:\n",
-        "    s = fm.sample(model, (N,16,ls,ls), device, STEPS)\n",
-        "    with torch.no_grad(): i = ((vae.decode(s.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n",
-        "    save_image(i, f\"{OUTPUT_DIR}/gen_uncond.png\", nrow=4)\n",
-        "ema.restore(model)\n",
-        "print(f\"\u2705 Saved to {OUTPUT_DIR}/\")"
+        "ema.apply(model); model.eval()\nN, STEPS, G = 8, 50, 2.5; ls = IMAGE_SIZE // 8\nif NUM_CLASSES > 0:\n    for ci in range(min(NUM_CLASSES, 6)):\n        l = torch.full((N,), ci, device=device, dtype=torch.long)\n        s = fm.sample(model, (N,16,ls,ls), device, STEPS, l, G)\n        with torch.no_grad(): i = ((vae.decode(s.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n        save_image(i, f\"{OUTPUT_DIR}/gen_class{ci}.png\", nrow=4)\n        print(f\"Class {ci}\")\nelse:\n    s = fm.sample(model, (N,16,ls,ls), device, STEPS)\n    with torch.no_grad(): i = ((vae.decode(s.half()/SCALE+SHIFT).sample+1)/2).clamp(0,1).float()\n    save_image(i, f\"{OUTPUT_DIR}/gen_uncond.png\", nrow=4)\nema.restore(model)\nprint(f\"Saved to {OUTPUT_DIR}/\")\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## \ud83d\udce4 Display Results"
+        "## \ud83d\udce4 Step 10: Display"
       ]
     },
     {
@@ -376,35 +191,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from IPython.display import display\n",
-        "import glob\n",
-        "for f in sorted(glob.glob(f\"{OUTPUT_DIR}/samples/*.png\"))[-3:]:\n",
-        "    print(os.path.basename(f)); display(Image.open(f))\n",
-        "for f in sorted(glob.glob(f\"{OUTPUT_DIR}/gen_*.png\")):\n",
-        "    print(os.path.basename(f)); display(Image.open(f))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## \ud83d\udcdd Architecture Reference\n",
-        "\n",
-        "### Core Equation (CfC Liquid Dynamics)\n",
-        "\n",
-        "\n",
-        "### Flow Matching\n",
-        "\n",
-        "\n",
-        "### Sampling (Euler ODE)\n",
-        "\n",
-        "\n",
-        "### References\n",
-        "- Hasani et al., \"Liquid Time-constant Networks\" (NeurIPS 2020)\n",
-        "- Hasani et al., \"Closed-form Continuous-depth Models\" (Nature MI 2022)\n",
-        "- Lechner et al., \"Neural Circuit Policies\" (Nature MI 2020)\n",
-        "- ZigMa (ECCV 2024), DiMSUM (NeurIPS 2024)\n",
-        "- Lipman et al., \"Flow Matching\" (2023), SiT (2024)\n"
+        "from IPython.display import display\nimport glob\nfor f in sorted(glob.glob(f\"{OUTPUT_DIR}/samples/*.png\"))[-3:]:\n    print(os.path.basename(f)); display(Image.open(f))\nfor f in sorted(glob.glob(f\"{OUTPUT_DIR}/gen_*.png\")):\n    print(os.path.basename(f)); display(Image.open(f))\n"
       ]
     }
   ]