OBS-Diff-SD3.5-Large

OBS-Diff Structured Pruning for Stable Diffusion 3.5-Large

<div style=" display: flex; flex-wrap: wrap; align-items: flex-start; gap: 20px; border: 1px solid #e0e0e0; padding: 20px; border-radius: 10px; margin-bottom: 20px; background-color: #fff; "> <div style="flex: 1; min-width: 280px; max-width: 100%;"> <img src="teaser.jpg" alt="OBS-Diff" style="width: 100%; height: auto; border-radius: 5px;" /> </div> <div style="flex: 2; min-width: 300px;"> <h4 style="margin-top: 0;">✂️ <a href="https://alrightlone.github.io/OBS-Diff-Webpage/">OBS-Diff: Accurate Pruning for Diffusion Models in One-Shot</a></h4> <p> <em><b>Junhan Zhu</b>, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang*</em> <br> <a href="https://arxiv.org/abs/2510.06751"><img src="https://img.shields.io/badge/Preprint-arXiv-b31b1b.svg?style=flat-square"></a> <a href="https://github.com/Alrightlone/OBS-Diff"><img src="https://img.shields.io/github/stars/Alrightlone/OBS-Diff?style=flat-square&logo=github"></a> </p> <p> The <b>first training-free, one-shot pruning framework</b> for Diffusion Models, supporting diverse architectures and pruning granularities. Uses Optimal Brain Surgeon (OBS) to achieve <b>SOTA</b> compression with high generative quality. </p> </div> </div>

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

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