Upload notebook.ipynb with huggingface_hub

notebook.ipynb

@@ -0,0 +1,563 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🔬 LatentRecurrentFlow (LRF) — A Novel Mobile-First Image Generation Architecture\n",
+    "\n",
+    "**A complete implementation of a novel image generation architecture designed for consumer devices.**\n",
+    "\n",
+    "## Key Innovations\n",
+    "\n",
+    "1. **Recursive Latent Refinement (RLR) Core** — HRM-inspired iterative reasoning on image latents with O(1) memory backpropagation\n",
+    "2. **Gated Linear Diffusion (GLD) Blocks** — O(N) subquadratic spatial mixing replacing quadratic self-attention\n",
+    "3. **Compact f=16 VAE** with SnapGen-inspired tiny decoder (1-2M params)\n",
+    "4. **Rectified Flow** training with consistency distillation readiness\n",
+    "5. **Editing-ready architecture** — same latent core supports text-to-image, inpainting, style editing, and more\n",
+    "\n",
+    "### Memory Budget\n",
+    "| Component | FP32 | INT8 (Mobile) |\n",
+    "|-----------|------|---------------|\n",
+    "| VAE Decoder | 4 MB | 1 MB |\n",
+    "| Text Encoder | 44 MB | 11 MB |\n",
+    "| Denoising Core | 2.5 MB | 0.6 MB |\n",
+    "| Activations (256²) | ~200 MB | ~100 MB |\n",
+    "| **Total** | **~250 MB** | **~113 MB** |\n",
+    "\n",
+    "This notebook demonstrates:\n",
+    "1. Architecture design and parameter counting\n",
+    "2. End-to-end VAE training\n",
+    "3. Flow matching denoiser training\n",
+    "4. Sample generation\n",
+    "5. Model saving and loading"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 0. Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install dependencies\n",
+    "!pip install -q torch torchvision einops safetensors huggingface_hub pillow matplotlib"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Clone the LRF repo (if not already available)\n",
+    "import os, sys\n",
+    "\n",
+    "# If running from the repo, just add to path\n",
+    "if os.path.exists('lrf'):\n",
+    "    sys.path.insert(0, '.')\n",
+    "else:\n",
+    "    # Clone from HF Hub\n",
+    "    !git clone https://huggingface.co/krystv/LatentRecurrentFlow\n",
+    "    sys.path.insert(0, 'LatentRecurrentFlow')\n",
+    "\n",
+    "from lrf.model import LatentRecurrentFlow, RecursiveLatentCore, CompactVAE, GatedLinearAttention\n",
+    "from lrf.training import LRFTrainer, RectifiedFlowScheduler, SyntheticImageTextDataset\n",
+    "from lrf.pipeline import LRFPipeline, LRFTrainingPipeline\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "# Device\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "print(f'Using device: {device}')\n",
+    "if device.type == 'cuda':\n",
+    "    print(f'GPU: {torch.cuda.get_device_name()}')\n",
+    "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Architecture Overview & Parameter Counting"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create model with different configs\n",
+    "configs = {\n",
+    "    'Tiny (5.7M)': LatentRecurrentFlow.tiny_config(),\n",
+    "    'Default (16.3M)': LatentRecurrentFlow.default_config(),\n",
+    "}\n",
+    "\n",
+    "for name, config in configs.items():\n",
+    "    model = LatentRecurrentFlow(config)\n",
+    "    counts = model.count_parameters()\n",
+    "    \n",
+    "    print(f'\\n=== {name} ===')\n",
+    "    print(f'Config: T_outer={config[\"T_outer\"]}, T_inner={config[\"T_inner\"]}, '\n",
+    "          f'num_blocks={config[\"num_blocks\"]}')\n",
+    "    print(f'Effective depth: {config[\"T_outer\"] * config[\"T_inner\"] * config[\"num_blocks\"]} layers '\n",
+    "          f'(from {config[\"num_blocks\"]} unique blocks)')\n",
+    "    for module, count in counts.items():\n",
+    "        mb = count * 4 / 1e6\n",
+    "        print(f'  {module:20s}: {count:>12,} params ({mb:.1f} MB FP32)')\n",
+    "    del model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Stage 1: VAE Training\n",
+    "\n",
+    "The VAE learns to compress images into a compact latent space.\n",
+    "- f=16 spatial compression: 256×256 → 16×16 latents\n",
+    "- C=16 or C=32 latent channels\n",
+    "- Tiny decoder (~280K params) inspired by SnapGen"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create model for training\n",
+    "config = LatentRecurrentFlow.tiny_config()\n",
+    "model = LatentRecurrentFlow(config).to(device)\n",
+    "\n",
+    "# Create synthetic dataset (replace with real data for actual training)\n",
+    "dataset = SyntheticImageTextDataset(\n",
+    "    num_samples=500,\n",
+    "    image_size=64,\n",
+    "    max_text_length=32\n",
+    ")\n",
+    "dataloader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=0)\n",
+    "\n",
+    "# Create trainer\n",
+    "trainer = LRFTrainer(model, device, './lrf_checkpoints')\n",
+    "\n",
+    "print(f'Dataset size: {len(dataset)}')\n",
+    "print(f'Batch size: 8')\n",
+    "print(f'Image size: 64x64')\n",
+    "print(f'Latent size: {64//16}x{64//16}x{config[\"latent_channels\"]}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Train VAE\n",
+    "vae_optimizer = torch.optim.AdamW(model.vae.parameters(), lr=1e-3, weight_decay=0.01)\n",
+    "\n",
+    "vae_losses = []\n",
+    "num_vae_steps = 100\n",
+    "\n",
+    "print('Training VAE...')\n",
+    "step = 0\n",
+    "for epoch in range(10):  # Multiple epochs over small dataset\n",
+    "    for batch in dataloader:\n",
+    "        if step >= num_vae_steps:\n",
+    "            break\n",
+    "        losses = trainer.train_vae_step(batch['image'], vae_optimizer)\n",
+    "        vae_losses.append(losses['total'])\n",
+    "        if step % 20 == 0:\n",
+    "            print(f'  Step {step}: total={losses[\"total\"]:.4f}, '\n",
+    "                  f'recon={losses[\"recon\"]:.4f}, kl={losses[\"kl\"]:.4f}')\n",
+    "        step += 1\n",
+    "    if step >= num_vae_steps:\n",
+    "        break\n",
+    "\n",
+    "# Plot VAE loss\n",
+    "plt.figure(figsize=(10, 4))\n",
+    "plt.plot(vae_losses)\n",
+    "plt.xlabel('Step')\n",
+    "plt.ylabel('Loss')\n",
+    "plt.title('VAE Training Loss')\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.show()\n",
+    "\n",
+    "# Save checkpoint\n",
+    "trainer.save_checkpoint('./lrf_checkpoints/vae.pt', 'vae', 0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Visualize VAE reconstruction\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    sample_batch = next(iter(dataloader))\n",
+    "    images = sample_batch['image'].to(device)\n",
+    "    recon, _, _ = model.vae(images)\n",
+    "\n",
+    "fig, axes = plt.subplots(2, 4, figsize=(16, 8))\n",
+    "for i in range(4):\n",
+    "    # Original\n",
+    "    img = images[i].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+    "    axes[0][i].imshow(np.clip(img, 0, 1))\n",
+    "    axes[0][i].set_title(f'Original {i}')\n",
+    "    axes[0][i].axis('off')\n",
+    "    \n",
+    "    # Reconstruction\n",
+    "    rec = recon[i].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+    "    axes[1][i].imshow(np.clip(rec, 0, 1))\n",
+    "    axes[1][i].set_title(f'Reconstruction {i}')\n",
+    "    axes[1][i].axis('off')\n",
+    "\n",
+    "plt.suptitle('VAE Reconstruction Quality', fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Stage 2: Flow Matching Denoiser Training\n",
+    "\n",
+    "The denoising core learns to predict the velocity field for rectified flow.\n",
+    "- VAE is frozen\n",
+    "- Core + text encoder are trained\n",
+    "- Uses SNR-weighted flow matching loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Freeze VAE\n",
+    "for p in model.vae.parameters():\n",
+    "    p.requires_grad = False\n",
+    "\n",
+    "# Train flow matching\n",
+    "flow_params = list(model.core.parameters()) + list(model.text_encoder.parameters())\n",
+    "flow_optimizer = torch.optim.AdamW(flow_params, lr=1e-3, weight_decay=0.01)\n",
+    "\n",
+    "flow_losses = []\n",
+    "num_flow_steps = 100\n",
+    "\n",
+    "print('Training flow matching denoiser...')\n",
+    "model.core.train()\n",
+    "model.text_encoder.train()\n",
+    "\n",
+    "step = 0\n",
+    "for epoch in range(10):\n",
+    "    for batch in dataloader:\n",
+    "        if step >= num_flow_steps:\n",
+    "            break\n",
+    "        losses = trainer.train_flow_step(\n",
+    "            batch['image'], batch['token_ids'], batch['attention_mask'],\n",
+    "            flow_optimizer, cfg_dropout=0.1\n",
+    "        )\n",
+    "        flow_losses.append(losses['flow_loss'])\n",
+    "        if step % 20 == 0:\n",
+    "            print(f'  Step {step}: flow_loss={losses[\"flow_loss\"]:.4f}')\n",
+    "        step += 1\n",
+    "    if step >= num_flow_steps:\n",
+    "        break\n",
+    "\n",
+    "# Plot flow loss\n",
+    "plt.figure(figsize=(10, 4))\n",
+    "plt.plot(flow_losses)\n",
+    "plt.xlabel('Step')\n",
+    "plt.ylabel('Flow Matching Loss')\n",
+    "plt.title('Denoiser Training Loss')\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.show()\n",
+    "\n",
+    "# Save checkpoint\n",
+    "trainer.save_checkpoint('./lrf_checkpoints/flow.pt', 'flow', 0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Generation & Visualization\n",
+    "\n",
+    "Generate images using the trained model with Euler ODE sampling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Generate samples\n",
+    "model.eval()\n",
+    "\n",
+    "# Create prompts (using simple tokenization for prototype)\n",
+    "prompts = [\n",
+    "    'a beautiful sunset over the ocean with golden light',\n",
+    "    'a cute cat sitting on a windowsill',\n",
+    "    'a mountain landscape with snow and trees',\n",
+    "    'a colorful abstract painting with swirls',\n",
+    "]\n",
+    "\n",
+    "pipe = LRFPipeline(model, device=device)\n",
+    "\n",
+    "# Generate with different step counts\n",
+    "for num_steps in [5, 10, 20]:\n",
+    "    images = pipe(\n",
+    "        prompts,\n",
+    "        num_steps=num_steps,\n",
+    "        cfg_scale=1.0,  # Low cfg for untrained model\n",
+    "        height=64,\n",
+    "        width=64,\n",
+    "        seed=42,\n",
+    "    )\n",
+    "    \n",
+    "    fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
+    "    for i in range(4):\n",
+    "        img = images[i].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n",
+    "        axes[i].imshow(np.clip(img, 0, 1))\n",
+    "        axes[i].set_title(prompts[i][:30] + '...')\n",
+    "        axes[i].axis('off')\n",
+    "    plt.suptitle(f'Generated Images ({num_steps} steps)', fontsize=14)\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Save & Load Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save the complete model\n",
+    "pipe.save_pretrained('./lrf_model')\n",
+    "print('Model saved to ./lrf_model/')\n",
+    "\n",
+    "# List saved files\n",
+    "for f in os.listdir('./lrf_model'):\n",
+    "    size = os.path.getsize(f'./lrf_model/{f}')\n",
+    "    print(f'  {f}: {size/1024:.1f} KB')\n",
+    "\n",
+    "# Reload and verify\n",
+    "pipe_loaded = LRFPipeline.from_pretrained('./lrf_model', device=str(device))\n",
+    "images_loaded = pipe_loaded('test prompt', num_steps=5, height=64, width=64, seed=42)\n",
+    "print(f'\\nReloaded model generates: {images_loaded.shape}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Training Curriculum for Real Data\n",
+    "\n",
+    "The full training curriculum for production-quality models:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display the full training curriculum\n",
+    "curriculum = LRFTrainingPipeline.get_curriculum()\n",
+    "\n",
+    "print('Full Training Curriculum')\n",
+    "print('=' * 70)\n",
+    "for i, stage_name in enumerate(curriculum):\n",
+    "    stage = LRFTrainingPipeline.get_stage_config(stage_name)\n",
+    "    print(f'\\nStage {i+1}: {stage_name}')\n",
+    "    print(f'  Description: {stage[\"description\"]}')\n",
+    "    print(f'  Freeze: {stage[\"freeze\"]}')\n",
+    "    print(f'  Train:  {stage[\"train\"]}')\n",
+    "    print(f'  LR:     {stage[\"lr\"]}')\n",
+    "    print(f'  Min steps: {stage[\"min_steps\"]:,}')\n",
+    "    if 'resolution' in stage:\n",
+    "        print(f'  Resolution: {stage[\"resolution\"]}×{stage[\"resolution\"]}')\n",
+    "\n",
+    "print('\\n' + '=' * 70)\n",
+    "print('\\nRecommended datasets for each stage:')\n",
+    "print('  Stage 1 (VAE):       ImageNet, COCO, or any large image dataset')\n",
+    "print('  Stage 2 (Flow 64):   Synthetic captions from teacher (SDXL/SD3) + LAION-aesthetic')\n",
+    "print('  Stage 3 (Flow 256):  Filtered LAION-aesthetic (score > 6.0) + synthetic')\n",
+    "print('  Stage 4 (Flow 512):  High-quality curated dataset + JourneyDB')\n",
+    "print('  Stage 5 (Distill):   Same as Stage 4 (distill from own multi-step model)')\n",
+    "print('  Stage 6 (Editing):   InstructPix2Pix + MagicBrush + synthetic edit pairs')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Architecture Deep Dive\n",
+    "\n",
+    "### The Recursive Latent Refinement Loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Demonstrate the recursive refinement\n",
+    "core = RecursiveLatentCore(\n",
+    "    dim=32, cond_dim=64, num_blocks=2, num_heads=2, head_dim=16,\n",
+    "    T_inner=4, T_outer=2, use_ift_training=False\n",
+    ")\n",
+    "\n",
+    "print('Recursive Latent Core Architecture')\n",
+    "print('=' * 50)\n",
+    "print(f'Unique GLD blocks: {core.num_blocks}')\n",
+    "print(f'T_outer (abstract updates): {core.T_outer}')\n",
+    "print(f'T_inner (refinement steps): {core.T_inner}')\n",
+    "print(f'Total recursions: {core.T_outer * core.T_inner}')\n",
+    "print(f'Effective depth: {core.T_outer * core.T_inner * core.num_blocks} layers')\n",
+    "print(f'Parameter reuse factor: {core.T_outer * core.T_inner}x')\n",
+    "print(f'\\nParameters: {sum(p.numel() for p in core.parameters()):,}')\n",
+    "\n",
+    "# Show memory savings from IFT\n",
+    "print('\\nMemory comparison:')\n",
+    "eff_depth = core.T_outer * core.T_inner * core.num_blocks\n",
+    "print(f'  Standard backprop: O({eff_depth}) activation memory')\n",
+    "print(f'  IFT backprop:      O(1) activation memory')\n",
+    "print(f'  Memory savings:    {eff_depth}x')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Demonstrate GLA complexity\n",
+    "import time\n",
+    "\n",
+    "gla = GatedLinearAttention(dim=64, num_heads=4, head_dim=16)\n",
+    "\n",
+    "print('GLA Complexity Scaling')\n",
+    "print('=' * 50)\n",
+    "\n",
+    "sizes = [4, 8, 16, 32, 64]\n",
+    "times = []\n",
+    "\n",
+    "for s in sizes:\n",
+    "    x = torch.randn(1, s*s, 64)\n",
+    "    \n",
+    "    # Warmup\n",
+    "    _ = gla(x, h=s, w=s)\n",
+    "    \n",
+    "    # Time\n",
+    "    t0 = time.time()\n",
+    "    for _ in range(10):\n",
+    "        _ = gla(x, h=s, w=s)\n",
+    "    dt = (time.time() - t0) / 10\n",
+    "    times.append(dt)\n",
+    "    print(f'  {s}×{s} = {s*s:>5} tokens: {dt*1000:.2f}ms')\n",
+    "\n",
+    "# Plot scaling\n",
+    "plt.figure(figsize=(8, 4))\n",
+    "tokens = [s*s for s in sizes]\n",
+    "plt.plot(tokens, [t*1000 for t in times], 'bo-', label='GLA (O(N))')\n",
+    "# Reference quadratic line\n",
+    "t_ref = times[0] * 1000\n",
+    "quadratic = [t_ref * (n / tokens[0])**2 for n in tokens]\n",
+    "plt.plot(tokens, quadratic, 'r--', label='Quadratic attention (O(N²))', alpha=0.5)\n",
+    "plt.xlabel('Number of tokens')\n",
+    "plt.ylabel('Time (ms)')\n",
+    "plt.title('GLA vs Quadratic Attention Scaling')\n",
+    "plt.legend()\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Push to HuggingFace Hub (Optional)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Uncomment to push to Hub\n",
+    "# from huggingface_hub import HfApi\n",
+    "# api = HfApi()\n",
+    "# api.upload_folder(\n",
+    "#     folder_path='./lrf_model',\n",
+    "#     repo_id='your-username/LatentRecurrentFlow',\n",
+    "#     repo_type='model',\n",
+    "# )\n",
+    "print('To push to HF Hub, uncomment the code above and set your repo_id.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "\n",
+    "## Summary\n",
+    "\n",
+    "This notebook demonstrated the LatentRecurrentFlow architecture end-to-end:\n",
+    "\n",
+    "1. ✅ Model creation with parameter counting\n",
+    "2. ✅ VAE training for image compression\n",
+    "3. ✅ Flow matching denoiser training\n",
+    "4. ✅ Image generation with Euler ODE sampling\n",
+    "5. ✅ Model save/load with HF-compatible format\n",
+    "6. ✅ Training curriculum for production\n",
+    "\n",
+    "### Next Steps\n",
+    "- Replace synthetic data with real image-text pairs\n",
+    "- Scale to default config (16M params)\n",
+    "- Train on GPU for actual quality\n",
+    "- Add consistency distillation for 4-step generation\n",
+    "- Add editing fine-tuning stage"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}