# EdgeDiffusion — Distilled Final (647M params)

A pruned and distilled Stable Diffusion 1.5 UNet, reduced from **858M → 647M parameters** (24.6% reduction) while preserving generation quality.

![image](https://cdn-uploads.huggingface.co/production/uploads/698ab6f4af78a06a22666e2d/zoYFOeV7U9rXC1TiMve8K.png)


## Model Details

| | Value |
|---|---|
| Base model | Stable Diffusion v1.5 |
| Pruning method | Iterative structured Taylor pruning (4 rounds × ~7%) |
| Distillation teacher | DreamShaper v8 |
| Distillation steps | 40K constant LR (1e-5) + 10K cosine decay (1e-5 → 1e-6) |
| Loss | L_out (noise MSE) + 0.1 × L_feat (feature MSE) |
| Parameters | 647.2M (vs 858.5M baseline) |
| Dataset | 20K images (DiffusionDB 10K + COCO 2017 10K) |

## Usage

```python
import torch
from diffusers import StableDiffusionPipeline
from pruned_rebuild import create_unet_from_safetensors

# Rebuild the pruned UNet
unet = create_unet_from_safetensors(
    "pruned_unet.safetensors",
    "pruned_unet.config.json"
)
unet = unet.to(dtype=torch.float16, device="cuda")

# Load into SD1.5 pipeline
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    unet=unet,
    torch_dtype=torch.float16,
    safety_checker=None,
    requires_safety_checker=False,
).to("cuda")

image = pipe("a beautiful landscape, 4k", num_inference_steps=30).images[0]
image.save("output.png")
```

## Files

- `pruned_unet.safetensors` — Pruned + distilled UNet weights
- `pruned_unet.config.json` — UNet architecture config (channel dimensions)
- `pruned_rebuild.py` — Script to rebuild the pruned UNet from safetensors

## Pipeline

1. **Iterative Structured Pruning**: 4 rounds of Taylor importance-based channel pruning (~7% per round)
2. **Sensitivity-Guided**: Latent Divergence (LD) sensitivity analysis to protect critical blocks
3. **Knowledge Distillation**: BK-SDM style distillation with DreamShaper v8 teacher
4. **Fine-grained Recovery**: Final 10K steps with cosine LR decay for quality refinement