torch-neuronx/inference/hf_pretrained_sd2_512_inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HuggingFace Stable Diffusion 2.1 (512x512) Inference on Trn1 / Inf2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Introduction**\n",
    "\n",
    "This notebook demonstrates how to compile and run the HuggingFace Stable Diffusion 2.1 (512x512) model for accelerated inference on Neuron. To deploy on SageMaker, please refer to the [SageMaker samples repository](https://github.com/aws-neuron/aws-neuron-sagemaker-samples/blob/master/inference/stable-diffusion/StableDiffusion2_1.ipynb). For Stable Diffusion 768x768, please see the notebook named `hf_pretrained_sd2_768_inference.ipynb`. \n",
    "\n",
    "This Jupyter notebook should be run on an Inf2 instance (`inf2.8xlarge` or larger)"
   ]
  },
  {
   "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",
    "\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",
    "- `safetensors==0.3.1`\n",
    "\n",
    "This tutorial also supports an older version of the `diffusers`, i.e. the following pip packages,\n",
    "- `torch-neuronx`\n",
    "- `neuronx-cc`\n",
    "- `diffusers==0.14.0`\n",
    "- `transformers==4.30.2`\n",
    "- `accelerate==0.16.0`\n",
    "- `matplotlib`\n",
    "- `safetensors==0.3.1`\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 safetensors==0.3.1 matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**imports**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"NEURON_FUSE_SOFTMAX\"] = \"1\"\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch_neuronx\n",
    "import numpy as np\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",
    "from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler\n",
    "from diffusers.models.unet_2d_condition import UNet2DConditionOutput\n",
    "\n",
    "\n",
    "# Compatibility for diffusers<0.18.0\n",
    "from packaging import version\n",
    "import diffusers\n",
    "diffusers_version = version.parse(diffusers.__version__)\n",
    "use_new_diffusers = diffusers_version >= version.parse('0.18.0')\n",
    "if use_new_diffusers:\n",
    "    from diffusers.models.attention_processor import Attention\n",
    "else:\n",
    "    from diffusers.models.cross_attention import CrossAttention\n",
    "\n",
    "# Define datatype\n",
    "DTYPE = torch.bfloat16\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 another wrapper for the text encoder. These wrappers enable `torch_neuronx.trace` to trace the wrapped models for compilation with the Neuron compiler. In addition, the `get_attention_scores` 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). Finally, a `decode_latent` function is defined and replaces the original `decode_latent` function in StableDiffusionPipelines (see the last code block under for usage). The purpose for this replacement is to enable dtype-based optimizations - it allows us to run the UNet in BF16 and the VAE decoder in FP32."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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, cross_attention_kwargs=None):\n",
    "        out_tuple = self.unet(sample, timestep, encoder_hidden_states, return_dict=False)\n",
    "        return out_tuple\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.device = unetwrap.unet.device\n",
    "\n",
    "    def forward(self, sample, timestep, encoder_hidden_states, cross_attention_kwargs=None, return_dict=False):\n",
    "        sample = self.unetwrap(sample, timestep.to(dtype=DTYPE).expand((sample.shape[0],)), encoder_hidden_states)[0]\n",
    "        return UNet2DConditionOutput(sample=sample)\n",
    "\n",
    "class NeuronTextEncoder(nn.Module):\n",
    "    def __init__(self, text_encoder):\n",
    "        super().__init__()\n",
    "        self.neuron_text_encoder = text_encoder\n",
    "        self.config = text_encoder.config\n",
    "        self.dtype = text_encoder.dtype\n",
    "        self.device = text_encoder.device\n",
    "\n",
    "    def forward(self, emb, attention_mask = None):\n",
    "        return [self.neuron_text_encoder(emb)['last_hidden_state']]\n",
    "\n",
    "# Optimized attention\n",
    "def get_attention_scores(self, query, key, attn_mask):       \n",
    "    dtype = query.dtype\n",
    "\n",
    "    if self.upcast_attention:\n",
    "        query = query.float()\n",
    "        key = key.float()\n",
    "\n",
    "    # Check for square matmuls\n",
    "    if(query.size() == key.size()):\n",
    "        attention_scores = custom_badbmm(\n",
    "            key,\n",
    "            query.transpose(-1, -2)\n",
    "        )\n",
    "\n",
    "        if self.upcast_softmax:\n",
    "            attention_scores = attention_scores.float()\n",
    "\n",
    "        attention_probs = attention_scores.softmax(dim=1).permute(0,2,1)\n",
    "        attention_probs = attention_probs.to(dtype)\n",
    "\n",
    "    else:\n",
    "        attention_scores = custom_badbmm(\n",
    "            query,\n",
    "            key.transpose(-1, -2)\n",
    "        )\n",
    "\n",
    "        if self.upcast_softmax:\n",
    "            attention_scores = attention_scores.float()\n",
    "\n",
    "        attention_probs = attention_scores.softmax(dim=-1)\n",
    "        attention_probs = attention_probs.to(dtype)\n",
    "        \n",
    "    return attention_probs\n",
    "\n",
    "def custom_badbmm(a, b):\n",
    "    bmm = torch.bmm(a, b)\n",
    "    scaled = bmm * 0.125\n",
    "    return scaled\n",
    "\n",
    "def decode_latents(self, latents):\n",
    "    # latents = latents.to(torch.float)\n",
    "    latents = 1 / self.vae.config.scaling_factor * latents\n",
    "    image = self.vae.decode(latents).sample\n",
    "    image = (image / 2 + 0.5).clamp(0, 1)\n",
    "    image = image.cpu().permute(0, 2, 3, 1).float().numpy()\n",
    "    return image"
   ]
  },
  {
   "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 CLIP text encoder;\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` function to replace the original `get_attention_scores` function in the `CrossAttention` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For saving compiler artifacts\n",
    "COMPILER_WORKDIR_ROOT = 'sd2_compile_dir_512'\n",
    "\n",
    "# Model ID for SD version pipeline\n",
    "model_id = \"stabilityai/stable-diffusion-2-1-base\"\n",
    "\n",
    "# --- Compile UNet and save ---\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=DTYPE)\n",
    "\n",
    "# Replace original cross-attention module with custom cross-attention module for better performance\n",
    "if use_new_diffusers:\n",
    "    Attention.get_attention_scores = get_attention_scores\n",
    "else:\n",
    "    CrossAttention.get_attention_scores = get_attention_scores\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\n",
    "sample_1b = torch.randn([1, 4, 64, 64], dtype=DTYPE)\n",
    "timestep_1b = torch.tensor(999, dtype=DTYPE).expand((1,))\n",
    "encoder_hidden_states_1b = torch.randn([1, 77, 1024], dtype=DTYPE)\n",
    "example_inputs = sample_1b, timestep_1b, encoder_hidden_states_1b\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\", \"--enable-fast-loading-neuron-binaries\"]\n",
    ")\n",
    "\n",
    "# Enable asynchronous and lazy loading to speed up model load\n",
    "torch_neuronx.async_load(unet_neuron)\n",
    "torch_neuronx.lazy_load(unet_neuron)\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",
    "del unet_neuron\n",
    "\n",
    "\n",
    "\n",
    "# --- Compile CLIP text encoder and save ---\n",
    "\n",
    "# Only keep the model being compiled in RAM to minimze memory pressure\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=DTYPE)\n",
    "text_encoder = copy.deepcopy(pipe.text_encoder)\n",
    "del pipe\n",
    "\n",
    "# Apply the wrapper to deal with custom return type\n",
    "text_encoder = NeuronTextEncoder(text_encoder)\n",
    "# Compile text encoder\n",
    "# This is used for indexing a lookup table in torch.nn.Embedding,\n",
    "# so using random numbers may give errors (out of range).\n",
    "emb = torch.tensor([[49406, 18376,   525,  7496, 49407,     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",
    "text_encoder_neuron = torch_neuronx.trace(\n",
    "        text_encoder.neuron_text_encoder, \n",
    "        emb, \n",
    "        compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),\n",
    "        compiler_args=[\"--enable-fast-loading-neuron-binaries\"]\n",
    "        )\n",
    "\n",
    "# Enable asynchronous loading to speed up model load\n",
    "torch_neuronx.async_load(text_encoder_neuron)\n",
    "\n",
    "# Save the compiled text encoder\n",
    "text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')\n",
    "torch.jit.save(text_encoder_neuron, text_encoder_filename)\n",
    "\n",
    "# delete unused objects\n",
    "del text_encoder\n",
    "del text_encoder_neuron\n",
    "\n",
    "\n",
    "\n",
    "# --- Compile VAE decoder and save ---\n",
    "\n",
    "# Only keep the model being compiled in RAM to minimze memory pressure\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)\n",
    "decoder = copy.deepcopy(pipe.vae.decoder)\n",
    "del pipe\n",
    "\n",
    "# Compile vae decoder\n",
    "decoder_in = torch.randn([1, 4, 64, 64], dtype=torch.float32)\n",
    "decoder_neuron = torch_neuronx.trace(\n",
    "    decoder, \n",
    "    decoder_in, \n",
    "    compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder'),\n",
    "    compiler_args=[\"--enable-fast-loading-neuron-binaries\"]\n",
    ")\n",
    "\n",
    "# Enable asynchronous loading to speed up model load\n",
    "torch_neuronx.async_load(decoder_neuron)\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",
    "del decoder_neuron\n",
    "\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 = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)\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, 64, 64], dtype=torch.float32)\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",
    "# Enable asynchronous loading to speed up model load\n",
    "torch_neuronx.async_load(post_quant_conv_neuron)\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\n",
    "del post_quant_conv_neuron"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "# --- Load all compiled models ---\n",
    "COMPILER_WORKDIR_ROOT = 'sd2_compile_dir_512'\n",
    "model_id = \"stabilityai/stable-diffusion-2-1-base\"\n",
    "text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/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 = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=DTYPE)\n",
    "pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "\n",
    "# Replaces StableDiffusionPipeline's decode_latents method with our custom decode_latents method defined above.\n",
    "StableDiffusionPipeline.decode_latents = decode_latents\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.text_encoder = NeuronTextEncoder(pipe.text_encoder)\n",
    "pipe.text_encoder.neuron_text_encoder = torch.jit.load(text_encoder_filename)\n",
    "\n",
    "class NeuronTypeConversionWrapper(nn.Module):\n",
    "    def __init__(self, post_quant_conv):\n",
    "        super().__init__()\n",
    "        self.network = post_quant_conv\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.network(x.float())\n",
    "        \n",
    "pipe.vae.decoder = NeuronTypeConversionWrapper(torch.jit.load(decoder_filename))\n",
    "pipe.vae.post_quant_conv = NeuronTypeConversionWrapper(torch.jit.load(post_quant_conv_filename))\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": {},
   "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": 4
}