File size: 1,742 Bytes
5746967 | 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 | """
Minimal inference example for the CLAMP diffusion checkpoint (clamp_ESD).
Requires: diffusers, transformers, torch, accelerate, safetensors
pip install diffusers transformers torch accelerate safetensors
Note: pipe.load_lora_weights() and unet.load_lora_adapter() do NOT work with this
checkpoint (it was saved via the older UNet-level attn-procs format, not the newer
pipeline-level LoRA format, so the key prefixes don't match what those methods
expect - they will silently attach zero LoRA parameters). Use
unet.load_attn_procs() instead, as below - it is deprecated in newer diffusers
versions but is the one that actually works for this file.
"""
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import snapshot_download
REPO_ID = "nexh98/clamp-checkpoints"
# Base model: Stable Diffusion V1-4 with the paper's ESD erasure already applied
# to the sub-concepts. Using this (rather than vanilla SD v1.4) is required to
# reproduce the paper's reported harmful/benign adaptation numbers.
base_model_dir = snapshot_download(REPO_ID, allow_patterns="diffusion/base_model/*")
lora_dir = snapshot_download(REPO_ID, allow_patterns="diffusion/clamp_ESD/*")
base_model_path = f"{base_model_dir}/diffusion/base_model"
lora_path = f"{lora_dir}/diffusion/clamp_ESD"
pipe = StableDiffusionPipeline.from_pretrained(
base_model_path, torch_dtype=torch.float16, safety_checker=None
)
pipe.unet.load_attn_procs(lora_path, weight_name="pytorch_lora_weights.safetensors")
pipe = pipe.to("cuda")
image = pipe(
"A picture of a Kavri", # one of the paper's merged/harmful concepts (elephant+zebra)
num_inference_steps=25,
guidance_scale=7.5,
).images[0]
image.save("clamp_diffusion_sample.png")
|