{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-00",
   "metadata": {},
   "source": [
    "# Text to image with classifier-free guidance\n",
    "\n",
    "The model in [notebook 02](02-train-a-diffusion-model.ipynb) learns what flowers look like, but we cannot tell it which flower to draw. In this notebook we give the model a caption with every image, so at the end we can type \"a red rose\" and get something red and rose-like. We also look at classifier-free guidance, the trick that makes a model follow its caption more closely.\n",
    "\n",
    "The notebook expects one NVIDIA GPU. Training takes most of the time: about twenty minutes on an RTX 4080. The first run downloads the CLIP text encoder (about 1.7 GB) and the flower images (about 270 MB)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-01",
   "metadata": {},
   "source": [
    "## Install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-02",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"dew-ml[streaming] @ git+https://github.com/AshishKumar4/dew\" \"jax[cuda12]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-03",
   "metadata": {},
   "source": [
    "## Settings\n",
    "\n",
    "`UNCONDITIONAL_PROB` is the fraction of training captions we replace with an empty one; the guidance section explains why. `PROMPTS` are the captions we sample at the end, and `GUIDANCE_SCALES` the guidance strengths we compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-04",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_SIZE = 64\n",
    "BATCH_SIZE = 64\n",
    "STEPS = 8000\n",
    "LEARNING_RATE = 3e-4\n",
    "UNCONDITIONAL_PROB = 0.12\n",
    "SAMPLE_STEPS = 40\n",
    "PROMPTS = [\"a red rose\", \"a yellow sunflower\", \"a white daisy\", \"a purple iris\"]\n",
    "GUIDANCE_SCALES = [1.0, 3.0, 6.0]\n",
    "RUN_DIR = \"runs/03-text-to-image\"\n",
    "SEED = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-05",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "print(jax.devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-06",
   "metadata": {},
   "outputs": [],
   "source": [
    "import textwrap\n",
    "\n",
    "def show_images(images, titles, columns=4):\n",
    "    rows = (len(images) + columns - 1) // columns\n",
    "    figure, axes = plt.subplots(rows, columns, figsize=(columns * 1.8, rows * 2.1))\n",
    "    for axis, image, title in zip(np.ravel(axes), images, titles):\n",
    "        axis.imshow(image)\n",
    "        axis.set_title(textwrap.fill(title, 22), fontsize=7)\n",
    "        axis.axis(\"off\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-07",
   "metadata": {},
   "source": [
    "## Turning text into numbers\n",
    "\n",
    "A network cannot read words, so a text encoder turns each caption into a sequence of vectors first. We use the text half of CLIP (ViT-L/14), which was trained to match captions with images, so its vectors already say something about what a caption looks like. We keep CLIP frozen; only the diffusion model trains.\n",
    "\n",
    "The `InputSpec` now has a condition as well as the image. `Condition(encoder, field=\"text\", unconditional=\"\")` says: read tokens from the batch field `text`, encode them with CLIP, and use the empty caption as \"no caption\". The key `textcontext` is the keyword the model receives the encoded caption under."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-08",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dew import Condition, Field, InputSpec\n",
    "from dew.inputs import CLIPText\n",
    "\n",
    "encoder = CLIPText.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "inputs = InputSpec(\n",
    "    sample=Field(\"image\", (IMAGE_SIZE, IMAGE_SIZE, 3)),\n",
    "    conditions={\"textcontext\": Condition(encoder, field=\"text\", unconditional=\"\")},\n",
    ")\n",
    "print(\"context length:\", encoder.context)\n",
    "print(inputs.tokenize([\"a red rose\"])[\"text\"].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-09",
   "metadata": {},
   "source": [
    "## Captioned data\n",
    "\n",
    "The data is the same Hugging Face copy of Oxford Flowers as notebook 02, and every photo carries a caption written by an image captioning model (BLIP). Passing `tokenize=inputs.tokenize` to `load` turns each caption into CLIP tokens while the batch is built, so the batches carry token arrays under `text` instead of strings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dew.data import HFImages, Loading\n",
    "\n",
    "data = HFImages(\n",
    "    name=\"pranked03/flowers-blip-captions\",\n",
    "    image_size=IMAGE_SIZE,\n",
    "    val_batches=0,\n",
    "    loading=Loading(workers=0, threads=16, read_buffer=64),\n",
    ").load(batch=BATCH_SIZE, tokenize=inputs.tokenize)\n",
    "\n",
    "batch = next(iter(data.train()))\n",
    "captions = encoder.captions(batch[\"text\"])\n",
    "show_images(batch[\"image\"][:8], captions[:8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-11",
   "metadata": {},
   "source": [
    "## The model and the objective\n",
    "\n",
    "The model is the DiT from notebook 02. It takes the caption vectors through the `textcontext` keyword and mixes them into every layer.\n",
    "\n",
    "`DiffusionObjective` gets two new arguments. `unconditional_prob` makes it swap the caption for the empty one on 12% of training rows, and `guidance=CFG(3.0)` is how its previews sample. The CLIP weights ride along in the training state as frozen parameters, so they are saved in the checkpoint but never updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optax\n",
    "from dew import Checkpoints, LocalTracker, Trainer, models, presets\n",
    "from dew.objectives.diffusion import DiffusionObjective\n",
    "from dew.sampling import CFG, EulerAncestral\n",
    "\n",
    "process = presets.EDM()()\n",
    "model = models.build(\n",
    "    \"simple_dit\",\n",
    "    patch_size=4,\n",
    "    emb_features=256,\n",
    "    num_layers=6,\n",
    "    num_heads=4,\n",
    "    output_channels=3,\n",
    "    dtype=\"bfloat16\",\n",
    "    attention_impl=\"auto\",\n",
    ")\n",
    "objective = DiffusionObjective(\n",
    "    model, process, inputs,\n",
    "    unconditional_prob=UNCONDITIONAL_PROB,\n",
    "    ema_decay=0.999,\n",
    "    sampler=EulerAncestral(), guidance=CFG(3.0), steps=SAMPLE_STEPS,\n",
    ")\n",
    "trainer = Trainer(\n",
    "    objective, optax.adamw(LEARNING_RATE),\n",
    "    key=jax.random.key(SEED),\n",
    "    checkpoints=Checkpoints(RUN_DIR),\n",
    "    tracker=LocalTracker(f\"{RUN_DIR}/tracking\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-13",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-14",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = trainer.fit(data, steps=STEPS, log_every=1000, checkpoint_every=STEPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "rows = [json.loads(line) for line in open(f\"{RUN_DIR}/tracking/scalars.jsonl\")]\n",
    "rows = [row for row in rows if \"train/loss\" in row[\"scalars\"]]\n",
    "plt.figure(figsize=(6, 3))\n",
    "plt.plot([row[\"step\"] for row in rows], [row[\"scalars\"][\"train/loss\"] for row in rows], marker=\".\")\n",
    "plt.xlabel(\"step\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-16",
   "metadata": {},
   "source": [
    "## Classifier-free guidance\n",
    "\n",
    "Because some captions were blanked during training, one model has learned two things: how to denoise a flower given its caption, and how to denoise a flower with no caption at all. At sampling time we ask it both questions at every step and push the answer away from the uncaptioned one:\n",
    "\n",
    "$\\hat{\\epsilon} = \\epsilon_{\\text{uncond}} + s \\, (\\epsilon_{\\text{cond}} - \\epsilon_{\\text{uncond}})$\n",
    "\n",
    "With scale $s = 1$ this is the plain captioned prediction. Larger scales follow the caption harder and give up some variety. This is classifier-free guidance (Ho and Salimans, 2022), and `CFG(scale)` in Dew computes it.\n",
    "\n",
    "`objective.pipeline(state)` packages the trained model, the process, the text encoder and the EMA weights into a `TextToImage` pipeline. Calling it with a list of prompts returns the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-17",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = objective.pipeline(state)\n",
    "\n",
    "rows = []\n",
    "for scale in GUIDANCE_SCALES:\n",
    "    out = pipe(PROMPTS * 2, steps=SAMPLE_STEPS, guidance=CFG(scale), sampler=EulerAncestral(), seed=1)\n",
    "    rows.append(np.asarray(out.images))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-18",
   "metadata": {},
   "source": [
    "Each row below is one guidance scale, from 1 at the top to 6 at the bottom. Each column is one prompt, and each prompt appears twice per row with different starting noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-19",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = len(PROMPTS) * 2\n",
    "figure, axes = plt.subplots(len(GUIDANCE_SCALES), columns, figsize=(columns * 1.2, len(GUIDANCE_SCALES) * 1.3))\n",
    "for row, (scale, images) in enumerate(zip(GUIDANCE_SCALES, rows)):\n",
    "    for column in range(columns):\n",
    "        axis = axes[row, column]\n",
    "        axis.imshow(np.clip((images[column] + 1) / 2, 0, 1))\n",
    "        axis.set_xticks([])\n",
    "        axis.set_yticks([])\n",
    "        if row == 0:\n",
    "            axis.set_title(textwrap.fill((PROMPTS * 2)[column], 12), fontsize=7)\n",
    "    axes[row, 0].set_ylabel(f\"scale {scale:g}\", fontsize=8)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-20",
   "metadata": {},
   "source": [
    "When we ran this notebook, the colours at scale 1 only loosely followed the prompts: one \"white daisy\" came out dark red. At scale 3 every image had the colour its prompt asked for, and the sunflowers got a dark centre. At scale 6 the colours were the most saturated and the two samples of each prompt looked more alike, which is the variety that strong guidance gives up. Twenty minutes of training on 6,500 images is not enough for sharp flowers, but the caption already steered the colour and the rough shape.\n",
    "\n",
    "## Where to go next\n",
    "\n",
    "`recipes/diffusion/train.py` runs this setup from the command line and writes `run.json` next to the checkpoints, so `TextToImage.from_run(directory)` can rebuild the pipeline later. For larger images, `StableDiffusionVAE` in `dew.nn.autoencoders.sd_vae` lets the model denoise small latents instead of pixels. [Notebook 04](04-samplers-and-schedules.ipynb) compares the samplers."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}