README.md

metadata

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: 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.

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 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.

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:

@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}
}