File size: 2,033 Bytes
f9e0f2f eff0aee 4fc9f66 f9e0f2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | # 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.
## 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
|