notebooks/stable-diffusion-v3/stable-diffusion-v3-torch-fx.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image generation with Torch.FX Stable Diffusion v3 and OpenVINO\n",
    "\n",
    "Stable Diffusion V3 is next generation of latent diffusion image Stable Diffusion models family that  outperforms state-of-the-art text-to-image generation systems in typography and prompt adherence, based on human preference evaluations. In comparison with previous versions, it based on Multimodal Diffusion Transformer (MMDiT) text-to-image model that features greatly improved performance in image quality, typography, complex prompt understanding, and resource-efficiency.\n",
    "\n",
    "![mmdit.png](https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/dd079427-89f2-4d28-a10e-c80792d750bf)\n",
    "\n",
    "More details about model can be found in [model card](https://huggingface.co/stabilityai/stable-diffusion-3-medium), [research paper](https://stability.ai/news/stable-diffusion-3-research-paper) and [Stability.AI blog post](https://stability.ai/news/stable-diffusion-3-medium).\n",
    "In this tutorial, we will demonstrate the optimize stable diffusion 3 in a Torch FX representation using NNCF [NNCF](https://github.com/openvinotoolkit/nncf/) for model optimization. Additionally, we will accelerate the pipeline further by running with torch.compile using the openvino backend.\n",
    "If you want to run previous Stable Diffusion versions, please check our other notebooks:\n",
    "\n",
    "* [Stable Diffusion](../stable-diffusion-text-to-image)\n",
    "* [Stable Diffusion v2](../stable-diffusion-v2)\n",
    "* [Stable Diffusion v3](../stable-diffusion-v3)\n",
    "* [Stable Diffusion XL](../stable-diffusion-xl)\n",
    "* [LCM Stable Diffusion](../latent-consistency-models-image-generation)\n",
    "* [Turbo SDXL](../sdxl-turbo)\n",
    "* [Turbo SD](../sketch-to-image-pix2pix-turbo)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/stable-diffusion-v3/stable-diffusion-v3.ipynb\" />\n",
    "\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/stable-diffusion-v3/stable-diffusion-v3-torch-fx.ipynb\" />\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Build PyTorch pipeline](#Build-PyTorch-pipeline)\n",
    "    - [Store the Configs](#Store-the-Configs)\n",
    "- [Run FP Inference](#Run-FP-Inference)\n",
    "- [Convert models to Torch FX](#Convert-models-to-Torch-FX)\n",
    "- [Quantization](#Quantization)\n",
    "    - [Collect Calibration Dataset](#Collect-Calibration-Dataset)\n",
    "    - [Compress and Quantize models](#Compress-and-Quantize-models)\n",
    "    - [Create Optimized Pipeline](#Create-Optimized-Pipeline)\n",
    "    - [Check File Size](#Check-File-Size)\n",
    "    - [Optimized pipeline inference](#Optimized-pipeline-inference)\n",
    "    - [Visualize Results](#Visualize-Results)\n",
    "- [Interactive demo](#Interactive-demo)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"gradio>=4.19\" \"torch>=2.5\" \"torchvision>=0.20\" \"numpy<2.0\" \"transformers\" \"datasets>=2.14.6\" \"opencv-python\" \"pillow\" \"peft>=0.7.0\" \"diffusers>=0.31.0\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "%pip install -qU \"openvino>=2024.3.0\"\n",
    "%pip install -q \"nncf>=2.14.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "if not Path(\"sd3_torch_fx_helper.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/stable-diffusion-v3/sd3_torch_fx_helper.py\")\n",
    "    open(\"sd3_torch_fx_helper.py\", \"w\").write(r.text)\n",
    "\n",
    "if not Path(\"gradio_helper.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/stable-diffusion-v3/gradio_helper.py\")\n",
    "    open(\"gradio_helper.py\", \"w\").write(r.text)\n",
    "\n",
    "if not Path(\"notebook_utils.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\")\n",
    "    open(\"notebook_utils.py\", \"w\").write(r.text)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build PyTorch pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    ">**Note**: run model with notebook, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment these lines to login to huggingfacehub to get access to pretrained model\n",
    "\n",
    "# from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "# try:\n",
    "#     whoami()\n",
    "#     print('Authorization token already provided')\n",
    "# except OSError:\n",
    "#     notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "test_replace": {
     "get_sd3_pipeline()": "get_sd3_pipeline(\"katuni4ka/tiny-random-sd3\")"
    }
   },
   "outputs": [],
   "source": [
    "from sd3_torch_fx_helper import get_sd3_pipeline, init_pipeline\n",
    "\n",
    "pipe = get_sd3_pipeline()\n",
    "pipe.to(\"cpu\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Store the Configs\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "\n",
    "This will be used later when wrapping the Torch FX models to insert back into the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configs_dict = {}\n",
    "configs_dict[\"text_encoder\"] = pipe.text_encoder.config\n",
    "configs_dict[\"text_encoder_2\"] = pipe.text_encoder_2.config\n",
    "configs_dict[\"transformer\"] = pipe.transformer.config\n",
    "configs_dict[\"vae\"] = pipe.vae.config\n",
    "\n",
    "pipe_config = pipe.config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run FP Inference\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "generator = torch.Generator(device=\"cpu\").manual_seed(42)\n",
    "prompt = \"A raccoon trapped inside a glass jar full of colorful candies, the background is steamy with vivid colors\"\n",
    "num_inference_steps = 28\n",
    "with torch.no_grad():\n",
    "    image = pipe(\n",
    "        prompt=prompt,\n",
    "        negative_prompt=\"\",\n",
    "        num_inference_steps=num_inference_steps,\n",
    "        generator=generator,\n",
    "        guidance_scale=5,\n",
    "    ).images[0]\n",
    "image.resize(\n",
    "    (\n",
    "        512,\n",
    "        512,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_utils import device_widget\n",
    "\n",
    "device = device_widget()\n",
    "\n",
    "device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert models to Torch FX\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "This step converts the pytorch models in the hf pipeline to Torch FX representation using the `capture_pre_autograd()` function. \n",
    "\n",
    "\n",
    "The pipeline consists of four important parts:\n",
    "\n",
    "* Clip and T5 Text Encoders to create condition to generate an image from a text prompt.\n",
    "* Transformer for step-by-step denoising latent image representation.\n",
    "* Autoencoder (VAE) for decoding latent space to image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "test_replace": {
     "torch.ones((1, 16, 128, 128))": "torch.ones((1, 16, 32, 32))",
     "torch.ones((1, 3, 128, 128))": "torch.ones((1, 3, 32, 32))",
     "torch.ones((2, 154, 4096))": "torch.ones((2, 154, 32))",
     "torch.ones((2, 16, 128, 128))": "torch.ones((2, 16, 32, 32))",
     "torch.ones((2, 2048))": "torch.ones((2, 64))"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from nncf.torch.dynamic_graph.patch_pytorch import disable_patching\n",
    "\n",
    "text_encoder_input = torch.ones((1, 77), dtype=torch.long)\n",
    "text_encoder_kwargs = {}\n",
    "text_encoder_kwargs[\"output_hidden_states\"] = True\n",
    "\n",
    "vae_encoder_input = torch.ones((1, 3, 128, 128))\n",
    "vae_decoder_input = torch.ones((1, 16, 128, 128))\n",
    "\n",
    "unet_kwargs = {}\n",
    "unet_kwargs[\"hidden_states\"] = torch.ones((2, 16, 128, 128))\n",
    "unet_kwargs[\"timestep\"] = torch.from_numpy(np.array([1, 2], dtype=np.float32))\n",
    "unet_kwargs[\"encoder_hidden_states\"] = torch.ones((2, 154, 4096))\n",
    "unet_kwargs[\"pooled_projections\"] = torch.ones((2, 2048))\n",
    "\n",
    "with torch.no_grad():\n",
    "    with disable_patching():\n",
    "        text_encoder = torch.export.export_for_training(\n",
    "            pipe.text_encoder.eval(),\n",
    "            args=(text_encoder_input,),\n",
    "            kwargs=(text_encoder_kwargs),\n",
    "        ).module()\n",
    "        text_encoder_2 = torch.export.export_for_training(\n",
    "            pipe.text_encoder_2.eval(),\n",
    "            args=(text_encoder_input,),\n",
    "            kwargs=(text_encoder_kwargs),\n",
    "        ).module()\n",
    "        pipe.vae.decoder = torch.export.export_for_training(pipe.vae.decoder.eval(), args=(vae_decoder_input,)).module()\n",
    "        pipe.vae.encoder = torch.export.export_for_training(pipe.vae.encoder.eval(), args=(vae_encoder_input,)).module()\n",
    "        vae = pipe.vae\n",
    "        transformer = torch.export.export_for_training(pipe.transformer.eval(), args=(), kwargs=(unet_kwargs)).module()\n",
    "models_dict = {}\n",
    "models_dict[\"transformer\"] = transformer\n",
    "models_dict[\"vae\"] = vae\n",
    "models_dict[\"text_encoder\"] = text_encoder\n",
    "models_dict[\"text_encoder_2\"] = text_encoder_2\n",
    "del unet_kwargs\n",
    "del vae_encoder_input\n",
    "del vae_decoder_input\n",
    "del text_encoder_input\n",
    "del text_encoder_kwargs\n",
    "del pipe"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "[NNCF](https://github.com/openvinotoolkit/nncf/) enables post-training quantization by adding quantization layers into model graph and then using a subset of the training dataset to initialize the parameters of these additional quantization layers. Quantized operations are executed in `INT8` instead of `FP32`/`FP16` making model inference faster.\n",
    "\n",
    "According to `StableDiffusion3Pipeline` structure, the `transformer` model takes up significant portion of the overall pipeline execution time. Now we will show you how to optimize the transformer part using [NNCF](https://github.com/openvinotoolkit/nncf/) to reduce computation cost and speed up the pipeline. Quantizing the rest of the pipeline does not significantly improve inference performance but can lead to a substantial degradation of accuracy. That's why we use 8-bit weight compression for the rest of the pipeline to reduce memory footprint.\n",
    "\n",
    "Please select below whether you would like to run quantization to improve model inference speed.\n",
    "\n",
    "> **NOTE**: Quantization is time and memory consuming operation. Running quantization code below may take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_utils import quantization_widget\n",
    "\n",
    "to_quantize = quantization_widget()\n",
    "\n",
    "to_quantize"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load `skip magic` extension to skip quantization if `to_quantize` is not selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch `skip_kernel_extension` module\n",
    "import requests\n",
    "\n",
    "r = requests.get(\n",
    "    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
    ")\n",
    "open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
    "\n",
    "%load_ext skip_kernel_extension"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Collect Calibration Dataset\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "test_replace": {
     "calibration_dataset_size = 300": "calibration_dataset_size = 1",
     "init_pipeline(models_dict, configs_dict)": "init_pipeline(models_dict, configs_dict, \"katuni4ka/tiny-random-sd3\")"
    }
   },
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import datasets\n",
    "from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "\n",
    "def disable_progress_bar(pipeline, disable=True):\n",
    "    if not hasattr(pipeline, \"_progress_bar_config\"):\n",
    "        pipeline._progress_bar_config = {\"disable\": disable}\n",
    "    else:\n",
    "        pipeline._progress_bar_config[\"disable\"] = disable\n",
    "\n",
    "\n",
    "class UNetWrapper(SD3Transformer2DModel):\n",
    "    def __init__(self, transformer, config):\n",
    "        super().__init__(**config)\n",
    "        self.transformer = transformer\n",
    "        self.captured_args = []\n",
    "\n",
    "    def forward(self, *args, **kwargs):\n",
    "        del kwargs[\"joint_attention_kwargs\"]\n",
    "        del kwargs[\"return_dict\"]\n",
    "        self.captured_args.append((*args, *tuple(kwargs.values())))\n",
    "        return self.transformer(*args, **kwargs)\n",
    "\n",
    "\n",
    "def collect_calibration_data(\n",
    "    pipe, calibration_dataset_size: int, num_inference_steps: int\n",
    ") -> List[Dict]:\n",
    "\n",
    "    original_unet = pipe.transformer\n",
    "    calibration_data = []\n",
    "    disable_progress_bar(pipe)\n",
    "\n",
    "    dataset = datasets.load_dataset(\n",
    "        \"google-research-datasets/conceptual_captions\",\n",
    "        split=\"train\",\n",
    "        trust_remote_code=True,\n",
    "    ).shuffle(seed=42)\n",
    "\n",
    "    transformer_config = dict(pipe.transformer.config)\n",
    "    del transformer_config[\"model\"]\n",
    "    wrapped_unet = UNetWrapper(pipe.transformer.model, transformer_config)\n",
    "    pipe.transformer = wrapped_unet\n",
    "    # Run inference for data collection\n",
    "    pbar = tqdm(total=calibration_dataset_size)\n",
    "    for i, batch in enumerate(dataset):\n",
    "        prompt = batch[\"caption\"]\n",
    "        if len(prompt) > pipe.tokenizer.model_max_length:\n",
    "            continue\n",
    "        # Run the pipeline\n",
    "        pipe(prompt, num_inference_steps=num_inference_steps)\n",
    "        calibration_data.extend(wrapped_unet.captured_args)\n",
    "        wrapped_unet.captured_args = []\n",
    "        pbar.update(len(calibration_data) - pbar.n)\n",
    "        if pbar.n >= calibration_dataset_size:\n",
    "            break\n",
    "\n",
    "    disable_progress_bar(pipe, disable=False)\n",
    "    pipe.transformer = original_unet\n",
    "    return calibration_data\n",
    "\n",
    "\n",
    "if to_quantize:\n",
    "    pipe = init_pipeline(models_dict, configs_dict)\n",
    "    calibration_dataset_size = 300\n",
    "    unet_calibration_data = collect_calibration_data(\n",
    "        pipe, calibration_dataset_size=calibration_dataset_size, num_inference_steps=28\n",
    "    )\n",
    "    del pipe"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compress and Quantize models\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "import nncf\n",
    "from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters\n",
    "from nncf.quantization.range_estimator import RangeEstimatorParametersSet\n",
    "\n",
    "text_encoder = models_dict[\"text_encoder\"]\n",
    "text_encoder_2 = models_dict[\"text_encoder_2\"]\n",
    "vae_encoder = models_dict[\"vae\"].encoder\n",
    "vae_decoder = models_dict[\"vae\"].decoder\n",
    "original_transformer = models_dict[\"transformer\"]\n",
    "if to_quantize:\n",
    "    with disable_patching():\n",
    "        with torch.no_grad():\n",
    "            nncf.compress_weights(text_encoder)\n",
    "            nncf.compress_weights(text_encoder_2)\n",
    "            nncf.compress_weights(vae_encoder)\n",
    "            nncf.compress_weights(vae_decoder)\n",
    "            quantized_transformer = nncf.quantize(\n",
    "                model=original_transformer,\n",
    "                calibration_dataset=nncf.Dataset(unet_calibration_data),\n",
    "                subset_size=len(unet_calibration_data),\n",
    "                model_type=nncf.ModelType.TRANSFORMER,\n",
    "                ignored_scope=nncf.IgnoredScope(names=[\"conv2d\"]),\n",
    "                advanced_parameters=nncf.AdvancedQuantizationParameters(\n",
    "                    weights_range_estimator_params=RangeEstimatorParametersSet.MINMAX,\n",
    "                    activations_range_estimator_params=RangeEstimatorParametersSet.MINMAX,\n",
    "                ),\n",
    "            )\n",
    "\n",
    "optimized_models_dict = {}\n",
    "optimized_models_dict[\"transformer\"] = quantized_transformer\n",
    "optimized_models_dict[\"vae\"] = vae\n",
    "optimized_models_dict[\"text_encoder\"] = text_encoder\n",
    "optimized_models_dict[\"text_encoder_2\"] = text_encoder_2\n",
    "del models_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "import openvino.torch\n",
    "\n",
    "optimized_models_dict[\"text_encoder\"] = torch.compile(\n",
    "    optimized_models_dict[\"text_encoder\"], backend=\"openvino\"\n",
    ")\n",
    "optimized_models_dict[\"text_encoder_2\"] = torch.compile(\n",
    "    optimized_models_dict[\"text_encoder_2\"], backend=\"openvino\"\n",
    ")\n",
    "optimized_models_dict[\"vae\"].encoder = torch.compile(\n",
    "    optimized_models_dict[\"vae\"].encoder, backend=\"openvino\"\n",
    ")\n",
    "optimized_models_dict[\"vae\"].decoder = torch.compile(\n",
    "    optimized_models_dict[\"vae\"].decoder, backend=\"openvino\"\n",
    ")\n",
    "optimized_models_dict[\"transformer\"] = torch.compile(\n",
    "    optimized_models_dict[\"transformer\"], backend=\"openvino\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Optimized Pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "\n",
    "Initialize the optimized pipeline using the optimized models "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "test_replace": {
     "init_pipeline(optimized_models_dict, configs_dict)": "init_pipeline(optimized_models_dict, configs_dict, \"katuni4ka/tiny-random-sd3\")"
    }
   },
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "opt_pipe = init_pipeline(optimized_models_dict, configs_dict)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check File Size\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "\n",
    "def get_model_size(models):\n",
    "    total_size = 0\n",
    "    for model in models:\n",
    "        param_size = 0\n",
    "        for param in model.parameters():\n",
    "            param_size += param.nelement() * param.element_size()\n",
    "        buffer_size = 0\n",
    "        for buffer in model.buffers():\n",
    "            buffer_size += buffer.nelement() * buffer.element_size()\n",
    "\n",
    "        model_size_mb = (param_size + buffer_size) / 1024**2\n",
    "\n",
    "        total_size += model_size_mb\n",
    "    return total_size\n",
    "\n",
    "\n",
    "optimized_model_size = get_model_size([opt_pipe.transformer])\n",
    "original_model_size = get_model_size([original_transformer])\n",
    "\n",
    "print(f\"Original Transformer Size: {original_model_size} MB\")\n",
    "print(f\"Optimized Transformer Size: {optimized_model_size} MB\")\n",
    "print(f\"Compression Rate: {original_model_size / optimized_model_size:.3f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimized pipeline inference\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "\n",
    "Run inference with single step to compile the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "# Warmup the model for initial compile\n",
    "with torch.no_grad():\n",
    "    image = opt_pipe(\n",
    "        prompt=prompt, negative_prompt=\"\", num_inference_steps=1, generator=generator\n",
    "    ).images[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize Results\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "from sd3_torch_fx_helper import visualize_results\n",
    "\n",
    "generator = torch.Generator(device=\"cpu\").manual_seed(42)\n",
    "opt_image = opt_pipe(\n",
    "    prompt,\n",
    "    negative_prompt=\"\",\n",
    "    num_inference_steps=28,\n",
    "    guidance_scale=5,\n",
    "    generator=generator,\n",
    ").images[0]\n",
    "\n",
    "visualize_results(image, opt_image)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive demo\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Please select below whether you would like to use the quantized models to launch the interactive demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_quantized_models = quantization_widget()\n",
    "\n",
    "use_quantized_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "test_replace": {
     "configs_dict": "configs_dict, \"katuni4ka/tiny-random-sd3\""
    }
   },
   "outputs": [],
   "source": [
    "from gradio_helper import make_demo\n",
    "\n",
    "fx_pipe = init_pipeline(models_dict if not to_quantize.value else optimized_models_dict, configs_dict)\n",
    "demo = make_demo(fx_pipe, False)\n",
    "\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "#  demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# if you have any issue to launch on your platform, you can pass share=True to launch method:\n",
    "# demo.launch(share=True)\n",
    "# it creates a publicly shareable link for the interface. Read more in the docs: https://gradio.app/docs/\n",
    "try:\n",
    "    demo.launch(debug=True)\n",
    "except Exception:\n",
    "    demo.launch(debug=True, share=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.7"
  },
  "openvino_notebooks": {
   "imageUrl": "https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/ac99098c-66ec-4b7b-9e01-e80625f1dc3f",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [
     "Stable Diffusion"
    ],
    "tasks": [
     "Text-to-Image"
    ]
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}