CH2_breakdown.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vPde5sjXpS3W"
      },
      "source": [
        "<span style=\"color: red;\">Requirement when running in Goolge Colab</span>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DY7pxP4VpS3X"
      },
      "outputs": [],
      "source": [
        "!pip install diffusers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "We1qVCk9pS3Y"
      },
      "source": [
        "#  Chapter 2 - Breakdown\n",
        "\n",
        "You've probably seen the image describing the architecture behind Stable Diffusion. In this chapter, we attempt to extract the internal sampling process that happens within the Diffusers pipeline. This breakdown will align closer to the main architecture for better understanding and will be utilized later for fine-tuning or prompt2prompt editing of images.\n",
        "\n",
        "- **VAE (Variational Auto-Encoder)**\n",
        "\n",
        "    The VAE handles the pixel space of the image. It consists of two parts:\n",
        "\n",
        "    - Encoder (E): *Compresses the image into a lower-dimensional latent space.*\n",
        "    - Decoder (D): *Reconstructs the image from the latent representation. The VAE enables efficient processing by working in this compressed latent space.*\n",
        "\n",
        "\n",
        "- **UNET and Scheduler**\n",
        "    \n",
        "    These components operate in the latent space:\n",
        "\n",
        "    - Scheduler: *Also known as the Diffusion Process or denoising step. It guides the gradual transformation from noise to image.*\n",
        "    - UNET: *Predicts the noise in the current state of the image, allowing for its removal.*\n",
        "\n",
        "\n",
        "- **Tokenizer and Text Encoder**\n",
        "    \n",
        "    These handle the text input (prompt):\n",
        "\n",
        "    - Tokenizer: *Converts the text prompt into tokens (numerical representations).*\n",
        "    - Text Encoder: *Transforms these tokens into a format that can guide the image generation process*\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hiyYiGgOpS3Z"
      },
      "source": [
        "![image1.png](assets/image1.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LrhvN0aVpS3Z"
      },
      "source": [
        "Importing the relevant libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yTCoqm8XpS3Z"
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "import diffusers\n",
        "from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler\n",
        "from transformers import CLIPTextModel, CLIPTokenizer\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4j3LYY-tpS3a"
      },
      "source": [
        "This time around instead of loading the pretrained pipeline altogther we will load individual elements one by one in order explained above, all of this would be equivalent to pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) that was done in the previous chapter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "frxGrd55pS3a"
      },
      "outputs": [],
      "source": [
        "model_id = \"stabilityai/stable-diffusion-2-1-base\"\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(\n",
        "    model_id, subfolder=\"vae\", revision=None, variant=\"fp16\"\n",
        ").to(\"cuda\")\n",
        "\n",
        "\n",
        "unet = UNet2DConditionModel.from_pretrained(\n",
        "    model_id, subfolder=\"unet\", revision=None, variant=\"fp16\"\n",
        ").to(\"cuda\")\n",
        "scheduler = DDIMScheduler.from_pretrained(model_id, subfolder=\"scheduler\")\n",
        "\n",
        "\n",
        "tokenizer = CLIPTokenizer.from_pretrained(\n",
        "    model_id, subfolder=\"tokenizer\", revision=None, variant=\"fp16\"\n",
        ")\n",
        "text_encoder = CLIPTextModel.from_pretrained(\n",
        "    model_id, subfolder=\"text_encoder\", revision=None, variant=\"fp16\"\n",
        ").to(\"cuda\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ly-sT7mLpS3a"
      },
      "source": [
        "now we have to manually tokenise our prompt and encode it with our text encoder implementing Conditiong section, generating $\\tau_\\theta$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6CSEd9QEpS3b"
      },
      "outputs": [],
      "source": [
        "prompt = \"A photo of a woman, straight hair, light blonde and pink hair, smiling expression, grey background\"\n",
        "\n",
        "text_inputs = tokenizer(\n",
        "                prompt,\n",
        "                padding=\"max_length\",\n",
        "                max_length=tokenizer.model_max_length,\n",
        "                truncation=True,\n",
        "                return_tensors=\"pt\",\n",
        "            ).input_ids\n",
        "prompt_embeds = text_encoder(text_inputs.to(\"cuda\"), attention_mask=None)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ctaKko-qpS3b"
      },
      "source": [
        "Generate a random lantet noise based of the unet config sample size ($\\mathbb{Z}_T$)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yaaXc5TxpS3b"
      },
      "outputs": [],
      "source": [
        "latents = torch.randn((1, unet.in_channels, unet.config.sample_size, unet.config.sample_size), generator=torch.Generator().manual_seed(220)).to(\"cuda\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xTfNMJVhpS3c"
      },
      "source": [
        "In the previous chapter the pipeline runs inference 50 times and that's what we will be doing by providing the latent and the prompt to the unet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGt7_91wpS3c"
      },
      "source": [
        "![image1.png](assets/image2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdXCmEttpS3c"
      },
      "source": [
        "and then providing the unet output to the scheduler for the denoising step"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fktp_3PrpS3c"
      },
      "source": [
        "![image1.png](assets/image3.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YlKRTcCypS3c"
      },
      "outputs": [],
      "source": [
        "num_inference_steps = 50\n",
        "scheduler.set_timesteps(num_inference_steps, device=\"cuda\")\n",
        "timesteps = scheduler.timesteps\n",
        "\n",
        "with torch.no_grad():\n",
        "    for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=\"Inference steps\"):\n",
        "\n",
        "        noise_pred = unet(\n",
        "            latents,\n",
        "            t,\n",
        "            encoder_hidden_states=prompt_embeds,\n",
        "            cross_attention_kwargs=None,\n",
        "            return_dict=False,\n",
        "        )[0]\n",
        "\n",
        "        latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z_Ab36P_pS3c"
      },
      "source": [
        "And afterwards we decode the latent back to the Pixel Space and normalise the output to display the image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "suhIk-x4pS3d"
      },
      "source": [
        "![image1.png](assets/image4.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-ceRP6SWpS3d"
      },
      "outputs": [],
      "source": [
        "with torch.no_grad():\n",
        "    image = vae.decode(latents / vae.config.scaling_factor, return_dict=False)[0]\n",
        "    image_np = image.squeeze(0).float().permute(1,2,0).detach().cpu()\n",
        "    image_np = (image_np / 2 + 0.5).clamp(0, 1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dnMP1EoxpS3d"
      },
      "source": [
        "The results of this process would be not as good as the image generated in the previous chapter, this is due to a key missing component of the the sampling process which will be explained in the next chapter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c-80p9hWpS3d"
      },
      "outputs": [],
      "source": [
        "plt.imshow(image_np)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "L4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}