ChenHe727/EdgeDiffusion_distilled_final

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

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