{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3451405a",
   "metadata": {},
   "source": [
    "# Spike 01 — SD Turbo via diffusers backend\n",
    "\n",
    "**Objetivo:** validar contrato `GenerationConfig_py_ml` + backend `diffusers_generate_py_ml` con SD Turbo. Sanity check antes de invertir en sd.cpp / FLUX.\n",
    "\n",
    "**Criterio PASS:**\n",
    "- 4 imagenes generadas, seeds deterministas, mismo seed → imagen identica entre runs.\n",
    "- `vram_peak_mb` < 6000 MB.\n",
    "- `duration_ms` < 5000 ms/imagen.\n",
    "- `image_grid` 2x2 visible.\n",
    "\n",
    "**Criterio FAIL:** cualquiera de los anteriores no se cumple → replantear contrato o backend.\n",
    "\n",
    "**Lo que NO se valida aqui:** LoRA loading, comparacion vs sd.cpp, samplers != euler_a (SD Turbo solo soporta 1-step euler_a).\n",
    "\n",
    "**Funciones del registry usadas:**\n",
    "- `cuda_available_py_ml`, `gpu_info_py_ml`, `vram_budget_py_ml`\n",
    "- `diffusers_load_pipeline_py_ml`, `diffusers_generate_py_ml`, `diffusers_unload_py_ml`\n",
    "- `image_grid_py_ml`, `image_save_png_py_ml`\n",
    "- Tipos: `GenerationConfig`, `ModelRef`, `ImageGenResult`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b43e6a2",
   "metadata": {},
   "source": [
    "## 1. Hardware check\n",
    "\n",
    "Verificar CUDA + GPU antes de empezar. Si no hay GPU, el resto del notebook correra en CPU (lento)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f395460",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "FN_ROOT = os.environ.get(\"FN_REGISTRY_ROOT\", \"/home/lucas/fn_registry\")\n",
    "sys.path.insert(0, os.path.join(FN_ROOT, \"python/functions/ml\"))\n",
    "\n",
    "from cuda_available import cuda_available\n",
    "from gpu_info import gpu_info\n",
    "from vram_budget import vram_budget\n",
    "\n",
    "cuda = cuda_available()\n",
    "gpus = gpu_info()\n",
    "print(\"CUDA:\", cuda)\n",
    "print(\"GPUs:\", gpus)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b7c5272",
   "metadata": {},
   "source": [
    "## 2. Hipotesis VRAM\n",
    "\n",
    "SD Turbo es derivado de SD 1.5 (~1.5GB fp16). Usamos `vram_budget` con `sd15/fp16` como proxy. Esperamos `fits=True` y `required_mb` ~2000-4000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b7eb904",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_mb = gpus[0][\"vram_total_mb\"] if gpus else 8000  # fallback CPU/sim\n",
    "budget = vram_budget(\n",
    "    gpu_vram_total_mb=total_mb,\n",
    "    model_type=\"sd15\",\n",
    "    quantization=\"fp16\",\n",
    "    width=512, height=512,\n",
    ")\n",
    "print(\"VRAM total:\", total_mb, \"MB\")\n",
    "print(\"Budget:\", budget)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a89fec8c",
   "metadata": {},
   "source": [
    "## 3. Construir 4 `GenerationConfig`\n",
    "\n",
    "Seeds fijos: 42, 123, 7, 999. Prompts variados. SD Turbo: `steps=1`, `cfg_scale=0.0` (no usa guidance), `sampler=euler_a`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b210dfc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Forzar imports limpios para evitar double-import pydantic\n",
    "for _mod in [\"sampler_name\", \"model_ref\", \"lora_ref\", \"generation_config\", \"image_gen_result\"]:\n",
    "    sys.modules.pop(_mod, None)\n",
    "\n",
    "from model_ref import ModelRef\n",
    "from generation_config import GenerationConfig\n",
    "\n",
    "VAULT = \"/home/lucas/vaults/imagegen_models/diffusers/sd-turbo\"\n",
    "\n",
    "model = ModelRef(\n",
    "    name=\"stabilityai/sd-turbo\",\n",
    "    model_type=\"sd15\",\n",
    "    quantization=\"fp16\",\n",
    "    path=VAULT,\n",
    ")\n",
    "\n",
    "PROMPTS = [\n",
    "    (\"a cinematic shot of a baby racoon wearing a tiny crown\", 42),\n",
    "    (\"watercolor of a cozy library with floating books\", 123),\n",
    "    (\"isometric pixel art of a robot fishing on a pier\", 7),\n",
    "    (\"oil painting of a fox playing chess against a cat\", 999),\n",
    "]\n",
    "\n",
    "configs = [\n",
    "    GenerationConfig(\n",
    "        prompt=p,\n",
    "        seed=s,\n",
    "        steps=1,\n",
    "        cfg_scale=0.0,\n",
    "        sampler=\"euler_a\",\n",
    "        width=512,\n",
    "        height=512,\n",
    "        model=model,\n",
    "    )\n",
    "    for p, s in PROMPTS\n",
    "]\n",
    "for c in configs:\n",
    "    print(c.seed, c.prompt[:60])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743b6bb2",
   "metadata": {},
   "source": [
    "## 4. Cargar pipeline SD Turbo\n",
    "\n",
    "Primera carga ~10-30s (mete pesos en VRAM). Llamadas subsiguientes: cacheadas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0f03b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers_load_pipeline import diffusers_load_pipeline\n",
    "\n",
    "pipe = diffusers_load_pipeline(model=model, device=\"auto\", dtype=\"fp16\")\n",
    "print(\"Pipeline:\", type(pipe).__name__)\n",
    "print(\"Scheduler:\", type(pipe.scheduler).__name__)\n",
    "print(\"Device:\", next(pipe.unet.parameters()).device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a328f117",
   "metadata": {},
   "source": [
    "## 5. Generar 4 imagenes\n",
    "\n",
    "Una a una. Captura `duration_ms` + `vram_peak_mb` por imagen. SD Turbo 1-step: esperado <2s en GPU NVIDIA moderna."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1780d91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Forzar imports limpios para alinear tipos con los que usa diffusers_generate\n",
    "for _mod in [\"image_gen_result\", \"generation_config\", \"model_ref\", \"lora_ref\", \"sampler_name\"]:\n",
    "    sys.modules.pop(_mod, None)\n",
    "\n",
    "from diffusers_generate import diffusers_generate\n",
    "from generation_config import GenerationConfig as GC2\n",
    "from model_ref import ModelRef as MR2\n",
    "\n",
    "# Reconstruir configs con los tipos recien re-importados (alineados con diffusers_generate)\n",
    "model2 = MR2(**model.model_dump())\n",
    "configs2 = [GC2(**c.model_dump()) for c in configs]\n",
    "\n",
    "results = []\n",
    "for c in configs2:\n",
    "    r = diffusers_generate(pipe, c)\n",
    "    print(f\"seed={c.seed} duration={r.duration_ms}ms vram_peak={r.vram_peak_mb}MB\")\n",
    "    results.append(r)\n",
    "print(\"Total:\", len(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffda5ccc",
   "metadata": {},
   "source": [
    "## 6. Grid 2x2 + display\n",
    "\n",
    "Imagenes lado a lado con labels (prompt acortado + seed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3dc597a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from image_grid import image_grid\n",
    "\n",
    "labels = [f\"seed={r.meta.get('seed', '?')} | {c.prompt[:40]}...\" for c, r in zip(configs2, results)]\n",
    "grid = image_grid(\n",
    "    images=[r.image for r in results],\n",
    "    cols=2,\n",
    "    labels=labels,\n",
    "    gap_px=12,\n",
    ")\n",
    "grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae4e3184",
   "metadata": {},
   "source": [
    "## 7. Guardar grid + configs a vault\n",
    "\n",
    "PNG con metadata embebida (prompt/seed) en `~/vaults/imagegen_models/outputs/`. JSON canonico en `configs/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ca8a69",
   "metadata": {},
   "outputs": [],
   "source": [
    "from image_save_png import image_save_png\n",
    "from genconfig_save_json import genconfig_save_json\n",
    "import datetime, os as _os\n",
    "\n",
    "OUT_DIR = \"/home/lucas/vaults/imagegen_models/outputs\"\n",
    "CFG_DIR = \"/home/lucas/vaults/imagegen_models/configs\"\n",
    "stamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "\n",
    "grid_path = image_save_png(\n",
    "    grid,\n",
    "    _os.path.join(OUT_DIR, f\"spike01_grid_{stamp}.png\"),\n",
    "    metadata={\"experiment\": \"spike01_sd_turbo\", \"model\": model2.name, \"n\": str(len(results))},\n",
    ")\n",
    "print(\"grid:\", grid_path)\n",
    "\n",
    "for i, (c, r) in enumerate(zip(configs2, results)):\n",
    "    img_path = image_save_png(\n",
    "        r.image,\n",
    "        _os.path.join(OUT_DIR, f\"spike01_seed{c.seed}_{stamp}.png\"),\n",
    "        metadata={\"prompt\": c.prompt, \"seed\": str(c.seed), \"steps\": str(c.steps), \"sampler\": c.sampler},\n",
    "    )\n",
    "    cfg_path = genconfig_save_json(c, _os.path.join(CFG_DIR, f\"spike01_seed{c.seed}_{stamp}.json\"))\n",
    "    print(f\"  [{i}] {img_path}\")\n",
    "    print(f\"      {cfg_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "071982d6",
   "metadata": {},
   "source": [
    "## 8. Tabla resumen + veredicto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a4a18c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "rows = []\n",
    "for c, r in zip(configs2, results):\n",
    "    rows.append({\n",
    "        \"seed\": c.seed,\n",
    "        \"prompt\": c.prompt[:50] + \"...\",\n",
    "        \"duration_ms\": r.duration_ms,\n",
    "        \"vram_peak_mb\": r.vram_peak_mb,\n",
    "    })\n",
    "\n",
    "try:\n",
    "    import pandas as pd\n",
    "    df = pd.DataFrame(rows)\n",
    "    display(df)\n",
    "except ImportError:\n",
    "    for row in rows:\n",
    "        print(row)\n",
    "\n",
    "max_dur = max(r.duration_ms for r in results)\n",
    "max_vram = max((r.vram_peak_mb or 0) for r in results)\n",
    "pass_dur = max_dur < 5000\n",
    "pass_vram = max_vram < 6000\n",
    "verdict = \"PASS\" if (pass_dur and pass_vram) else \"FAIL\"\n",
    "print(f\"\\nMax duration: {max_dur} ms  (target <5000) -> {'OK' if pass_dur else 'FAIL'}\")\n",
    "print(f\"Max VRAM peak: {max_vram} MB  (target <6000) -> {'OK' if pass_vram else 'FAIL'}\")\n",
    "print(f\"\\nVEREDICTO: {verdict}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19701ac7",
   "metadata": {},
   "source": [
    "## 9. Cleanup VRAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d25cfb9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers_unload import diffusers_unload\n",
    "diffusers_unload(pipe)\n",
    "diffusers_unload(None)  # clear cache global\n",
    "print(\"unloaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f331dc00",
   "metadata": {},
   "source": [
    "## 10. Siguiente paso\n",
    "\n",
    "Si PASS:\n",
    "- Notebook `02_seed_reproducibility.ipynb` — mismo seed, dos runs, verificar hash identico.\n",
    "- Notebook `03_sdxl_turbo.ipynb` cuando se descargue (~6.5GB).\n",
    "- Lanzar Ola 3.B (sd.cpp Python bindings) + descargar FLUX schnell GGUF para validacion cruzada.\n",
    "\n",
    "Si FAIL:\n",
    "- Revisar `vram_budget` table si VRAM peak excede.\n",
    "- Revisar `diffusers_generate` si duration excede (offloading, dtype incorrecto).\n",
    "- Replantear contrato `GenerationConfig` si la firma no cuadra."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
