CH4_attention_breakdown.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O_putZlTpVgI"
      },
      "source": [
        "<span style=\"color: red;\">Requirement when running in Goolge Colab</span>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "053w39xVpVgJ",
        "outputId": "fee77e3d-a40d-4537-9816-78e2b7ecbf5d"
      },
      "outputs": [],
      "source": [
        "!pip install diffusers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z4BdpYtFpVgK"
      },
      "source": [
        "#  Chapter 4 - Attention Layer Breakdown\n",
        "\n",
        "Attention layers, a key component of the transformer architecture, were introduced to Stable Diffusion to enhance its ability to capture long-range dependencies and context in image generation. These layers, inspired by the seminal \"Attention Is All You Need\" (https://arxiv.org/abs/1706.03762) paper by Vaswani et al. (2017), allow the model to focus on relevant parts of the input when generating each part of the output. In Stable Diffusion, attention layers play a crucial role in connecting text prompts to visual elements. Understanding and breaking down these attention layers is vital for advanced users and researchers who aim to modify or optimise the model's behavior. By overriding and analysing the attention mechanism, one can gain insights into how the model interprets prompts and constructs images, paving the way for targeted improvements, custom behaviors, or even novel applications of the technology."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k22FQE91pVgL"
      },
      "source": [
        "The original architecture is only displaying the Cross Attention which the Key is the prompt condition and the Query is the latents, but in reality for each Cross-Attention in Stable Diffusion there is a prior Self-Attention which the Key and Query are the image latents."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PPrvMBwhpVgL"
      },
      "source": [
        "![image.png](https://github.com/OutofAi/StableFace/blob/main/assets/image5.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LzHzrZvTpVgL"
      },
      "source": [
        "This part, is technically the same as previous chapter, but for simplification we removed the breakdown and shorten it, all individual elements of the Stable Diffusion architecture can be accessed through the pipeline variable itself, for example the u-net can be accessed by pipe.unet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "d154f0cab9b74aaab3e02215bb8e8f17",
            "59c4eaa914454328afb8992042c4d4ee",
            "11921ea8218e48ebae1b62f953d29b2e",
            "bc7fb817e254442b9fc44b48451bebf6",
            "3fa56930756141fe8b9952ebb628ec55",
            "5be2a24d3bf34e41b5f73b7a1f7f7331",
            "b331efc4535f47458476bce296b90113",
            "27d69dfce9144ba2b2b99b2517c7a6d6",
            "4515100ca2184b1a9e42fb0a795a9450",
            "748eb0d3361d40378f9cf59c0d92e203",
            "b59e55ba67be4e2a968ef07a9c755fd7"
          ]
        },
        "id": "Tc4seE4YpVgL",
        "outputId": "b4fb1d92-a28c-48e0-cd53-8d3b9d7a6166"
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "from diffusers import StableDiffusionPipeline, DDIMScheduler\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "from typing import Optional\n",
        "from tqdm import tqdm\n",
        "from diffusers.models.attention_processor import Attention, AttnProcessor2_0\n",
        "\n",
        "model_id = \"stabilityai/stable-diffusion-2-1-base\"\n",
        "\n",
        "pipe = StableDiffusionPipeline.from_pretrained(model_id)\n",
        "scheduler = DDIMScheduler.from_pretrained(model_id, subfolder=\"scheduler\")\n",
        "pipe = pipe.to(\"cuda\")\n",
        "\n",
        "prompt = \"A photo of a woman, straight hair, light blonde and pink hair, smiling expression, grey background\"\n",
        "\n",
        "org_prompt_embeds = pipe.encode_prompt(prompt=prompt, device=\"cuda\", num_images_per_prompt=1, do_classifier_free_guidance=False)[0]\n",
        "uncond_prompt_embeds = pipe.encode_prompt(prompt=\"\", device=\"cuda\", num_images_per_prompt=1, do_classifier_free_guidance=False)[0]\n",
        "prompt_embeds_combined = torch.cat([uncond_prompt_embeds, org_prompt_embeds])\n",
        "\n",
        "latents = torch.randn((1, pipe.unet.in_channels, pipe.unet.config.sample_size, pipe.unet.config.sample_size), generator=torch.Generator().manual_seed(220)).to(\"cuda\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BIA3tEjnpVgM"
      },
      "source": [
        "First we expand a version of a attention layer calculation. We could then replace it with the main caclulation in the model layers.\n",
        "\n",
        "If you are not familiar with math behind self-attention or cross-attention layers, I highly recommend this 3Blue1Brown YouTube video https://youtu.be/eMlx5fFNoYc explaining in details of how it works.\n",
        "\n",
        "The default calcuation function in Stable Diffusion is an optimised c version which has python binding, so the pure python version will introduce a slight overhead to the current performance of sampling."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qtHy-R8KDr6l"
      },
      "outputs": [],
      "source": [
        "class AttnBreakdownProcessor(AttnProcessor2_0):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "    def __call__(\n",
        "        self,\n",
        "        attn: Attention,\n",
        "        hidden_states: torch.FloatTensor,\n",
        "        encoder_hidden_states: Optional[torch.FloatTensor] = None,\n",
        "        attention_mask: Optional[torch.FloatTensor] = None,\n",
        "        temb: Optional[torch.FloatTensor] = None,\n",
        "        *args,\n",
        "        **kwargs,\n",
        "    ) -> torch.FloatTensor:\n",
        "\n",
        "        residual = hidden_states\n",
        "\n",
        "        input_ndim = hidden_states.ndim\n",
        "\n",
        "        if input_ndim == 4:\n",
        "            batch_size, channel, height, width = hidden_states.shape\n",
        "            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n",
        "\n",
        "        batch_size, _, _ = (\n",
        "            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n",
        "        )\n",
        "\n",
        "        if attn.group_norm is not None:\n",
        "            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n",
        "\n",
        "        query = attn.to_q(hidden_states)\n",
        "\n",
        "        if encoder_hidden_states is None:\n",
        "            encoder_hidden_states = hidden_states\n",
        "        elif attn.norm_cross:\n",
        "            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n",
        "\n",
        "        key = attn.to_k(encoder_hidden_states)\n",
        "        value = attn.to_v(encoder_hidden_states)\n",
        "\n",
        "        query = attn.head_to_batch_dim(query)\n",
        "        key = attn.head_to_batch_dim(key)\n",
        "        value = attn.head_to_batch_dim(value)\n",
        "\n",
        "        attention_scores = attn.scale * torch.bmm(query, key.transpose(-1, -2))\n",
        "\n",
        "        attention_probs = attention_scores.softmax(dim=-1)\n",
        "        del attention_scores\n",
        "\n",
        "        hidden_states = torch.bmm(attention_probs, value)\n",
        "        hidden_states = attn.batch_to_head_dim(hidden_states)\n",
        "        del attention_probs\n",
        "\n",
        "        hidden_states = attn.to_out[0](hidden_states)\n",
        "\n",
        "        if input_ndim == 4:\n",
        "            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n",
        "\n",
        "        if attn.residual_connection:\n",
        "            hidden_states = hidden_states + residual\n",
        "\n",
        "        hidden_states = hidden_states / attn.rescale_output_factor\n",
        "\n",
        "        return hidden_states"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z_w58AnJpVgN"
      },
      "source": [
        "We now replace the default attention processer with our custom python version."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0so-CMR-pVgN"
      },
      "outputs": [],
      "source": [
        "def replace_attention_processor(unet):\n",
        "\n",
        "  replace_processor = AttnBreakdownProcessor()\n",
        "\n",
        "  for name, module in unet.named_modules():\n",
        "    if 'attn1' in name and 'to' not in name:\n",
        "      module.processor = replace_processor\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N0EkHlGTpVgN"
      },
      "source": [
        "Now before running inference steps for sampling we overwrite our current attention layer implementation in the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fHLSnqLipVgO",
        "outputId": "22c75a0a-dee9-4b1b-e335-78f3204b675f"
      },
      "outputs": [],
      "source": [
        "num_inference_steps = 50\n",
        "guidance_scale = 7.5\n",
        "scheduler.set_timesteps(num_inference_steps, device=\"cuda\")\n",
        "timesteps = scheduler.timesteps\n",
        "\n",
        "with torch.no_grad():\n",
        "\n",
        "    # prior to running our model we replace the unet forward function for all attention layer to our custom version\n",
        "    replace_attention_processor(pipe.unet)\n",
        "\n",
        "    for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=\"Inference steps\"):\n",
        "\n",
        "        latent_model_input = torch.cat([latents] * 2)\n",
        "\n",
        "        noise_pred = pipe.unet(\n",
        "            latent_model_input,\n",
        "            t,\n",
        "            encoder_hidden_states=prompt_embeds_combined,\n",
        "            cross_attention_kwargs=None,\n",
        "            return_dict=False,\n",
        "        )[0]\n",
        "\n",
        "\n",
        "        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
        "        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
        "\n",
        "        latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "evUvoz5MpVgO"
      },
      "source": [
        "The results should theoratically look exactly the same as the previous chapter as we only overriden the attention layers forward function for our breakdown"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VKDAeoMkpVgO"
      },
      "outputs": [],
      "source": [
        "with torch.no_grad():\n",
        "    image = pipe.vae.decode(latents / pipe.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": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 452
        },
        "id": "NKiqDlUvpVgP",
        "outputId": "b1c6a425-d2cc-4e7a-882c-f4590835e1db"
      },
      "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"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "11921ea8218e48ebae1b62f953d29b2e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_27d69dfce9144ba2b2b99b2517c7a6d6",
            "max": 6,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4515100ca2184b1a9e42fb0a795a9450",
            "value": 6
          }
        },
        "27d69dfce9144ba2b2b99b2517c7a6d6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3fa56930756141fe8b9952ebb628ec55": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4515100ca2184b1a9e42fb0a795a9450": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "59c4eaa914454328afb8992042c4d4ee": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5be2a24d3bf34e41b5f73b7a1f7f7331",
            "placeholder": "​",
            "style": "IPY_MODEL_b331efc4535f47458476bce296b90113",
            "value": "Loading pipeline components...: 100%"
          }
        },
        "5be2a24d3bf34e41b5f73b7a1f7f7331": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "748eb0d3361d40378f9cf59c0d92e203": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b331efc4535f47458476bce296b90113": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b59e55ba67be4e2a968ef07a9c755fd7": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "bc7fb817e254442b9fc44b48451bebf6": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_748eb0d3361d40378f9cf59c0d92e203",
            "placeholder": "​",
            "style": "IPY_MODEL_b59e55ba67be4e2a968ef07a9c755fd7",
            "value": " 6/6 [00:00&lt;00:00,  8.38it/s]"
          }
        },
        "d154f0cab9b74aaab3e02215bb8e8f17": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_59c4eaa914454328afb8992042c4d4ee",
              "IPY_MODEL_11921ea8218e48ebae1b62f953d29b2e",
              "IPY_MODEL_bc7fb817e254442b9fc44b48451bebf6"
            ],
            "layout": "IPY_MODEL_3fa56930756141fe8b9952ebb628ec55"
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}