Upload SDG pipeline + classifier weights

README.md

@@ -1,54 +1,11 @@
----
-library_name: diffusers
-pipeline_tag: text-to-image
-tags:
-  - safety
-  - classifier-guidance
-  - stable-diffusion
-  - plug-and-play
-license: apache-2.0
----
+# Safe Diffusion Guidance (SDG)
 
-# Safe Diffusion Guidance (SDG) — plug-and-play safety layer for Stable Diffusion
+Custom Diffusers pipeline that applies a mid-UNet safety classifier as guidance during denoising.
+- Plug-and-play: works with any Stable Diffusion checkpoint (e.g., SD 1.5).
+- No retraining needed; classifier runs on mid-UNet features.
+- Tunable: `safety_scale`, `mid_fraction`, `safe_class_index`.
 
-**Safe Diffusion Guidance (SDG)** is a *classifier-guided denoising* layer that steers the sampling trajectory away from unsafe content **without retraining** the base model.  
-It works **standalone** with SD 1.4 / 1.5 / 2.1 and **composes** cleanly with ESD/UCE/SLD.
-
-- **Safety signal:** a 5-class mid-UNet feature classifier (classes: `gore, hate, medical, safe, sexual`) trained on (1280×8×8) features.
-- **Controls:** `safety_scale` (strength), `mid_fraction` (fraction of steps guided).
-- **Plug-in:** drop into any SD pipeline, or stack on top of ESD/UCE/SLD.
-- **No retraining:** small gradient nudges to latents during denoising.
-
-> **Note on metrics** (matching our paper): FID/KID are computed vs. _baseline model outputs_ rather than real images; baseline FID/KID are ≈0 by construction.
-
-## Quickstart (SD 1.5)
-
-```python
-import torch
-from diffusers import StableDiffusionPipeline
-
-# 1) Load base SD pipeline (disable default safety checker)
-base = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16,
-    safety_checker=None
-).to("cuda")
-
-# 2) Load SDG custom pipeline from Hub (this repo)
-sdg = StableDiffusionPipeline.from_pretrained(
-    "your-org/safe-diffusion-guidance",
-    custom_pipeline="safe_diffusion_guidance",
-    torch_dtype=torch.float16
-).to("cuda")
-
-img = sdg(
-    base_pipe=base,
-    prompt="portrait photograph, studio light, 85mm, realistic",
-    num_inference_steps=50,
-    guidance_scale=7.5,
-    safety_scale=5.0,        # strength: ~2–8 (Light→Strong)
-    mid_fraction=1.0,        # guide fraction of steps: 0.5, 0.8, 1.0
-    safe_class_index=3       # index of 'safe' in [gore,hate,medical,safe,sexual]
-).images[0]
-
-img.save("sdg_safe_output.png")
+## Install
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt

__init__.py

@@ -1,3 +1,2 @@
-# __init__.py
 from .safe_diffusion_guidance import SafeDiffusionGuidance
 __all__ = ["SafeDiffusionGuidance"]

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.1
+transformers>=4.41
+diffusers>=0.30
+accelerate>=0.30
+safetensors>=0.4
+huggingface_hub>=0.23
+numpy
+Pillow

safe_diffusion_guidance.py

@@ -1,25 +1,30 @@
 # safe_diffusion_guidance.py
 import os
-import torch
 from typing import Optional, List
+
+import torch
+import torch.nn as nn
 from diffusers import DiffusionPipeline, StableDiffusionPipeline
 from diffusers.utils import BaseOutput
 
-# ---- Classifier (unchanged) ----
-import torch.nn as nn
 
-CLASS_NAMES = ['gore', 'hate', 'medical', 'safe', 'sexual']
+# ----------------------------- Classifier ------------------------------------
+CLASS_NAMES = ["gore", "hate", "medical", "safe", "sexual"]
 
 class SafetyClassifier1280(nn.Module):
+    """
+    Safety classifier for mid-UNet features of shape (B, 1280, H, W).
+    Robust to HxW via AdaptiveAvgPool2d((8,8)) before the head.
+    """
     def __init__(self, num_classes: int = 5):
         super().__init__()
         self.pre = nn.AdaptiveAvgPool2d((8, 8))
         self.net = nn.Sequential(
             nn.Conv2d(1280, 512, 3, padding=1),
-            nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(2),
+            nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(2),  # 512x4x4
             nn.Conv2d(512, 256, 3, padding=1),
-            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),
-            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),  # 256x2x2
+            nn.AdaptiveAvgPool2d(1), nn.Flatten(),                         # 256
             nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.3),
             nn.Linear(128, num_classes)
         )
@@ -28,39 +33,18 @@ class SafetyClassifier1280(nn.Module):
     @staticmethod
     def _init_weights(m):
         if isinstance(m, nn.Linear):
-            nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None: nn.init.zeros_(m.bias)
         elif isinstance(m, nn.Conv2d):
-            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
             if m.bias is not None: nn.init.zeros_(m.bias)
         elif isinstance(m, nn.BatchNorm2d):
             nn.init.ones_(m.weight); nn.init.zeros_(m.bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.pre(x)
+        x = self.pre(x)  # (B, 1280, 8, 8)
         return self.net(x)
 
-# ---- NEW: robust path resolution for weights ----
-def _resolve_repo_path(rel_path: str) -> str:
-    """Return an absolute path inside the cached repo; fallback to hf_hub_download."""
-    here = os.path.dirname(__file__)
-    local_path = os.path.join(here, rel_path)
-    if os.path.exists(local_path):
-        return local_path
-    # Fallback: try hub download (works even if code is executed outside repo root)
-    try:
-        from huggingface_hub import hf_hub_download
-        # Best effort to get repo id; default to your public repo if unknown
-        repo_id = getattr(_resolve_repo_path, "_repo_id", None)
-        if repo_id is None:
-            # Diffusers stores name or path in internal dict sometimes:
-            repo_id = getattr(SafeDiffusionGuidance, "__repo_id__", "basimazam/safe-diffusion-guidance")
-        return hf_hub_download(repo_id=repo_id, filename=rel_path)
-    except Exception as e:
-        raise FileNotFoundError(
-            f"Could not find classifier weights at '{rel_path}'. "
-            f"Make sure the file exists in the repo, or pass `classifier_weights=...`. "
-            f"Original error: {e}"
-        )
 
 def load_classifier_1280(
     weights_path: str,
@@ -75,35 +59,66 @@ def load_classifier_1280(
     model.eval()
     return model
 
-def pick_weights_for_pipe(pipe) -> str:
-    # Always use the standard path inside the repo
-    return _resolve_repo_path("classifiers/safety_classifier_1280.pth")
 
+def _here(*paths: str) -> str:
+    return os.path.join(os.path.dirname(__file__), *paths)
+
+
+def pick_weights_path() -> str:
+    """
+    Try common locations; allow env override. Raise if not found.
+    """
+    candidates = [
+        os.getenv("SDG_CLASSIFIER_WEIGHTS", ""),
+        _here("classifiers", "safety_classifier_1280.pth"),
+        _here("safety_classifier_1280.pth"),
+        "classifiers/safety_classifier_1280.pth",
+        "safety_classifier_1280.pth",
+    ]
+    for p in candidates:
+        if p and os.path.exists(p):
+            return p
+    raise FileNotFoundError(
+        "Safety-classifier weights not found. Place 'safety_classifier_1280.pth' "
+        "in repo root or 'classifiers/' (or set SDG_CLASSIFIER_WEIGHTS, or pass "
+        "`classifier_weights=...` to the call())."
+    )
+
+
+# ----------------------------- Pipeline --------------------------------------
 class SDGOutput(BaseOutput):
-    images: List
+    images: List  # list of PIL Images
+
 
 class SafeDiffusionGuidance(DiffusionPipeline):
-    """Pure custom pipeline; loads base SD lazily at runtime."""
+    """
+    Minimal custom pipeline that loads a base Stable Diffusion pipeline on demand
+    and applies mid-UNet classifier-guided denoising for safety.
+    """
 
-    def __init__(self):                    # <-- IMPORTANT: no **kwargs
+    def __init__(self):  # IMPORTANT: no **kwargs (diffusers inspects this)
         super().__init__()
-        self.base_pipe_ = None
-        # Hint for the fallback downloader (optional)
-        try:
-            SafeDiffusionGuidance.__repo_id__ = self.config._name_or_path  # diffusers sets this sometimes
-        except Exception:
-            pass
-
-    def _ensure_base(self, base_pipe, base_model_id, torch_dtype):
+        self.base_pipe_ = None  # lazy cache
+
+    def _ensure_base(
+        self,
+        base_pipe: Optional[StableDiffusionPipeline],
+        base_model_id: str,
+        torch_dtype: torch.dtype,
+    ) -> StableDiffusionPipeline:
         if base_pipe is not None:
             self.base_pipe_ = base_pipe
             return self.base_pipe_
         if self.base_pipe_ is None:
             self.base_pipe_ = StableDiffusionPipeline.from_pretrained(
-                base_model_id, torch_dtype=torch_dtype, safety_checker=None
+                base_model_id,
+                torch_dtype=torch_dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
             ).to(self.device)
         return self.base_pipe_
 
+    @torch.no_grad()
     def __call__(
         self,
         prompt: str,
@@ -111,19 +126,21 @@ class SafeDiffusionGuidance(DiffusionPipeline):
         num_inference_steps: int = 50,
         guidance_scale: float = 7.5,
         safety_scale: float = 5.0,
-        mid_fraction: float = 1.0,
-        safe_class_index: int = 3,
+        mid_fraction: float = 1.0,   # 0..1 fraction of steps to guide
+        safe_class_index: int = 3,   # "safe" in CLASS_NAMES
         classifier_weights: Optional[str] = None,
         base_pipe: Optional[StableDiffusionPipeline] = None,
         base_model_id: str = "runwayml/stable-diffusion-v1-5",
         generator: Optional[torch.Generator] = None,
-        **kwargs
+        **kwargs,
     ) -> SDGOutput:
+
+        # 1) prepare base SD
         base = self._ensure_base(base_pipe, base_model_id, torch_dtype=torch.float16)
         device = getattr(base, "_execution_device", base.device)
-        dtype  = base.unet.dtype
+        dtype = base.unet.dtype
 
-        # text embeds (CFG)
+        # 2) text embeddings (classifier-free guidance)
         tok = base.tokenizer
         max_len = tok.model_max_length
         txt = tok([prompt], padding="max_length", max_length=max_len, return_tensors="pt")
@@ -135,66 +152,72 @@ class SafeDiffusionGuidance(DiffusionPipeline):
         uncond = base.text_encoder(uncond_txt.input_ids.to(device)).last_hidden_state
         cond_embeds = torch.cat([uncond, cond], dim=0)
 
-        # latents
+        # 3) latents
         h = kwargs.pop("height", 512); w = kwargs.pop("width", 512)
-        latents = torch.randn((1, base.unet.in_channels, h // 8, w // 8),
-                              device=device, generator=generator, dtype=dtype)
+        latents = torch.randn(
+            (1, base.unet.in_channels, h // 8, w // 8),
+            device=device, generator=generator, dtype=dtype
+        )
 
         base.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = base.scheduler.timesteps
 
-        # classifier (fp32) — use provided path or default resolved path
-        weights_file = classifier_weights or pick_weights_for_pipe(base)
-        clf = load_classifier_1280(weights_file, device=device, dtype=torch.float32).eval()
+        # 4) classifier (run in fp32)
+        weights_file = classifier_weights or pick_weights_path()
+        clf = load_classifier_1280(weights_file, device=device, dtype=torch.float32)
 
-        # mid-block hook
+        # 5) mid-block hook
         mid = {}
         def hook(_, __, out): mid["feat"] = out[0] if isinstance(out, tuple) else out
         handle = base.unet.mid_block.register_forward_hook(hook)
 
-        base_alpha = 1e-3
+        base_alpha = 1e-3  # step size factor for safety update
+
+        # 6) denoising loop
+        for i, t in enumerate(timesteps):
+            # standard SD forward
+            lat_in = base.scheduler.scale_model_input(latents, t)
+            lat_cat = torch.cat([lat_in, lat_in], dim=0)  # for CFG
+            do_guide = (i / len(timesteps)) <= mid_fraction and safety_scale > 0
+
+            if do_guide:
+                # safety gradient w.r.t latents
+                with torch.enable_grad():
+                    lg = latents.detach().clone().requires_grad_(True)
+                    lin = base.scheduler.scale_model_input(lg, t)
+                    lcat = torch.cat([lin, lin], dim=0)
+
+                    _ = base.unet(lcat, t, encoder_hidden_states=cond_embeds).sample
+                    feat = mid["feat"].detach().float()  # (B*2, 1280, H, W)
+                    logits = clf(feat)
+                    probs = torch.softmax(logits, dim=-1)
+                    unsafe = 1.0 - probs[:, safe_class_index].mean()  # encourage "safe"
 
-        with torch.no_grad():
-            for i, t in enumerate(timesteps):
+                    loss = safety_scale * unsafe
+                    loss.backward()
+
+                    alpha = base_alpha
+                    if hasattr(base.scheduler, "sigmas"):  # DDIM/PNDM/… support
+                        idx = min(i, len(base.scheduler.sigmas) - 1)
+                        alpha = base_alpha * float(base.scheduler.sigmas[idx])
+
+                    latents = (lg - alpha * lg.grad).detach()
+
+                # resume SD denoising with updated latents
                 lat_in = base.scheduler.scale_model_input(latents, t)
                 lat_cat = torch.cat([lat_in, lat_in], dim=0)
 
-                do_guide = (i / len(timesteps)) <= mid_fraction and safety_scale > 0
-                if do_guide:
-                    with torch.enable_grad():
-                        lg = latents.detach().clone().requires_grad_(True)
-                        lin = base.scheduler.scale_model_input(lg, t)
-                        lcat = torch.cat([lin, lin], dim=0)
-
-                        _ = base.unet(lcat, t, encoder_hidden_states=cond_embeds).sample
-                        feat = mid["feat"].detach().float()
-                        logits = clf(feat)
-                        probs  = torch.softmax(logits, dim=-1)
-                        unsafe = 1.0 - probs[:, safe_class_index].mean()
-
-                        loss = safety_scale * unsafe
-                        loss.backward()
-
-                        alpha = base_alpha
-                        if hasattr(base.scheduler, "sigmas"):
-                            idx = min(i, len(base.scheduler.sigmas) - 1)
-                            alpha = base_alpha * float(base.scheduler.sigmas[idx])
-                        latents = (lg - alpha * lg.grad).detach()
-
-                    lat_in = base.scheduler.scale_model_input(latents, t)
-                    lat_cat = torch.cat([lat_in, lat_in], dim=0)
-                    noise_pred = base.unet(lat_cat, t, encoder_hidden_states=cond_embeds).sample
-                else:
-                    noise_pred = base.unet(lat_cat, t, encoder_hidden_states=cond_embeds).sample
-
-                n_uncond, n_text = noise_pred.chunk(2)
-                noise = n_uncond + guidance_scale * (n_text - n_uncond)
-                latents = base.scheduler.step(noise, t, latents).prev_sample
+            noise_pred = base.unet(lat_cat, t, encoder_hidden_states=cond_embeds).sample
+            n_uncond, n_text = noise_pred.chunk(2)
+            noise = n_uncond + guidance_scale * (n_text - n_uncond)
+            latents = base.scheduler.step(noise, t, latents).prev_sample
 
         handle.remove()
-        with torch.no_grad():
-            img = base.decode_latents(latents)
-            img = base.image_processor.postprocess(img, output_type="pil")[0]
-        return SDGOutput(images=[img])
+
+        # 7) decode
+        img = base.decode_latents(latents)
+        pil = base.image_processor.postprocess(img, output_type="pil")[0]
+        return SDGOutput(images=[pil])
+
 
 __all__ = ["SafeDiffusionGuidance"]

utils/__init__.py

@@ -1 +1,4 @@
-from .adaptive_classifiers import SafetyClassifier1280, load_classifier_1280, CLASS_NAMES
+# Namespace init for utils
+from .adaptive_classifiers import (
+    SafetyClassifier1280, load_classifier_1280, CLASS_NAMES
+)