torch-neuronx/inference/hf_pretrained_sdxl_base_1024_inference.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## HuggingFace Stable Diffusion XL 1.0 (1024x1024) Inference on Inf2"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Introduction**\n",
                "\n",
                "This notebook demonstrates how to compile and run the HuggingFace Stable Diffusion XL (1024x1024) model for accelerated inference on Neuron.\n",
                "\n",
                "This Jupyter notebook should be run on an Inf2 instance (`inf2.8xlarge` or larger)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Verify that this Jupyter notebook is running the Python kernel environment that was set up according to the [PyTorch Installation Guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/torch-neuronx.html#setup-torch-neuronx). You can select the kernel from the 'Kernel -> Change Kernel' option on the top of this Jupyter notebook page."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Install Dependencies**\n",
                "This tutorial requires the following pip packages to be installed:\n",
                "- `torch-neuronx`\n",
                "- `neuronx-cc`\n",
                "- `diffusers==0.20.2`\n",
                "- `transformers==4.33.1`\n",
                "- `accelerate==0.22.0`\n",
                "- `matplotlib`\n",
                "\n",
                "`torch-neuronx` and `neuronx-cc` will be installed when you configure your environment following the Inf2 setup guide. The remaining dependencies can be installed below:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
                "!pip install diffusers==0.20.2 transformers==4.33.1 accelerate==0.22.0 matplotlib"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**imports**"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                " \n",
                "import numpy as np\n",
                "import torch\n",
                "import torch.nn as nn\n",
                "import torch_neuronx\n",
                "from diffusers import DiffusionPipeline\n",
                "from diffusers.models.unet_2d_condition import UNet2DConditionOutput\n",
                "from diffusers.models.attention_processor import Attention\n",
                "from transformers.models.clip.modeling_clip import CLIPTextModelOutput\n",
                "\n",
                "from matplotlib import pyplot as plt\n",
                "from matplotlib import image as mpimg\n",
                "import time\n",
                "import copy\n",
                "from IPython.display import clear_output\n",
                "\n",
                "clear_output(wait=False)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Define utility classes and functions**\n",
                "\n",
                "The following section defines some utility classes and functions. In particular, we define a double-wrapper for the UNet and text encoders. These wrappers enable `torch_neuronx.trace` to trace the wrapped models for compilation with the Neuron compiler. The second wrapper enables the compiled model (which is a TorchScript object so loses the pre compilation attributes) to be used in the pipeline without having to modify the pipeline source code. In addition, the `get_attention_scores_neuron` utility function performs optimized attention score calculation and is used to replace the origianl `get_attention_scores` function in the `diffusers` package via a monkey patch (see the next code block under \"Compile UNet and save\" for usage)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [],
            "source": [
                "def get_attention_scores_neuron(self, query, key, attn_mask):    \n",
                "    if(query.size() == key.size()):\n",
                "        attention_scores = custom_badbmm(\n",
                "            key,\n",
                "            query.transpose(-1, -2),\n",
                "            self.scale\n",
                "        )\n",
                "        attention_probs = attention_scores.softmax(dim=1).permute(0,2,1)\n",
                "\n",
                "    else:\n",
                "        attention_scores = custom_badbmm(\n",
                "            query,\n",
                "            key.transpose(-1, -2),\n",
                "            self.scale\n",
                "        )\n",
                "        attention_probs = attention_scores.softmax(dim=-1)\n",
                "  \n",
                "    return attention_probs\n",
                " \n",
                "\n",
                "def custom_badbmm(a, b, scale):\n",
                "    bmm = torch.bmm(a, b)\n",
                "    scaled = bmm * scale\n",
                "    return scaled\n",
                " \n",
                "\n",
                "class UNetWrap(nn.Module):\n",
                "    def __init__(self, unet):\n",
                "        super().__init__()\n",
                "        self.unet = unet\n",
                " \n",
                "    def forward(self, sample, timestep, encoder_hidden_states, text_embeds=None, time_ids=None):\n",
                "        out_tuple = self.unet(sample,\n",
                "                              timestep,\n",
                "                              encoder_hidden_states,\n",
                "                              added_cond_kwargs={\"text_embeds\": text_embeds, \"time_ids\": time_ids},\n",
                "                              return_dict=False)\n",
                "        return out_tuple\n",
                "    \n",
                "    \n",
                "class NeuronUNet(nn.Module):\n",
                "    def __init__(self, unetwrap):\n",
                "        super().__init__()\n",
                "        self.unetwrap = unetwrap\n",
                "        self.config = unetwrap.unet.config\n",
                "        self.in_channels = unetwrap.unet.in_channels\n",
                "        self.add_embedding = unetwrap.unet.add_embedding\n",
                "        self.device = unetwrap.unet.device\n",
                " \n",
                "    def forward(self, sample, timestep, encoder_hidden_states, added_cond_kwargs=None, return_dict=False, cross_attention_kwargs=None):\n",
                "        sample = self.unetwrap(sample,\n",
                "                               timestep.float().expand((sample.shape[0],)),\n",
                "                               encoder_hidden_states,\n",
                "                               added_cond_kwargs[\"text_embeds\"],\n",
                "                               added_cond_kwargs[\"time_ids\"])[0]\n",
                "        return UNet2DConditionOutput(sample=sample)\n",
                "    \n",
                "\n",
                "class TextEncoderOutputWrapper(nn.Module):\n",
                "    def __init__(self, traceable_text_encoder, original_text_encoder):\n",
                "        super().__init__()\n",
                "        self.traceable_text_encoder = traceable_text_encoder\n",
                "        self.config = original_text_encoder.config\n",
                "        self.dtype = original_text_encoder.dtype\n",
                "        self.device = original_text_encoder.device\n",
                "\n",
                "    def forward(self, text_input_ids, output_hidden_states=True):\n",
                "        out_tuple = self.traceable_text_encoder(text_input_ids)\n",
                "        return CLIPTextModelOutput(text_embeds=out_tuple[0], last_hidden_state=out_tuple[1], hidden_states=out_tuple[2])\n",
                "    \n",
                "class TraceableTextEncoder(nn.Module):\n",
                "    def __init__(self, text_encoder):\n",
                "        super().__init__()\n",
                "        self.text_encoder = text_encoder\n",
                "\n",
                "    def forward(self, text_input_ids):\n",
                "        out_tuple = self.text_encoder(text_input_ids, output_hidden_states=True, return_dict=False)\n",
                "        return out_tuple"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Compile the model into an optimized TorchScript and save the TorchScript**\n",
                "\n",
                "In the following section, we will compile parts of the Stable Diffusion pipeline for execution on Neuron. Note that this only needs to be done once: After you have compiled and saved the model by running the following section of code, you can reuse it any number of times without having to recompile. In particular, we will compile:\n",
                "1. The text encoders (text_encoder, text_encoder_2)\n",
                "2. The VAE decoder;\n",
                "3. The UNet, and\n",
                "4. The VAE_post_quant_conv\n",
                "These blocks are chosen because they represent the bulk of the compute in the pipeline, and performance benchmarking has shown that running them on Neuron yields significant performance benefit.\n",
                "\n",
                "Several points worth noting are:\n",
                "1. In order to save RAM (these compiles need lots of RAM!), before tracing each model, we make a deepcopy of the part of the pipeline (i.e. the UNet or the VAE decoder) that is to be traced, and then delete the pipeline object from memory with `del pipe`. This trick allows the compile to succeed on instance types with a smaller amount of RAM.\n",
                "2. When compiling each part of the pipeline, we need to pass `torch_neuronx.trace` sample input(s), When there are multiple inputs, they are passed together as a tuple. For details on how to use `torch_neuronx.trace`, please refer to our documentation here: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/api-reference-guide/inference/api-torch-neuronx-trace.html\n",
                "3. Note that while compiling the UNet, we make use of the double-wrapper structure defined above. In addition, we also use the optimized `get_attention_scores_neuron` function to replace the original `get_attention_scores` function in the `Attention` class."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "COMPILER_WORKDIR_ROOT = 'sdxl_compile_dir_1024'\n",
                "\n",
                "# Model ID for SD XL version pipeline\n",
                "model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
                "\n",
                "# --- Compile Text Encoders and save ---\n",
                "\n",
                "pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)\n",
                "\n",
                "\n",
                "# Apply wrappers to make text encoders traceable\n",
                "traceable_text_encoder = copy.deepcopy(TraceableTextEncoder(pipe.text_encoder))\n",
                "traceable_text_encoder_2 = copy.deepcopy(TraceableTextEncoder(pipe.text_encoder_2))\n",
                "\n",
                "del pipe\n",
                "\n",
                "text_input_ids_1 = torch.tensor([[49406,   736,  1615, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
                "         49407, 49407, 49407, 49407, 49407, 49407, 49407]])\n",
                "\n",
                "\n",
                "text_input_ids_2 = torch.tensor([[49406,   736,  1615, 49407,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
                "             0,     0,     0,     0,     0,     0,     0]])\n",
                "\n",
                "\n",
                "# Text Encoder 1\n",
                "neuron_text_encoder = torch_neuronx.trace(\n",
                "    traceable_text_encoder,\n",
                "    text_input_ids_1,\n",
                "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),\n",
                ")\n",
                "\n",
                "text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')\n",
                "torch.jit.save(neuron_text_encoder, text_encoder_filename)\n",
                "\n",
                "\n",
                "# Text Encoder 2\n",
                "neuron_text_encoder_2 = torch_neuronx.trace(\n",
                "    traceable_text_encoder_2,\n",
                "    text_input_ids_2,\n",
                "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder_2'),\n",
                ")\n",
                "\n",
                "text_encoder_2_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder_2/model.pt')\n",
                "torch.jit.save(neuron_text_encoder_2, text_encoder_2_filename)\n",
                "\n",
                "# --- Compile VAE decoder and save ---\n",
                "\n",
                "# Only keep the model being compiled in RAM to minimze memory pressure\n",
                "pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True)\n",
                "decoder = copy.deepcopy(pipe.vae.decoder)\n",
                "del pipe\n",
                "\n",
                "# # Compile vae decoder\n",
                "decoder_in = torch.randn([1, 4, 128, 128])\n",
                "decoder_neuron = torch_neuronx.trace(\n",
                "    decoder, \n",
                "    decoder_in, \n",
                "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder'),\n",
                ")\n",
                "\n",
                "# Save the compiled vae decoder\n",
                "decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder/model.pt')\n",
                "torch.jit.save(decoder_neuron, decoder_filename)\n",
                "\n",
                "# delete unused objects\n",
                "del decoder\n",
                "\n",
                "\n",
                "# --- Compile UNet and save ---\n",
                "\n",
                "pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True)\n",
                "\n",
                "# Replace original cross-attention module with custom cross-attention module for better performance\n",
                "Attention.get_attention_scores = get_attention_scores_neuron\n",
                "\n",
                "# Apply double wrapper to deal with custom return type\n",
                "pipe.unet = NeuronUNet(UNetWrap(pipe.unet))\n",
                "\n",
                "# Only keep the model being compiled in RAM to minimze memory pressure\n",
                "unet = copy.deepcopy(pipe.unet.unetwrap)\n",
                "del pipe\n",
                "\n",
                "# Compile unet - FP32\n",
                "sample_1b = torch.randn([1, 4, 128, 128])\n",
                "timestep_1b = torch.tensor(999).float().expand((1,))\n",
                "encoder_hidden_states_1b = torch.randn([1, 77, 2048])\n",
                "added_cond_kwargs_1b = {\"text_embeds\": torch.randn([1, 1280]),\n",
                "                        \"time_ids\": torch.randn([1, 6])}\n",
                "example_inputs = (sample_1b, timestep_1b, encoder_hidden_states_1b, added_cond_kwargs_1b[\"text_embeds\"], added_cond_kwargs_1b[\"time_ids\"],)\n",
                "\n",
                "unet_neuron = torch_neuronx.trace(\n",
                "    unet,\n",
                "    example_inputs,\n",
                "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'unet'),\n",
                "    compiler_args=[\"--model-type=unet-inference\"]\n",
                ")\n",
                "\n",
                "# save compiled unet\n",
                "unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet/model.pt')\n",
                "torch.jit.save(unet_neuron, unet_filename)\n",
                "\n",
                "# delete unused objects\n",
                "del unet\n",
                "\n",
                "\n",
                "# --- Compile VAE post_quant_conv and save ---\n",
                "\n",
                "# Only keep the model being compiled in RAM to minimze memory pressure\n",
                "pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True)\n",
                "post_quant_conv = copy.deepcopy(pipe.vae.post_quant_conv)\n",
                "del pipe\n",
                "\n",
                "# Compile vae post_quant_conv\n",
                "post_quant_conv_in = torch.randn([1, 4, 128, 128])\n",
                "post_quant_conv_neuron = torch_neuronx.trace(\n",
                "    post_quant_conv, \n",
                "    post_quant_conv_in,\n",
                "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv'),\n",
                ")\n",
                "\n",
                "# Save the compiled vae post_quant_conv\n",
                "post_quant_conv_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv/model.pt')\n",
                "torch.jit.save(post_quant_conv_neuron, post_quant_conv_filename)\n",
                "\n",
                "# delete unused objects\n",
                "del post_quant_conv"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Load the saved model and run it**\n",
                "\n",
                "Now that the model is compiled, you can reload it with any number of prompts. Note the use of the `torch_neuronx.DataParallel` API to load the UNet onto two neuron cores for data-parallel inference. Currently the UNet is the only part of the pipeline that runs data-parallel on two cores. All other parts of the pipeline runs on a single Neuron core.\n",
                "\n",
                "Edit the Prompts below to see what you can create."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "scrolled": false
            },
            "outputs": [],
            "source": [
                "# --- Load all compiled models and run pipeline ---\n",
                "COMPILER_WORKDIR_ROOT = 'sdxl_compile_dir_1024'\n",
                "model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
                "text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')\n",
                "text_encoder_2_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder_2/model.pt')\n",
                "decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder/model.pt')\n",
                "unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet/model.pt')\n",
                "post_quant_conv_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv/model.pt')\n",
                "\n",
                "pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)\n",
                "\n",
                "# Load the compiled UNet onto two neuron cores.\n",
                "pipe.unet = NeuronUNet(UNetWrap(pipe.unet))\n",
                "device_ids = [0,1]\n",
                "pipe.unet.unetwrap = torch_neuronx.DataParallel(torch.jit.load(unet_filename), device_ids, set_dynamic_batching=False)\n",
                "\n",
                "# Load other compiled models onto a single neuron core.\n",
                "pipe.vae.decoder = torch.jit.load(decoder_filename)\n",
                "pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)\n",
                "pipe.text_encoder = TextEncoderOutputWrapper(torch.jit.load(text_encoder_filename), pipe.text_encoder)\n",
                "pipe.text_encoder_2 = TextEncoderOutputWrapper(torch.jit.load(text_encoder_2_filename), pipe.text_encoder_2)\n",
                "\n",
                "# Run pipeline\n",
                "prompt = [\"a photo of an astronaut riding a horse on mars\",\n",
                "          \"sonic on the moon\",\n",
                "          \"elvis playing guitar while eating a hotdog\",\n",
                "          \"saved by the bell\",\n",
                "          \"engineers eating lunch at the opera\",\n",
                "          \"panda eating bamboo on a plane\",\n",
                "          \"A digital illustration of a steampunk flying machine in the sky with cogs and mechanisms, 4k, detailed, trending in artstation, fantasy vivid colors\",\n",
                "          \"kids playing soccer at the FIFA World Cup\"\n",
                "         ]\n",
                "\n",
                "# First do a warmup run so all the asynchronous loads can finish\n",
                "image_warmup = pipe(prompt[0]).images[0]\n",
                "\n",
                "plt.title(\"Image\")\n",
                "plt.xlabel(\"X pixel scaling\")\n",
                "plt.ylabel(\"Y pixels scaling\")\n",
                "\n",
                "total_time = 0\n",
                "for x in prompt:\n",
                "    start_time = time.time()\n",
                "    image = pipe(x).images[0]\n",
                "    total_time = total_time + (time.time()-start_time)\n",
                "    image.save(\"image.png\")\n",
                "    image = mpimg.imread(\"image.png\")\n",
                "    #clear_output(wait=True)\n",
                "    plt.imshow(image)\n",
                "    plt.show()\n",
                "print(\"Average time: \", np.round((total_time/len(prompt)), 2), \"seconds\")\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Now have Fun**\n",
                "\n",
                "Uncomment the cell below for interactive experiment with different prompts."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "scrolled": false
            },
            "outputs": [],
            "source": [
                "# user_input = \"\"\n",
                "# print(\"Enter Prompt, type exit to quit\")\n",
                "# while user_input != \"exit\": \n",
                "#     total_time = 0\n",
                "#     user_input = input(\"What prompt would you like to give?  \")\n",
                "#     if user_input == \"exit\":\n",
                "#         break\n",
                "#     start_time = time.time()\n",
                "#     image = pipe(user_input).images[0]\n",
                "#     total_time = total_time + (time.time()-start_time)\n",
                "#     image.save(\"image.png\")\n",
                "\n",
                "#     plt.title(\"Image\")\n",
                "#     plt.xlabel(\"X pixel scaling\")\n",
                "#     plt.ylabel(\"Y pixels scaling\")\n",
                "\n",
                "#     image = mpimg.imread(\"image.png\")\n",
                "#     plt.imshow(image)\n",
                "#     plt.show()\n",
                "#     print(\"time: \", np.round(total_time, 2), \"seconds\")\n"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "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.8.10"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}