{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PixelModel v2 — Kaggle training and evaluation\n",
    "\n",
    "This notebook trains the 200,259-parameter v2 model on the same approximately 20K disjoint COCO caption/image pairs as v1, then evaluates it on the held-out 5K COCO rows using FID and CLIP Score.\n",
    "\n",
    "Before running: enable **Internet** and **GPU** in Kaggle. Add this repository as a Kaggle Dataset named `pixelmodel-v2`, so it is mounted at `/kaggle/input/pixelmodel-v2`. The notebook copies it to the writable working directory; the trained `model.png`, `model.safetensors`, and evaluation results will be in `/kaggle/working/pixelmodel-v2-output`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip -q install datasets safetensors 'torchmetrics[image]'\n",
    "!pip -q install --upgrade 'torch-fidelity'\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "SOURCE_DIR = Path('/kaggle/input/pixelmodel-v2')\n",
    "PROJECT_DIR = Path('/kaggle/working/pixelmodel-v2')\n",
    "WORK_DIR = Path('/kaggle/working/pm-work')\n",
    "OUTPUT_DIR = Path('/kaggle/working/pixelmodel-v2-output')\n",
    "\n",
    "assert SOURCE_DIR.exists(), f'Attach the v2 repository as a Kaggle Dataset: {SOURCE_DIR}'\n",
    "shutil.rmtree(PROJECT_DIR, ignore_errors=True)\n",
    "shutil.copytree(SOURCE_DIR, PROJECT_DIR)\n",
    "WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
    "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
    "os.chdir(PROJECT_DIR)\n",
    "!python model.py\n",
    "!cat config.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creates the v1-compatible split: ~20K 64x64 training pairs and 5K held-out eval pairs.\n",
    "# This downloads COCO through Hugging Face, so Kaggle Internet must stay enabled.\n",
    "!python fetch_coco_subset.py --out {WORK_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start fresh: do not pass --resume for the first run.\n",
    "# The v2 model is larger than v1, so use a GPU and give it a longer run.\n",
    "!python train.py --data {WORK_DIR / 'coco_train.npz'} --epochs 50 --batch 128 --pixels 1024 --lr 0.002\n",
    "!python convert_to_safetensors.py --model model.png --out model.safetensors\n",
    "# Persist the trained model immediately; Kaggle exposes /kaggle/working as notebook output.\n",
    "!cp model.png model.safetensors config.json {OUTPUT_DIR}/\n",
    "!ls -lh {OUTPUT_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Held-out evaluation: generate 5K images, then report FID (lower is better) and CLIP Score (higher is better).\n",
    "!python eval/run_eval.py --work {WORK_DIR} --model model.png --n 5000\n",
    "!cp {WORK_DIR / 'eval_results.json'} {OUTPUT_DIR}/\n",
    "!ls -lh {OUTPUT_DIR}\n",
    "!cat {WORK_DIR / 'eval_results.json'}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}