--- language: - en license: other license_name: stabilityai-community license_link: https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE library_name: diffusers tags: - text-to-image - stable-diffusion - stable-diffusion-3.5 - stable-diffusion-3.5-large - pruning - structural-pruning - obs-diff - pytorch base_model: stabilityai/stable-diffusion-3.5-large pipeline_tag: text-to-image --- # OBS-Diff Structured Pruning for Stable Diffusion 3.5-Large
OBS-Diff

✂️ OBS-Diff: Accurate Pruning for Diffusion Models in One-Shot

Junhan Zhu, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang*

The first training-free, one-shot pruning framework for Diffusion Models, supporting diverse architectures and pruning granularities. Uses Optimal Brain Surgeon (OBS) to achieve SOTA compression with high generative quality.

This repository contains the structured pruned checkpoints for Stable Diffusion 3.5 Large. These models were compressed using OBS-Diff, an accurate one-shot pruning method designed to reduce model size and accelerate inference while preserving high-quality image generation capabilities. By removing redundant parameters from the Transformer backbone, we offer variants with different sparsity levels (15% - 30%), allowing for a flexible trade-off between efficiency and performance. ![](sd3-5-2.png) ![](sd3-5-1.png) ![](sd3-5-3.png) ### Pruned Transformer Variants | Sparsity (%) | 0 (Dense) | 15 | 20 | 25 | 30 | | :--- | :---: | :---: | :---: | :---: | :---: | | **Params (B)** | 8.06 | 7.28 | 7.02 | 6.76 | 6.54 | ### How to use the pruned model 1. Download the base model (SD3.5-Large) from [huggingface](https://huggingface.co/stabilityai/stable-diffusion-3.5-large) or ModelScope. 2. Download the pruned weights (.pth files) and use `torch.load` to replace the original Transformer in the pipeline. 3. Run inference using the code below. ``` python import os import torch from diffusers import StableDiffusion3Pipeline from PIL import Image # 1. Load the base SD3.5-Large model pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.float16) # 2. Swap the original Transformer with the pruned Transformer checkpoint # Note: Ensure the path points to your downloaded .pth file pruned_transformer_path = "/path/to/sparsity_30/pruned_model.pth" pipe.transformer = torch.load(pruned_transformer_path, weights_only=False) pipe = pipe.to("cuda") total_params = sum(p.numel() for p in pipe.transformer.parameters()) print(f"Total Transformer parameters: {total_params / 1e6:.2f} M") image = pipe( prompt="photo of a delicious hamburger with fries and a coke on a wooden table, professional food photography, bokeh", negative_prompt=None, height=1024, width=1024, num_inference_steps=30, guidance_scale=7.0, generator=torch.Generator("cuda").manual_seed(42) ).images[0] image.save("output_pruned.png") ``` ### Citation If you find this work useful, please consider citing: ```bibtex @article{zhu2025obs, title={OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot}, author={Zhu, Junhan and Wang, Hesong and Su, Mingluo and Wang, Zefang and Wang, Huan}, journal={arXiv preprint arXiv:2510.06751}, year={2025} } ```