README.md

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

<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.
![](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}
}
```