Optimizing Stable Diffusion for Intel CPUs with NNCF and 🤗 Optimum

Overview

Hugging Face presented a method to optimize Stable Diffusion for Intel CPUs by combining Neural Network Compression Framework (NNCF) quantization-aware training with Token Merging, resulting in up to 5.1x faster inference and a model size reduced to 0.25x of the original PyTorch checkpoint.

Stable Diffusion Optimization

The UNet component of the Stable Diffusion pipeline is the most computationally expensive part, so optimizing it yields large speed gains. Traditional post‑training 8‑bit quantization does not work well for this model because pixel‑level prediction tasks are highly sensitive to parameter changes and the model contains little redundancy due to training on hundreds of millions of samples. To preserve accuracy, more sophisticated quantization methods such as Quantization‑Aware Training (QAT) are required.

Optimization Workflow

We started from a Stable Diffusion model fine‑tuned on the Pokemon dataset (svjack/Stable-Diffusion-Pokemon-en). Using the Diffusers text‑to‑image fine‑tuning example, we integrated NNCF‑based QAT into the training script, added a knowledge‑distillation loss where the original model acts as a teacher, and applied Exponential Moving Average (EMA) to model parameters (excluding quantizers) for training stability. Gradient checkpointing and keeping the EMA model in RAM allowed the whole optimization to run on a single GPU with 24 GB VRAM in less than a day for 4096 iterations.

Going Beyond Quantization‑Aware Training

Quantization alone reduces model footprint, load time, memory consumption and latency. We stacked 8‑bit quantization with the Token Merging (ToME) method, which merges redundant tokens before the self‑attention block to cut computation. The combined workflow includes knowledge distillation, EMA, and gradient checkpointing as described above. Starting again from the Pokemon‑fine‑tuned model, we applied ToME with a merging ratio of 0.4 on top of quantization. The resulting model is intended for inference on client or edge CPUs.

Results

Converting the PyTorch baseline to OpenVINO FP32 gave a 1.9x speedup. Adding 8‑bit quantization increased the speedup to 3.9x versus PyTorch and reduced the model footprint to 0.25x of the original checkpoint. Stacking Token Merging on top of quantization yielded a 5.1x inference speedup while keeping the footprint at the same 0.25x level. All measurements were performed with OpenVINO 2022.3 on a Hugging Face Spaces CPU upgrade instance using 3rd Generation Intel® Xeon® Scalable processors with Intel® Deep Learning Boost technology, using the default 50 inference steps. The blog notes that fewer steps increase speed but may affect image quality, and recommends experimenting with step counts and schedulers.

An example inference pipeline is shown below:

from optimum.intel import OVStableDiffusionPipeline

# Load and compile the pipeline for performance.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil\)).images[0]
output.save("image.png

The training and quantization code are available in the Optimum Intel repository, a demonstration notebook is provided, and optimized models can be found on the Hugging Face Hub under the OpenVINO organization. A live demo runs on Hugging Face Spaces.

What about the general‑purpose Stable Diffusion model?

The workflow demonstrated on the Pokemon model shows that substantial optimization is achievable with modest training resources. Training a general‑purpose Stable Diffusion model from scratch is expensive, but with sufficient budget and hardware the same approach can be applied. The caveat is that Token Merging reduces model capacity; therefore, for more complex datasets a lower merging ratio should be used during optimization.

For further reading on complementary approaches for 4th‑generation Intel Xeon CPUs, see the related Hugging Face blog post on Stable Diffusion inference with Intel.

Sources

Related