Add complete Colab/Kaggle training notebook

IRIS_Training_Notebook.ipynb

@@ -0,0 +1,956 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# \ud83d\udd2e IRIS: Iterative Recurrent Image Synthesis \u2014 Training Notebook",
+    "",
+    "**Train a novel mobile-first image generation model from scratch on free Colab/Kaggle GPUs.**",
+    "",
+    "This notebook runs the complete 2-stage training pipeline:",
+    "1. **Stage 1 \u2014 Wavelet VAE Training**: Learn to encode/decode images via wavelet-frequency latent space",
+    "2. **Stage 2 \u2014 Generator Training**: Train the recurrent-depth denoiser with rectified flow on captioned images",
+    "",
+    "### Hardware Requirements",
+    "| Platform | GPU | VRAM | Estimated Time |",
+    "|----------|-----|------|----------------|",
+    "| **Colab Free** | T4 | 16GB | ~2-3 hours total |",
+    "| **Colab Pro** | A100 | 40GB | ~45 min total |",
+    "| **Kaggle** | P100/T4\u00d72 | 16GB | ~2-3 hours total |",
+    "",
+    "### What You Get",
+    "- A trained Wavelet VAE that compresses 256\u00d7256 images to 16\u00d716 latent (48\u00d7 compression)",
+    "- A trained IRIS generator that can denoise latents conditioned on text (CLIP embeddings)",
+    "- Visualization of reconstructions, generation samples, and loss curves",
+    "- Saved checkpoints you can continue training from"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Setup & Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install dependencies\n",
+    "!pip install -q torch torchvision datasets transformers accelerate matplotlib Pillow tqdm huggingface_hub\n",
+    "\n",
+    "# Check GPU\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 / 1024**3:.1f} GB\")\n",
+    "    device = torch.device('cuda')\n",
+    "else:\n",
+    "    print(\"\u26a0\ufe0f  No GPU detected! Training will be very slow on CPU.\")\n",
+    "    device = torch.device('cpu')"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Download IRIS Architecture from Hugging Face"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Download the IRIS architecture code from HF Hub\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "import shutil, os\n",
+    "\n",
+    "repo_id = \"asdf98/IRIS-architecture\"\n",
+    "for fname in [\"iris_model.py\", \"train_iris.py\", \"test_iris.py\"]:\n",
+    "    path = hf_hub_download(repo_id=repo_id, filename=fname)\n",
+    "    shutil.copy(path, f\"./{fname}\")\n",
+    "    print(f\"\u2705 Downloaded {fname}\")\n",
+    "\n",
+    "# Import IRIS\n",
+    "from iris_model import (\n",
+    "    IRIS, IRISConfig, WaveletVAE, IRISGenerator,\n",
+    "    HaarDWT2D, HaarIDWT2D,\n",
+    "    create_iris_small, create_iris_tiny, create_iris_base,\n",
+    "    count_parameters, estimate_memory_mb,\n",
+    ")\n",
+    "print(\"\\n\u2705 IRIS architecture imported successfully!\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Model Architecture Overview",
+    "",
+    "Let's inspect the three model variants and their parameter counts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Show model variants\n",
+    "for name, fn in [(\"IRIS-Tiny (ultra-mobile)\", create_iris_tiny),\n",
+    "                 (\"IRIS-Small (mobile)\", create_iris_small),\n",
+    "                 (\"IRIS-Base (desktop)\", create_iris_base)]:\n",
+    "    model = fn()\n",
+    "    counts = count_parameters(model)\n",
+    "    mem16 = estimate_memory_mb(model, torch.float16)\n",
+    "\n",
+    "    core_params = sum(p.numel() for p in model.generator.core.parameters())\n",
+    "    print(f\"\\n{'='*55}\")\n",
+    "    print(f\"  {name}\")\n",
+    "    print(f\"{'='*55}\")\n",
+    "    print(f\"  Total params:       {counts['total']:>12,}\")\n",
+    "    print(f\"  Generator params:   {counts['total'] - sum(p.numel() for p in model.vae.parameters()):>12,}\")\n",
+    "    print(f\"  Core (shared):      {core_params:>12,}\")\n",
+    "    print(f\"  Model memory fp16:  {mem16:>10.1f} MB\")\n",
+    "    print(f\"  + CLIP-L/14 text:        156.0 MB\")\n",
+    "    print(f\"  + Overhead:              350.0 MB\")\n",
+    "    print(f\"  \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\")\n",
+    "    print(f\"  Total inference:    {mem16+156+350:>10.1f} MB  {'\u2705 <3GB' if mem16+506 < 3000 else ''}\")\n",
+    "\n",
+    "del model  # Free memory"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Load Dataset \u2014 Pok\u00e9mon BLIP Captions",
+    "",
+    "We use `reach-vb/pokemon-blip-captions` \u2014 a small, high-quality dataset with ~833 image-caption pairs. ",
+    "Perfect for free-tier training to validate the architecture works end-to-end.",
+    "",
+    "**For serious training later**, swap in larger datasets:",
+    "- `ILSVRC/imagenet-1k` (Stage 2 class-conditional)",
+    "- `laion/laion-art` (Text-image alignment)",
+    "- `caidas/JourneyDB` (Aesthetic fine-tuning)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from datasets import load_dataset\n",
+    "from torchvision import transforms\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "\n",
+    "# Load Pok\u00e9mon dataset\n",
+    "print(\"Loading dataset...\")\n",
+    "raw_dataset = load_dataset(\"reach-vb/pokemon-blip-captions\", split=\"train\")\n",
+    "print(f\"\u2705 Loaded {len(raw_dataset)} image-caption pairs\")\n",
+    "\n",
+    "# Show a few examples\n",
+    "import matplotlib.pyplot as plt\n",
+    "fig, axes = plt.subplots(1, 5, figsize=(20, 4))\n",
+    "for i, ax in enumerate(axes):\n",
+    "    item = raw_dataset[i]\n",
+    "    ax.imshow(item[\"image\"])\n",
+    "    ax.set_title(item[\"text\"][:40] + \"...\", fontsize=9)\n",
+    "    ax.axis(\"off\")\n",
+    "plt.suptitle(\"Sample Training Images\", fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.1 Create PyTorch Dataset with Transforms"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2500\u2500\u2500 Training configuration \u2500\u2500\u2500\n",
+    "IMAGE_SIZE = 256      # Input image resolution\n",
+    "BATCH_SIZE = 4        # Fits on T4 16GB; increase on A100\n",
+    "NUM_WORKERS = 2       # Dataloader workers\n",
+    "\n",
+    "# \u2500\u2500\u2500 Image transforms \u2500\u2500\u2500\n",
+    "train_transform = transforms.Compose([\n",
+    "    transforms.Resize(IMAGE_SIZE, interpolation=transforms.InterpolationMode.LANCZOS),\n",
+    "    transforms.CenterCrop(IMAGE_SIZE),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
+    "    transforms.ToTensor(),             # [0, 1]\n",
+    "    transforms.Normalize([0.5]*3, [0.5]*3),  # [-1, 1]\n",
+    "])\n",
+    "\n",
+    "class ImageCaptionDataset(Dataset):\n",
+    "    \"\"\"Wraps a HF dataset with transforms. Returns (image_tensor, caption_string).\"\"\"\n",
+    "    def __init__(self, hf_dataset, transform):\n",
+    "        self.dataset = hf_dataset\n",
+    "        self.transform = transform\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.dataset)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        item = self.dataset[idx]\n",
+    "        image = item[\"image\"].convert(\"RGB\")\n",
+    "        image = self.transform(image)\n",
+    "        caption = item[\"text\"]\n",
+    "        return image, caption\n",
+    "\n",
+    "train_dataset = ImageCaptionDataset(raw_dataset, train_transform)\n",
+    "train_loader = DataLoader(\n",
+    "    train_dataset, batch_size=BATCH_SIZE, shuffle=True,\n",
+    "    num_workers=NUM_WORKERS, pin_memory=True, drop_last=True,\n",
+    ")\n",
+    "print(f\"\u2705 DataLoader ready: {len(train_loader)} batches of {BATCH_SIZE}\")\n",
+    "\n",
+    "# Quick sanity check\n",
+    "imgs, caps = next(iter(train_loader))\n",
+    "print(f\"   Image batch: {imgs.shape}, range [{imgs.min():.2f}, {imgs.max():.2f}]\")\n",
+    "print(f\"   Caption[0]:  {caps[0]}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Load CLIP Text Encoder (Frozen)",
+    "",
+    "We use CLIP-L/14 (~150MB) as the text encoder. It's frozen during training \u2014 ",
+    "only the IRIS generator learns. This is the same encoder used in aMUSEd, Meissonic, and SnapGen."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from transformers import CLIPTextModel, CLIPTokenizer\n",
+    "\n",
+    "print(\"Loading CLIP-L/14 text encoder...\")\n",
+    "clip_model_name = \"openai/clip-vit-large-patch14\"\n",
+    "tokenizer = CLIPTokenizer.from_pretrained(clip_model_name)\n",
+    "text_encoder = CLIPTextModel.from_pretrained(clip_model_name).to(device).eval()\n",
+    "\n",
+    "# Freeze text encoder\n",
+    "for p in text_encoder.parameters():\n",
+    "    p.requires_grad = False\n",
+    "\n",
+    "clip_params = sum(p.numel() for p in text_encoder.parameters())\n",
+    "print(f\"\u2705 CLIP-L/14 loaded: {clip_params/1e6:.1f}M params (frozen)\")\n",
+    "print(f\"   Text embedding dim: {text_encoder.config.hidden_size}\")\n",
+    "print(f\"   Max tokens: {tokenizer.model_max_length}\")\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def encode_text(captions, max_length=77):\n",
+    "    \"\"\"Encode a list of captions to CLIP text embeddings.\"\"\"\n",
+    "    tokens = tokenizer(\n",
+    "        captions, padding=\"max_length\", truncation=True,\n",
+    "        max_length=max_length, return_tensors=\"pt\"\n",
+    "    ).to(device)\n",
+    "    outputs = text_encoder(**tokens)\n",
+    "    return outputs.last_hidden_state  # [B, 77, 768]\n",
+    "\n",
+    "# Test encoding\n",
+    "test_emb = encode_text([\"a cute dragon breathing fire\"])\n",
+    "print(f\"   Test encoding shape: {test_emb.shape}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Stage 1 \u2014 Wavelet VAE Training",
+    "",
+    "Train the lightweight Wavelet VAE to reconstruct images through the wavelet-frequency latent space.",
+    "",
+    "**Architecture**: `Image \u2192 HaarDWT \u2192 Encoder \u2192 Latent(16ch, 16\u00d716) \u2192 Decoder \u2192 HaarIDWT \u2192 Image`",
+    "",
+    "**Losses**:",
+    "- MSE reconstruction loss",
+    "- KL divergence (variational regularization)",
+    "- Wavelet frequency loss (preserves high-frequency details)",
+    "- Perceptual loss via LPIPS-like gradient matching"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# STAGE 1: WAVELET VAE TRAINING\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "\n",
+    "# Build IRIS-Tiny config for free-tier training\n",
+    "# DWT(2\u00d7) + 3 down-blocks(8\u00d7) = 16\u00d7 total compression\n",
+    "# 256px input \u2192 128px after DWT \u2192 64\u219232\u219216 after encoder = 16\u00d716 latent\n",
+    "config = IRISConfig(\n",
+    "    latent_channels=8,        # Smaller for memory efficiency\n",
+    "    latent_spatial=16,        # 16\u00d716 spatial latent\n",
+    "    hidden_dim=384,           # IRIS-Tiny\n",
+    "    num_heads=6,\n",
+    "    head_dim=64,\n",
+    "    ffn_ratio=2.667,\n",
+    "    num_prelude_blocks=1,\n",
+    "    num_core_layers=3,\n",
+    "    num_coda_blocks=1,\n",
+    "    default_iterations=6,\n",
+    "    max_iterations=16,\n",
+    "    fourier_num_blocks=6,\n",
+    "    sparsity_threshold=0.01,\n",
+    "    recurrence_dim=192,\n",
+    "    manhattan_window=12,\n",
+    "    text_dim=768,             # CLIP-L/14\n",
+    "    max_text_tokens=77,\n",
+    "    patch_size=2,\n",
+    "    vae_channels=[32, 64, 128, 256],\n",
+    ")\n",
+    "\n",
+    "# Create VAE\n",
+    "vae = WaveletVAE(config).to(device)\n",
+    "vae_params = sum(p.numel() for p in vae.parameters())\n",
+    "print(f\"Wavelet VAE: {vae_params:,} params ({vae_params*4/1024/1024:.1f} MB fp32)\")\n",
+    "print(f\"Encoder: {sum(p.numel() for p in vae.encoder.parameters()):,}\")\n",
+    "print(f\"Decoder: {sum(p.numel() for p in vae.decoder.parameters()):,}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2500\u2500\u2500 VAE Training Loop \u2500\u2500\u2500\n",
+    "import time\n",
+    "from torch.cuda.amp import autocast, GradScaler\n",
+    "\n",
+    "VAE_EPOCHS = 80            # Enough to get good reconstructions\n",
+    "VAE_LR = 1e-4\n",
+    "KL_WEIGHT = 1e-4           # Light KL to avoid posterior collapse\n",
+    "FREQ_WEIGHT = 0.1          # Wavelet frequency preservation\n",
+    "\n",
+    "optimizer_vae = torch.optim.AdamW(vae.parameters(), lr=VAE_LR, weight_decay=0.01)\n",
+    "scheduler_vae = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_vae, T_max=VAE_EPOCHS)\n",
+    "scaler = GradScaler()\n",
+    "dwt = HaarDWT2D()\n",
+    "\n",
+    "# Logging\n",
+    "vae_losses = {\"total\": [], \"recon\": [], \"kl\": [], \"freq\": []}\n",
+    "\n",
+    "print(f\"Training VAE for {VAE_EPOCHS} epochs on {len(train_loader)} batches...\")\n",
+    "print(f\"{'Epoch':>6} {'Loss':>10} {'Recon':>10} {'KL':>10} {'Freq':>10} {'LR':>10} {'Time':>8}\")\n",
+    "print(\"\u2500\" * 70)\n",
+    "\n",
+    "vae.train()\n",
+    "for epoch in range(VAE_EPOCHS):\n",
+    "    epoch_losses = {\"total\": 0, \"recon\": 0, \"kl\": 0, \"freq\": 0}\n",
+    "    t0 = time.time()\n",
+    "\n",
+    "    for images, _ in train_loader:\n",
+    "        images = images.to(device)\n",
+    "\n",
+    "        with autocast(dtype=torch.float16):\n",
+    "            x_recon, mean, logvar = vae(images)\n",
+    "\n",
+    "            # Reconstruction loss\n",
+    "            recon_loss = torch.nn.functional.mse_loss(x_recon, images)\n",
+    "\n",
+    "            # KL divergence\n",
+    "            kl_loss = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp()).mean()\n",
+    "\n",
+    "            # Wavelet frequency loss \u2014 preserve high-freq details\n",
+    "            with torch.no_grad():\n",
+    "                target_wv = dwt(images)\n",
+    "            recon_wv = dwt(x_recon)\n",
+    "            freq_loss = torch.nn.functional.l1_loss(recon_wv, target_wv)\n",
+    "\n",
+    "            loss = recon_loss + KL_WEIGHT * kl_loss + FREQ_WEIGHT * freq_loss\n",
+    "\n",
+    "        optimizer_vae.zero_grad(set_to_none=True)\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.unscale_(optimizer_vae)\n",
+    "        torch.nn.utils.clip_grad_norm_(vae.parameters(), 1.0)\n",
+    "        scaler.step(optimizer_vae)\n",
+    "        scaler.update()\n",
+    "\n",
+    "        epoch_losses[\"total\"] += loss.item()\n",
+    "        epoch_losses[\"recon\"] += recon_loss.item()\n",
+    "        epoch_losses[\"kl\"] += kl_loss.item()\n",
+    "        epoch_losses[\"freq\"] += freq_loss.item()\n",
+    "\n",
+    "    # Average losses\n",
+    "    n = len(train_loader)\n",
+    "    for k in epoch_losses:\n",
+    "        epoch_losses[k] /= n\n",
+    "        vae_losses[k].append(epoch_losses[k])\n",
+    "\n",
+    "    scheduler_vae.step()\n",
+    "    dt = time.time() - t0\n",
+    "\n",
+    "    if (epoch + 1) % 10 == 0 or epoch == 0:\n",
+    "        lr = optimizer_vae.param_groups[0][\"lr\"]\n",
+    "        print(f\"{epoch+1:>6} {epoch_losses['total']:>10.4f} {epoch_losses['recon']:>10.4f} \"\n",
+    "              f\"{epoch_losses['kl']:>10.4f} {epoch_losses['freq']:>10.4f} {lr:>10.2e} {dt:>7.1f}s\")\n",
+    "\n",
+    "print(\"\\n\u2705 VAE training complete!\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 6.1 Visualize VAE Reconstructions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Visualize reconstructions\n",
+    "vae.eval()\n",
+    "fig, axes = plt.subplots(3, 8, figsize=(20, 8))\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    imgs_sample, _ = next(iter(train_loader))\n",
+    "    imgs_sample = imgs_sample[:8].to(device)\n",
+    "    recon, _, _ = vae(imgs_sample)\n",
+    "\n",
+    "    # Also show latent statistics\n",
+    "    z, mean, logvar = vae.encode(imgs_sample)\n",
+    "    print(f\"Latent shape: {z.shape}\")\n",
+    "    print(f\"Latent mean: {z.mean():.3f}, std: {z.std():.3f}\")\n",
+    "    print(f\"Latent range: [{z.min():.3f}, {z.max():.3f}]\")\n",
+    "\n",
+    "def show_img(ax, tensor, title=\"\"):\n",
+    "    img = tensor.cpu().clamp(-1, 1) * 0.5 + 0.5  # [-1,1] \u2192 [0,1]\n",
+    "    ax.imshow(img.permute(1, 2, 0).numpy())\n",
+    "    ax.set_title(title, fontsize=8)\n",
+    "    ax.axis(\"off\")\n",
+    "\n",
+    "for i in range(8):\n",
+    "    show_img(axes[0, i], imgs_sample[i], f\"Original {i}\")\n",
+    "    show_img(axes[1, i], recon[i], f\"Recon {i}\")\n",
+    "    axes[2, i].imshow(z[i, :3].cpu().permute(1, 2, 0).numpy() * 0.3 + 0.5)\n",
+    "    axes[2, i].set_title(f\"Latent ch0-2\", fontsize=8)\n",
+    "    axes[2, i].axis(\"off\")\n",
+    "\n",
+    "axes[0, 0].set_ylabel(\"Original\", fontsize=12)\n",
+    "axes[1, 0].set_ylabel(\"Reconstructed\", fontsize=12)\n",
+    "axes[2, 0].set_ylabel(\"Latent\", fontsize=12)\n",
+    "plt.suptitle(\"Wavelet VAE Reconstructions\", fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "# Plot loss curves\n",
+    "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
+    "for ax, key, color in zip(axes, [\"total\", \"recon\", \"freq\"], [\"blue\", \"green\", \"red\"]):\n",
+    "    ax.plot(vae_losses[key], color=color)\n",
+    "    ax.set_title(f\"VAE {key.title()} Loss\")\n",
+    "    ax.set_xlabel(\"Epoch\")\n",
+    "    ax.set_ylabel(\"Loss\")\n",
+    "    ax.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Stage 2 \u2014 IRIS Generator Training (Rectified Flow)",
+    "",
+    "Now we train the **recurrent-depth generator** to denoise latent representations conditioned on CLIP text embeddings.",
+    "",
+    "**Key features of this training**:",
+    "- **Rectified Flow**: Linear noise schedule, velocity prediction, logit-normal timestep sampling",
+    "- **Recurrent Depth**: Core block is iterated randomly 4-8\u00d7 per step (training robustness)",
+    "- **adaLN-Zero**: Stable training start via zero-initialized gating",
+    "- **Mixed precision (fp16)**: Fits on 16GB VRAM",
+    "- **Gradient checkpointing**: Optional, for very tight memory",
+    "",
+    "**The magic**: Because the core block shares weights across iterations, ",
+    "we get deep effective network capacity from tiny parameter count!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# STAGE 2: IRIS GENERATOR TRAINING (RECTIFIED FLOW)\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "\n",
+    "# Build full IRIS model (reusing config from VAE stage)\n",
+    "iris = IRIS(config).to(device)\n",
+    "\n",
+    "# Load trained VAE weights\n",
+    "iris.vae.load_state_dict(vae.state_dict())\n",
+    "\n",
+    "# Freeze VAE\n",
+    "for p in iris.vae.parameters():\n",
+    "    p.requires_grad = False\n",
+    "iris.vae.eval()\n",
+    "\n",
+    "gen_params = sum(p.numel() for p in iris.generator.parameters())\n",
+    "core_params = sum(p.numel() for p in iris.generator.core.parameters())\n",
+    "print(f\"IRIS Generator: {gen_params:,} trainable params\")\n",
+    "print(f\"  Core block (shared): {core_params:,} ({core_params/gen_params*100:.1f}%)\")\n",
+    "print(f\"  Effective at r=6: ~{gen_params + 5*core_params:,} effective params\")\n",
+    "print(f\"  Memory fp16: {gen_params*2/1024/1024:.1f} MB\")\n",
+    "\n",
+    "# Free standalone VAE to save memory\n",
+    "del vae, optimizer_vae, scheduler_vae\n",
+    "torch.cuda.empty_cache()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2500\u2500\u2500 Generator Training Loop \u2500\u2500\u2500\n",
+    "GEN_EPOCHS = 150           # More epochs for small dataset\n",
+    "GEN_LR = 2e-4             # Higher LR works well with AdamW + cosine\n",
+    "GRAD_ACCUM = 2             # Effective batch = BATCH_SIZE \u00d7 GRAD_ACCUM = 8\n",
+    "WARMUP_STEPS = 100\n",
+    "\n",
+    "optimizer_gen = torch.optim.AdamW(\n",
+    "    iris.generator.parameters(),\n",
+    "    lr=GEN_LR,\n",
+    "    weight_decay=0.03,\n",
+    "    betas=(0.9, 0.95),\n",
+    ")\n",
+    "\n",
+    "total_steps = GEN_EPOCHS * len(train_loader) // GRAD_ACCUM\n",
+    "\n",
+    "def lr_lambda(step):\n",
+    "    if step < WARMUP_STEPS:\n",
+    "        return step / max(1, WARMUP_STEPS)\n",
+    "    progress = (step - WARMUP_STEPS) / max(1, total_steps - WARMUP_STEPS)\n",
+    "    return 0.5 * (1 + __import__('math').cos(__import__('math').pi * progress))\n",
+    "\n",
+    "scheduler_gen = torch.optim.lr_scheduler.LambdaLR(optimizer_gen, lr_lambda)\n",
+    "scaler_gen = GradScaler()\n",
+    "\n",
+    "# Logging\n",
+    "gen_losses = {\"total\": [], \"velocity\": [], \"kl\": []}\n",
+    "\n",
+    "print(f\"Training generator for {GEN_EPOCHS} epochs ({total_steps} optimizer steps)\")\n",
+    "print(f\"Effective batch size: {BATCH_SIZE} \u00d7 {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}\")\n",
+    "print(f\"Warmup: {WARMUP_STEPS} steps, then cosine decay to 0\")\n",
+    "print()\n",
+    "print(f\"{'Epoch':>6} {'Loss':>10} {'VelLoss':>10} {'MeanT':>8} {'LR':>10} {'Time':>8}\")\n",
+    "print(\"\u2500\" * 60)\n",
+    "\n",
+    "iris.generator.train()\n",
+    "global_step = 0\n",
+    "best_loss = float('inf')\n",
+    "\n",
+    "for epoch in range(GEN_EPOCHS):\n",
+    "    epoch_vel = 0\n",
+    "    epoch_total = 0\n",
+    "    n_batches = 0\n",
+    "    t0 = time.time()\n",
+    "\n",
+    "    optimizer_gen.zero_grad(set_to_none=True)\n",
+    "\n",
+    "    for batch_idx, (images, captions) in enumerate(train_loader):\n",
+    "        images = images.to(device)\n",
+    "\n",
+    "        # Encode text with CLIP\n",
+    "        with torch.no_grad():\n",
+    "            text_emb = encode_text(list(captions))  # [B, 77, 768]\n",
+    "\n",
+    "        # Forward pass with mixed precision\n",
+    "        with autocast(dtype=torch.float16):\n",
+    "            # Randomly sample iteration count for robustness\n",
+    "            r = [4, 5, 6, 7, 8][torch.randint(0, 5, (1,)).item()]\n",
+    "            result = iris.train_step(images, text_emb, num_iterations=r)\n",
+    "            loss = result[\"loss\"] / GRAD_ACCUM\n",
+    "\n",
+    "        scaler_gen.scale(loss).backward()\n",
+    "\n",
+    "        # Gradient accumulation\n",
+    "        if (batch_idx + 1) % GRAD_ACCUM == 0:\n",
+    "            scaler_gen.unscale_(optimizer_gen)\n",
+    "            torch.nn.utils.clip_grad_norm_(iris.generator.parameters(), 1.0)\n",
+    "            scaler_gen.step(optimizer_gen)\n",
+    "            scaler_gen.update()\n",
+    "            optimizer_gen.zero_grad(set_to_none=True)\n",
+    "            scheduler_gen.step()\n",
+    "            global_step += 1\n",
+    "\n",
+    "        epoch_vel += result[\"velocity_loss\"]\n",
+    "        epoch_total += result[\"loss\"].item() if hasattr(result[\"loss\"], 'item') else result[\"velocity_loss\"]\n",
+    "        n_batches += 1\n",
+    "\n",
+    "    avg_vel = epoch_vel / n_batches\n",
+    "    avg_total = epoch_total / n_batches\n",
+    "    gen_losses[\"velocity\"].append(avg_vel)\n",
+    "    gen_losses[\"total\"].append(avg_total)\n",
+    "    dt = time.time() - t0\n",
+    "\n",
+    "    if avg_vel < best_loss:\n",
+    "        best_loss = avg_vel\n",
+    "\n",
+    "    if (epoch + 1) % 10 == 0 or epoch == 0:\n",
+    "        lr = optimizer_gen.param_groups[0][\"lr\"]\n",
+    "        print(f\"{epoch+1:>6} {avg_total:>10.4f} {avg_vel:>10.4f} \"\n",
+    "              f\"{result['mean_t']:>8.3f} {lr:>10.2e} {dt:>7.1f}s\")\n",
+    "\n",
+    "print(f\"\\n\u2705 Generator training complete! Best velocity loss: {best_loss:.4f}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Generate Images!",
+    "",
+    "Now let's generate images using the trained IRIS model. We'll test different iteration budgets ",
+    "to see the adaptive compute in action."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# GENERATION\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "\n",
+    "prompts = [\n",
+    "    \"a fire-breathing dragon pokemon\",\n",
+    "    \"a cute blue water pokemon\",\n",
+    "    \"a green grass-type pokemon with leaves\",\n",
+    "    \"a purple ghost pokemon floating\",\n",
+    "    \"a yellow electric pokemon with lightning\",\n",
+    "    \"a pink fairy pokemon with wings\",\n",
+    "    \"a red phoenix pokemon\",\n",
+    "    \"a small brown fox pokemon\",\n",
+    "]\n",
+    "\n",
+    "iris.eval()\n",
+    "\n",
+    "# Generate with different iteration counts to show adaptive compute\n",
+    "fig, axes = plt.subplots(len(prompts), 4, figsize=(16, len(prompts) * 4))\n",
+    "iteration_counts = [2, 4, 6, 8]\n",
+    "\n",
+    "for row, prompt in enumerate(prompts):\n",
+    "    with torch.no_grad():\n",
+    "        text_emb = encode_text([prompt])\n",
+    "\n",
+    "    for col, n_iter in enumerate(iteration_counts):\n",
+    "        with torch.no_grad():\n",
+    "            img = iris.generate(\n",
+    "                text_emb,\n",
+    "                num_steps=4,\n",
+    "                num_iterations=n_iter,\n",
+    "                cfg_scale=1.0,  # No CFG on untrained model\n",
+    "                seed=42,\n",
+    "            )\n",
+    "        # Convert to displayable\n",
+    "        img_np = img[0].cpu().clamp(-1, 1) * 0.5 + 0.5\n",
+    "        img_np = img_np.permute(1, 2, 0).numpy()\n",
+    "\n",
+    "        axes[row, col].imshow(img_np)\n",
+    "        axes[row, col].axis(\"off\")\n",
+    "        if row == 0:\n",
+    "            axes[row, col].set_title(f\"r={n_iter} iterations\", fontsize=11)\n",
+    "    axes[row, 0].set_ylabel(prompt[:25] + \"...\", fontsize=9, rotation=0, labelpad=120, va='center')\n",
+    "\n",
+    "plt.suptitle(\"IRIS Generated Images (Adaptive Compute Budget)\", fontsize=14, y=1.01)\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "print(\"\\nNote: With only ~800 training images and short training, outputs are noisy.\")\n",
+    "print(\"This demonstrates the architecture works. Quality improves dramatically with:\")\n",
+    "print(\"  \u2022 More training data (CC3M, LAION)\")\n",
+    "print(\"  \u2022 More epochs (1000+)\")\n",
+    "print(\"  \u2022 Larger model (IRIS-Small or IRIS-Base)\")\n",
+    "print(\"  \u2022 Stage 3-5 training (text alignment + aesthetics + distillation)\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 8.1 Training Loss Curves"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "# VAE losses\n",
+    "ax = axes[0]\n",
+    "ax.plot(vae_losses[\"recon\"], label=\"Reconstruction\", color=\"blue\")\n",
+    "ax.plot(vae_losses[\"freq\"], label=\"Wavelet Freq\", color=\"red\")\n",
+    "ax.set_title(\"Stage 1: VAE Losses\")\n",
+    "ax.set_xlabel(\"Epoch\")\n",
+    "ax.set_ylabel(\"Loss\")\n",
+    "ax.legend()\n",
+    "ax.grid(True, alpha=0.3)\n",
+    "ax.set_yscale(\"log\")\n",
+    "\n",
+    "# Generator losses\n",
+    "ax = axes[1]\n",
+    "ax.plot(gen_losses[\"velocity\"], label=\"Velocity Loss\", color=\"green\")\n",
+    "ax.set_title(\"Stage 2: Generator Velocity Loss\")\n",
+    "ax.set_xlabel(\"Epoch\")\n",
+    "ax.set_ylabel(\"Loss\")\n",
+    "ax.legend()\n",
+    "ax.grid(True, alpha=0.3)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Save Checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Save the trained model\n",
+    "import os\n",
+    "os.makedirs(\"iris_checkpoint\", exist_ok=True)\n",
+    "\n",
+    "checkpoint = {\n",
+    "    \"config\": config,\n",
+    "    \"iris_state_dict\": iris.state_dict(),\n",
+    "    \"epoch\": GEN_EPOCHS,\n",
+    "    \"best_velocity_loss\": best_loss,\n",
+    "    \"vae_losses\": vae_losses,\n",
+    "    \"gen_losses\": gen_losses,\n",
+    "}\n",
+    "torch.save(checkpoint, \"iris_checkpoint/iris_trained.pt\")\n",
+    "print(f\"\u2705 Checkpoint saved to iris_checkpoint/iris_trained.pt\")\n",
+    "print(f\"   File size: {os.path.getsize('iris_checkpoint/iris_trained.pt') / 1024 / 1024:.1f} MB\")\n",
+    "\n",
+    "# Optional: push to HF Hub\n",
+    "# from huggingface_hub import HfApi\n",
+    "# api = HfApi()\n",
+    "# api.upload_folder(folder_path=\"iris_checkpoint\", repo_id=\"YOUR_USERNAME/iris-trained\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 10. Inspect Learned Components",
+    "",
+    "Let's peek inside the trained model to understand what the different pathways learned."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Inspect GRFM pathway gating\n",
+    "print(\"=== GRFM Analysis ===\\n\")\n",
+    "\n",
+    "# Look at the blend gate \u2014 does it prefer Fourier or Recurrence?\n",
+    "with torch.no_grad():\n",
+    "    # Get a sample through the model\n",
+    "    imgs_sample, caps = next(iter(train_loader))\n",
+    "    imgs_sample = imgs_sample.to(device)\n",
+    "    text_emb = encode_text(list(caps))\n",
+    "\n",
+    "    z, _, _ = iris.encode(imgs_sample)\n",
+    "    noise = torch.randn_like(z)\n",
+    "    t = torch.tensor([0.5] * z.shape[0], device=device)\n",
+    "    z_t = iris.add_noise(z, noise, t)\n",
+    "\n",
+    "    # Trace through to get GRFM internal state\n",
+    "    x = iris.generator.patch_embed(iris.generator.patchify(z_t)) + iris.generator.pos_embed\n",
+    "    for block in iris.generator.prelude:\n",
+    "        x = block(x)\n",
+    "\n",
+    "    # Get first core layer's GRFM gate values\n",
+    "    core_layer = iris.generator.core.layers[0]\n",
+    "    H, W = iris.generator.patch_h, iris.generator.patch_w\n",
+    "\n",
+    "    # Compute adaLN modulation\n",
+    "    t_emb = iris.generator.time_embed(t * 1000)\n",
+    "    i_emb = iris.generator.iter_embed(torch.zeros(z.shape[0], dtype=torch.long, device=device))\n",
+    "    text_global = iris.generator.text_pool_proj(text_emb.mean(dim=1))\n",
+    "    c = t_emb + i_emb + text_global\n",
+    "\n",
+    "    s1, sh1, g1, *_ = core_layer.adaln(c)\n",
+    "    h_normed = core_layer._modulate(core_layer.norm1(x), s1, sh1)\n",
+    "\n",
+    "    # Get the blend gate value from GRFM\n",
+    "    gate = core_layer.grfm.blend_gate(h_normed)  # [B, N, D]\n",
+    "    gate_mean = gate.mean(dim=(0, 2))  # [N] \u2014 per-position gate\n",
+    "\n",
+    "    # Reshape to 2D\n",
+    "    gate_2d = gate_mean.reshape(H, W).cpu().numpy()\n",
+    "\n",
+    "    fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n",
+    "\n",
+    "    # Gate heatmap\n",
+    "    im = axes[0].imshow(gate_2d, cmap='RdBu_r', vmin=0, vmax=1)\n",
+    "    axes[0].set_title(\"GRFM Blend Gate\\n(red=Fourier, blue=Recurrence)\")\n",
+    "    plt.colorbar(im, ax=axes[0])\n",
+    "\n",
+    "    # Manhattan decay gammas\n",
+    "    gammas = torch.sigmoid(core_layer.grfm.spatial.gamma_logit).cpu().numpy()\n",
+    "    axes[1].bar(range(len(gammas)), gammas)\n",
+    "    axes[1].set_title(\"Manhattan Spatial Decay \u03b3 per Head\\n(lower=more local)\")\n",
+    "    axes[1].set_xlabel(\"Head\")\n",
+    "    axes[1].set_ylabel(\"\u03b3\")\n",
+    "    axes[1].set_ylim(0, 1)\n",
+    "\n",
+    "    # Fourier sparsity (how many coefficients survive soft-shrink)\n",
+    "    x_2d = h_normed.reshape(h_normed.shape[0], H, W, h_normed.shape[-1])\n",
+    "    x_freq = torch.fft.rfft2(x_2d, dim=(1, 2), norm='ortho')\n",
+    "    magnitude = x_freq.abs()\n",
+    "    threshold = core_layer.grfm.fourier.sparsity_threshold\n",
+    "    alive = (magnitude > threshold).float().mean().item()\n",
+    "    axes[2].text(0.5, 0.5, f\"Fourier coefficients\\nabove threshold:\\n{alive*100:.1f}%\",\n",
+    "                 ha='center', va='center', fontsize=16,\n",
+    "                 transform=axes[2].transAxes)\n",
+    "    axes[2].set_title(\"Fourier Domain Sparsity\")\n",
+    "    axes[2].axis(\"off\")\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 11. \ud83d\ude80 Next Steps \u2014 Scaling Up",
+    "",
+    "This notebook trained on ~800 images as a **proof of concept**. To get production quality:",
+    "",
+    "### Datasets for Each Training Stage",
+    "",
+    "| Stage | Dataset | Size | HF ID |",
+    "|-------|---------|------|-------|",
+    "| 1. VAE | ImageNet + CC3M | 4.2M images | `ILSVRC/imagenet-1k`, `pixparse/cc3m-wds` |",
+    "| 2. Class-Cond | ImageNet | 1.2M images | `ILSVRC/imagenet-1k` |",
+    "| 3. Text-Image | CC12M (VLM-recaptioned) | 12M images | `pixparse/cc12m-wds` |",
+    "| 4. Aesthetic | JourneyDB + LAION-art | ~1M images | `caidas/JourneyDB` |",
+    "| 5. Distillation | Self-distill from Stage 4 | Same data | \u2014 |",
+    "",
+    "### Optimization Tips for Larger Runs",
+    "```python",
+    "# On Kaggle with 2\u00d7 T4:",
+    "# Use accelerate for multi-GPU",
+    "# accelerate launch --num_processes 2 train.py",
+    "",
+    "# On Colab Pro (A100 40GB):",
+    "BATCH_SIZE = 16",
+    "GEN_EPOCHS = 500",
+    "config = create_iris_small().config  # Upgrade to IRIS-Small",
+    "",
+    "# For production (cloud GPUs):",
+    "# Use IRIS-Base with 8\u00d7 A100",
+    "# Add LADD adversarial distillation in Stage 5",
+    "# Train for 200k+ steps on CC12M",
+    "```",
+    "",
+    "### Model Size Recommendations",
+    "| Use Case | Model | Batch | Resolution | GPU |",
+    "|----------|-------|-------|-----------|-----|",
+    "| Demo/Proof | IRIS-Tiny | 4 | 256px | T4 16GB |",
+    "| Mobile deploy | IRIS-Small | 8 | 512px | A100 40GB |",
+    "| Quality focus | IRIS-Base | 16 | 512px | 2\u00d7A100 |",
+    "| Production | IRIS-Base | 64 | 1024px | 8\u00d7A100 |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 12. Kaggle Adaptation",
+    "",
+    "To run this on **Kaggle**, just change one thing:",
+    "",
+    "```python",
+    "# In Kaggle, GPU is already available. Just:",
+    "# 1. Copy this notebook to Kaggle",
+    "# 2. Enable \"GPU T4 \u00d72\" or \"GPU P100\" in accelerator settings",
+    "# 3. Run all cells!",
+    "",
+    "# For Kaggle's dual-T4 setup, use DataParallel:",
+    "if torch.cuda.device_count() > 1:",
+    "    print(f\"Using {torch.cuda.device_count()} GPUs!\")",
+    "    iris.generator = torch.nn.DataParallel(iris.generator)",
+    "```",
+    "",
+    "The training loop works identically on both platforms. \ud83c\udf89"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---",
+    "*Built with \u2764\ufe0f using the IRIS architecture. Repository: [asdf98/IRIS-architecture](https://huggingface.co/asdf98/IRIS-architecture)*"
+   ]
+  }
+ ]
+}