#!/usr/bin/env python
# coding: utf-8
# # High-resolution image generation with Segmind-VegaRT and OpenVINO
#
# The [Segmind Vega](https://huggingface.co/segmind/Segmind-Vega) Model is a distilled version of the [Stable Diffusion XL (SDXL)](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0), offering a remarkable 70% reduction in size and an impressive speedup while retaining high-quality text-to-image generation capabilities. Segmind Vega marks a significant milestone in the realm of text-to-image models, setting new standards for efficiency and speed. Engineered with a compact yet powerful design, it boasts only 745 million parameters. This streamlined architecture not only makes it the smallest in its class but also ensures lightning-fast performance, surpassing the capabilities of its predecessors. Vega represents a breakthrough in model optimization. Its compact size, compared to the 859 million parameters of the SD 1.5 and the hefty 2.6 billion parameters of SDXL, maintains a commendable balance between size and performance. Vega's ability to deliver high-quality images rapidly makes it a game-changer in the field, offering an unparalleled blend of speed, efficiency, and precision.
#
# Segmind Vega is a symmetrical, distilled version of the SDXL model; it is over 70% smaller and ~100% faster. The Down Block contains 247 million parameters, the Mid Block has 31 million, and the Up Block has 460 million. Apart from the size difference, the architecture is virtually identical to that of SDXL, ensuring compatibility with existing interfaces requiring no or minimal adjustments. Although smaller than the SD1.5 Model, Vega supports higher-resolution generation due to the SDXL architecture, making it an ideal replacement for [Stable Diffusion 1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5)
#
# Segmind VegaRT is a distilled LCM-LoRA adapter for the Vega model, that allowed us to reduce the number of inference steps required to generate a good quality image to somewhere between 2 - 8 steps. Latent Consistency Model (LCM) LoRA was proposed in [LCM-LoRA: A universal Stable-Diffusion Acceleration Module](https://arxiv.org/abs/2311.05556) by Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu et al.
#
# More details about models can be found in [Segmind blog post](https://blog.segmind.com/segmind-vega/)
#
# In this tutorial, we explore how to run and optimize Segmind-VegaRT with OpenVINO. We will use a pre-trained model from the [Hugging Face Diffusers](https://huggingface.co/docs/diffusers/index) library. To simplify the user experience, the [Hugging Face Optimum Intel](https://huggingface.co/docs/optimum/intel/index) library is used to convert the models to OpenVINO™ IR format. Additionally, we demonstrate how to improve pipeline latency with the quantization UNet model using [NNCF](https://github.com/openvinotoolkit/nncf).
#
#
#
#
# #### Table of contents:
#
# - [Prerequisites](#Prerequisites)
# - [Prepare PyTorch model](#Prepare-PyTorch-model)
# - [Convert model to OpenVINO format](#Convert-model-to-OpenVINO-format)
# - [Text-to-image generation](#Text-to-image-generation)
# - [Select inference device for text-to-image generation](#Select-inference-device-for-text-to-image-generation)
# - [Quantization](#Quantization)
# - [Prepare calibration dataset](#Prepare-calibration-dataset)
# - [Run quantization](#Run-quantization)
# - [Compare UNet file size](#Compare-UNet-file-size)
# - [Compare the inference time of the FP16 and INT8 models](#Compare-the-inference-time-of-the-FP16-and-INT8-models)
# - [Interactive Demo](#Interactive-Demo)
#
#
# ### Installation Instructions
#
# This is a self-contained example that relies solely on its own code.
#
# We recommend running the notebook in a virtual environment. You only need a Jupyter server to start.
# For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).
#
# 
#
# ## Prerequisites
# [back to top ⬆️](#Table-of-contents:)
# In[ ]:
get_ipython().run_line_magic('pip', 'install -q --extra-index-url https://download.pytorch.org/whl/cpu "torch>=2.1" transformers "diffusers>=0.24.0"')
get_ipython().run_line_magic('pip', 'install -q --extra-index-url https://download.pytorch.org/whl/cpu "git+https://github.com/huggingface/optimum-intel.git" "gradio>=4.19" "openvino>=2023.3.0" "peft>=0.6.2"')
# ## Prepare PyTorch model
# [back to top ⬆️](#Table-of-contents:)
#
#
# For preparing Segmind-VegaRT model for inference, we should create Segmind-Vega pipeline first. After that, for enabling Latent Consistency Model capability, we should integrate VegaRT LCM adapter using `add_lora_weights` method and replace scheduler with LCMScheduler.
# For simplification of these steps for next notebook running, we save created pipeline on disk.
# In[ ]:
import torch
from diffusers import LCMScheduler, AutoPipelineForText2Image
import gc
from pathlib import Path
model_id = "segmind/Segmind-Vega"
adapter_id = "segmind/Segmind-VegaRT"
pt_model_dir = Path("segmind-vegart")
model_dir = Path("openvino-segmind-vegart")
sdxl_model_id = pt_model_dir
tae_id = "madebyollin/taesdxl"
skip_convert_model = model_dir.exists()
if not pt_model_dir.exists() and not skip_convert_model:
pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(adapter_id)
pipe.fuse_lora()
pipe.save_pretrained(pt_model_dir)
del pipe
gc.collect()
# ## Convert model to OpenVINO format
# [back to top ⬆️](#Table-of-contents:)
#
#
# We will use optimum-cli interface for exporting it into OpenVINO Intermediate Representation (IR) format.
#
# Optimum CLI interface for converting models supports export to OpenVINO (supported starting optimum-intel 1.12 version).
# General command format:
#
# ```bash
# optimum-cli export openvino --model --task
# ```
#
# where task is task to export the model for, if not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, as Segmind-Vega uses interface compatible with SDXL, we should be selected `stable-diffusion-xl`
#
# You can find a mapping between tasks and model classes in Optimum TaskManager [documentation](https://huggingface.co/docs/optimum/exporters/task_manager).
#
# Additionally, you can specify weights compression `--weight-format` for the model compression. Please note, that for INT8/INT4, it is necessary to install nncf.
#
# Full list of supported arguments available via `--help`
# For more details and examples of usage, please check [optimum documentation](https://huggingface.co/docs/optimum/intel/inference#export).
#
# For Tiny Autoencoder, we will use `ov.convert_model` function for obtaining `ov.Model` and save it using `ov.save_model`. Model consists of 2 parts that used in pipeline separately:
# `vae_encoder` for encoding input image in latent space in image-to-image generation task and `vae_decoder` that responsible for decoding diffusion result back to image format.
# In[3]:
from pathlib import Path
import requests
if not Path("cmd_helper.py").exists():
r = requests.get("https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/cmd_helper.py")
with open("cmd_helper.py", "w") as f:
f.write(r.text)
# In[4]:
import torch
import openvino as ov
from diffusers import AutoencoderTiny
import gc
from cmd_helper import optimum_cli
class VAEEncoder(torch.nn.Module):
def __init__(self, vae):
super().__init__()
self.vae = vae
def forward(self, sample):
return self.vae.encode(sample)
class VAEDecoder(torch.nn.Module):
def __init__(self, vae):
super().__init__()
self.vae = vae
def forward(self, latent_sample):
return self.vae.decode(latent_sample)
def convert_tiny_vae(model_id, output_path):
tiny_vae = AutoencoderTiny.from_pretrained(model_id)
tiny_vae.eval()
vae_encoder = VAEEncoder(tiny_vae)
ov_model = ov.convert_model(vae_encoder, example_input=torch.zeros((1, 3, 512, 512)))
ov.save_model(ov_model, output_path / "vae_encoder/openvino_model.xml")
tiny_vae.save_config(output_path / "vae_encoder")
vae_decoder = VAEDecoder(tiny_vae)
ov_model = ov.convert_model(vae_decoder, example_input=torch.zeros((1, 4, 64, 64)))
ov.save_model(ov_model, output_path / "vae_decoder/openvino_model.xml")
tiny_vae.save_config(output_path / "vae_decoder")
del tiny_vae
del ov_model
gc.collect()
if not skip_convert_model:
optimum_cli(sdxl_model_id, model_dir, additional_args={"task": "text-to-image", "weight-format": "fp16"})
convert_tiny_vae(tae_id, model_dir)
# ## Text-to-image generation
# [back to top ⬆️](#Table-of-contents:)
#
#
# Text-to-image generation lets you create images using text description. To start generating images, we need to load models first.
# To load an OpenVINO model and run an inference with Optimum and OpenVINO Runtime, you need to replace diffusers `StableDiffusionXLPipeline` with Optimum `OVStableDiffusionXLPipeline`. Pipeline initialization starts with using `from_pretrained` method, where a directory with OpenVINO models should be passed. Additionally, you can specify an inference device.
#
# For saving time, we will not cover image-to-image generation in this notebook. As we already mentioned, Segmind-Vega is compatible with Stable Diffusion XL pipeline, the steps required to run Stable Diffusion XL inference for image-to-image task were discussed in this [notebook](stable-dffision-xl.ipynb).
# ### Select inference device for text-to-image generation
# [back to top ⬆️](#Table-of-contents:)
# In[5]:
if not Path("notebook_utils.py").exists():
r = requests.get(
url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py",
)
open("notebook_utils.py", "w").write(r.text)
from notebook_utils import device_widget
device = device_widget()
device
# In[6]:
from optimum.intel.openvino import OVStableDiffusionXLPipeline
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)
# In[7]:
from transformers import set_seed
set_seed(23)
prompt = "A cinematic highly detailed shot of a baby Yorkshire terrier wearing an intricate Italian priest robe, with crown"
image = text2image_pipe(prompt, num_inference_steps=4, height=512, width=512, guidance_scale=0.5).images[0]
image.save("dog.png")
image
# In[8]:
del text2image_pipe
gc.collect();
# ## Quantization
# [back to top ⬆️](#Table-of-contents:)
#
#
# [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.
#
# According to `Segmind-VEGAModel` structure, the UNet model takes up significant portion of the overall pipeline execution time. Now we will show you how to optimize the UNet part using [NNCF](https://github.com/openvinotoolkit/nncf/) to reduce computation cost and speed up the pipeline. Quantizing the rest of the SDXL pipeline does not significantly improve inference performance but can lead to a substantial degradation of accuracy.
#
# The optimization process contains the following steps:
#
# 1. Create a calibration dataset for quantization.
# 2. Run `nncf.quantize()` to obtain quantized model.
# 3. Save the `INT8` model using `openvino.save_model()` function.
#
# Please select below whether you would like to run quantization to improve model inference speed.
# In[9]:
from notebook_utils import quantization_widget
to_quantize = quantization_widget()
to_quantize
# In[10]:
# Fetch `skip_kernel_extension` module
if not Path("skip_kernel_extension.py").exists():
r = requests.get(
url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py",
)
open("skip_kernel_extension.py", "w").write(r.text)
int8_pipe = None
core = ov.Core()
def create_int8_pipe(model_dir, unet_int8_path, device, core, unet_device="CPU"):
int8_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device, compile=True)
del int8_pipe.unet.request
del int8_pipe.unet.model
gc.collect()
int8_pipe.unet.model = core.read_model(unet_int8_path)
int8_pipe.unet.request = core.compile_model(int8_pipe.unet.model, unet_device or device)
return int8_pipe
if to_quantize.value and "GPU" in device.value:
to_quantize.value = False
42
get_ipython().run_line_magic('load_ext', 'skip_kernel_extension')
# ### Prepare calibration dataset
# [back to top ⬆️](#Table-of-contents:)
#
#
# We use a portion of [conceptual_captions](https://huggingface.co/datasets/conceptual_captions) dataset from Hugging Face as calibration data.
# To collect intermediate model inputs for calibration we should customize `CompiledModel`.
# In[11]:
UNET_INT8_OV_PATH = model_dir / "optimized_unet" / "openvino_model.xml"
def disable_progress_bar(pipeline, disable=True):
if not hasattr(pipeline, "_progress_bar_config"):
pipeline._progress_bar_config = {"disable": disable}
else:
pipeline._progress_bar_config["disable"] = disable
# In[12]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nimport datasets\nimport numpy as np\nfrom tqdm.notebook import tqdm\nfrom transformers import set_seed\nfrom typing import Any, Dict, List\n\nset_seed(1)\n\nclass CompiledModelDecorator(ov.CompiledModel):\n def __init__(self, compiled_model: ov.CompiledModel, data_cache: List[Any] = None):\n super().__init__(compiled_model)\n self.data_cache = data_cache if data_cache else []\n\n def __call__(self, *args, **kwargs):\n self.data_cache.append(*args)\n return super().__call__(*args, **kwargs)\n\ndef collect_calibration_data(pipe, subset_size: int) -> List[Dict]:\n original_unet = pipe.unet.request\n pipe.unet.request = CompiledModelDecorator(original_unet)\n\n dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True).shuffle(seed=42)\n disable_progress_bar(pipe)\n\n # Run inference for data collection\n pbar = tqdm(total=subset_size)\n diff = 0\n for batch in dataset:\n prompt = batch["caption"]\n if len(prompt) > pipe.tokenizer.model_max_length:\n continue\n _ = pipe(\n prompt,\n num_inference_steps=1,\n height=512,\n width=512,\n guidance_scale=0.0,\n generator=np.random.RandomState(987)\n )\n collected_subset_size = len(pipe.unet.request.data_cache)\n if collected_subset_size >= subset_size:\n pbar.update(subset_size - pbar.n)\n break\n pbar.update(collected_subset_size - diff)\n diff = collected_subset_size\n\n calibration_dataset = pipe.unet.request.data_cache\n disable_progress_bar(pipe, disable=False)\n pipe.unet.request = original_unet\n return calibration_dataset\n')
# In[13]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nif not UNET_INT8_OV_PATH.exists():\n text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)\n unet_calibration_data = collect_calibration_data(text2image_pipe, subset_size=200)\n')
# ### Run quantization
# [back to top ⬆️](#Table-of-contents:)
#
#
# Create a quantized model from the pre-trained converted OpenVINO model. Quantization of the first and last `Convolution` layers impacts the generation results. We recommend using `IgnoredScope` to keep accuracy sensitive `Convolution` layers in FP16 precision.
#
# > **NOTE**: Quantization is time and memory consuming operation. Running quantization code below may take some time.
# In[14]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nimport nncf\nfrom nncf.scopes import IgnoredScope\n\nUNET_OV_PATH = model_dir / "unet" / "openvino_model.xml"\nif not UNET_INT8_OV_PATH.exists():\n unet = core.read_model(UNET_OV_PATH)\n quantized_unet = nncf.quantize(\n model=unet,\n model_type=nncf.ModelType.TRANSFORMER,\n calibration_dataset=nncf.Dataset(unet_calibration_data),\n ignored_scope=IgnoredScope(\n names=[\n "__module.model.conv_in/aten::_convolution/Convolution",\n "__module.model.up_blocks.2.resnets.2.conv_shortcut/aten::_convolution/Convolution",\n "__module.model.conv_out/aten::_convolution/Convolution"\n ],\n ),\n )\n ov.save_model(quantized_unet, UNET_INT8_OV_PATH)\n')
# In[15]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nint8_text2image_pipe = create_int8_pipe(model_dir, UNET_INT8_OV_PATH, device.value, core)\n\n\nset_seed(23)\n \nimage = int8_text2image_pipe(prompt, num_inference_steps=4, height=512, width=512, guidance_scale=0.5).images[0]\ndisplay(image)\n')
# #### Compare UNet file size
# [back to top ⬆️](#Table-of-contents:)
# In[16]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nfp16_ir_model_size = UNET_OV_PATH.with_suffix(".bin").stat().st_size / 1024\nquantized_model_size = UNET_INT8_OV_PATH.with_suffix(".bin").stat().st_size / 1024\n\nprint(f"FP16 model size: {fp16_ir_model_size:.2f} KB")\nprint(f"INT8 model size: {quantized_model_size:.2f} KB")\nprint(f"Model compression rate: {fp16_ir_model_size / quantized_model_size:.3f}")\n')
# ### Compare the inference time of the FP16 and INT8 models
# [back to top ⬆️](#Table-of-contents:)
#
#
# To measure the inference performance of the `FP16` and `INT8` pipelines, we use median inference time on the calibration subset.
#
# > **NOTE**: For the most accurate performance estimation, it is recommended to run `benchmark_app` in a terminal/command prompt after closing other applications.
# In[17]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nimport time\n\nvalidation_size = 7\ncalibration_dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True)\nvalidation_data = []\nfor idx, batch in enumerate(calibration_dataset):\n if idx >= validation_size:\n break\n prompt = batch["caption"]\n validation_data.append(prompt)\n\ndef calculate_inference_time(pipe, dataset):\n inference_time = []\n disable_progress_bar(pipe)\n\n for prompt in dataset:\n start = time.perf_counter()\n image = pipe(\n prompt,\n num_inference_steps=4,\n guidance_scale=1.0,\n generator=np.random.RandomState(23)\n ).images[0]\n end = time.perf_counter()\n delta = end - start\n inference_time.append(delta)\n disable_progress_bar(pipe, disable=False)\n return np.median(inference_time)\n')
# In[18]:
get_ipython().run_cell_magic('skip', 'not $to_quantize.value', '\nint8_latency = calculate_inference_time(int8_text2image_pipe, validation_data)\n\ndel int8_text2image_pipe\ngc.collect()\n\ntext2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)\nfp_latency = calculate_inference_time(text2image_pipe, validation_data)\n\ndel text2image_pipe\ngc.collect()\nprint(f"FP16 pipeline latency: {fp_latency:.3f}")\nprint(f"INT8 pipeline latency: {int8_latency:.3f}")\nprint(f"Text-to-Image generation speed up: {fp_latency / int8_latency:.3f}")\n')
# ## Interactive Demo
# [back to top ⬆️](#Table-of-contents:)
#
#
# Now, you can check model work using own text descriptions. Provide text prompt in the text box and launch generation using Run button. Additionally you can control generation with additional parameters:
# * Seed - random seed for initialization
# * Steps - number of generation steps
# * Height and Width - size of generated image
#
# Please select below whether you would like to use the quantized model to launch the interactive demo.
# In[19]:
import ipywidgets as widgets
quantized_model_present = UNET_INT8_OV_PATH.exists()
use_quantized_model = widgets.Checkbox(
value=quantized_model_present,
description="Use quantized model",
disabled=not quantized_model_present,
)
use_quantized_model
# In[ ]:
if not Path("gradio_helper.py").exists():
r = requests.get(url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/stable-diffusion-xl/gradio_helper.py")
open("gradio_helper.py", "w").write(r.text)
from gradio_helper import make_demo_segmind_vegart
if use_quantized_model.value:
if not quantized_model_present:
raise RuntimeError("Quantized model not found.")
text2image_pipe = create_int8_pipe(model_dir, UNET_INT8_OV_PATH, device.value, core)
else:
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)
demo = make_demo_segmind_vegart(text2image_pipe, use_quantized_model)
# if you are launching remotely, specify server_name and server_port
# demo.launch(server_name='your server name', server_port='server port in int')
# Read more in the docs: https://gradio.app/docs/
# if you want create public link for sharing demo, please add share=True
try:
demo.launch(debug=True)
except Exception:
demo.launch(share=True, debug=True)