Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,25 @@
+# Python
+__pycache__/
+*.pyc
+.venv/
+.env
+.envrc
+
+# OS
+.DS_Store
+Thumbs.db
+
+# HF cache
+~/.cache/huggingface/
+*.lock
+
+# Models & artifacts (we use LFS for these)
+*.pt
+*.pth
+*.bin
+*.safetensors
+*.onnx
+
+# Misc
+*.log
+outputs/

README.md

@@ -0,0 +1,54 @@
+---
+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) — plug-and-play safety layer for Stable Diffusion
+
+**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")

__init__.py

@@ -0,0 +1,2 @@
+# Makes the repo importable if someone clones it as a package.
+__all__ = ["safe_diffusion_guidance"]

model_index.json

@@ -0,0 +1,11 @@
+{
+  "library_name": "diffusers",
+  "pipeline": "StableDiffusionPipeline",
+  "tags": [
+    "safety",
+    "classifier-guidance",
+    "stable-diffusion",
+    "plug-and-play"
+  ],
+  "inference": "Use `custom_pipeline='safe_diffusion_guidance'` and pass your base SD pipeline via `base_pipe=...`."
+}

safe_diffusion_guidance.py

@@ -0,0 +1,143 @@
+# safe_diffusion_guidance.py
+import torch
+from typing import Optional, List
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+from diffusers.utils import BaseOutput
+
+from utils.adaptive_classifiers import (
+    load_classifier_1280, pick_weights_for_pipe
+)
+
+class SDGOutput(BaseOutput):
+    images: List
+
+class SafeDiffusionGuidancePipeline(StableDiffusionPipeline):
+    """
+    Plug-and-play safety guidance for Stable Diffusion.
+    - If `base_pipe` is None, we auto-load the base SD checkpoint specified by `base_model_id`.
+    - Classifier weights are shipped in this repo; no extra installs required.
+
+    Typical use:
+      sdg = DiffusionPipeline.from_pretrained("your-org/safe-diffusion-guidance",
+                                              custom_pipeline="safe_diffusion_guidance",
+                                              torch_dtype=torch.float16).to("cuda")
+      out = sdg(prompt="...", base_model_id="runwayml/stable-diffusion-v1-5",
+                safety_scale=5.0, mid_fraction=1.0, safe_class_index=3)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.base_pipe_ = None  # lazy-loaded base pipeline cache
+
+    def _ensure_base(self, base_pipe, base_model_id, torch_dtype):
+        if base_pipe is not None:
+            self.base_pipe_ = base_pipe
+            return self.base_pipe_
+        if self.base_pipe_ is None:
+            # lazy load chosen SD checkpoint
+            self.base_pipe_ = StableDiffusionPipeline.from_pretrained(
+                base_model_id,
+                torch_dtype=torch_dtype,
+                safety_checker=None
+            ).to(self.device)
+        return self.base_pipe_
+
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        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,
+        classifier_weights: Optional[str] = None,
+        # new convenience arg (optional):
+        base_pipe: Optional[StableDiffusionPipeline] = None,
+        base_model_id: str = "runwayml/stable-diffusion-v1-5",
+        generator: Optional[torch.Generator] = None,
+        **kwargs
+    ) -> SDGOutput:
+
+        # 0) choose / load base SD
+        base = self._ensure_base(base_pipe, base_model_id, torch_dtype=self.unet.dtype)
+        device = base._execution_device if hasattr(base, "_execution_device") else base.device
+        dtype   = base.unet.dtype
+
+        # 1) Text embeddings (CFG)
+        tok = base.tokenizer
+        max_len = tok.model_max_length
+        text_inputs = tok([prompt], padding="max_length", max_length=max_len, return_tensors="pt")
+        text_embeds = base.text_encoder(text_inputs.input_ids.to(device)).last_hidden_state
+
+        if negative_prompt is not None:
+            uncond_inputs = tok([negative_prompt], padding="max_length", max_length=max_len, return_tensors="pt")
+        else:
+            uncond_inputs = tok([""], padding="max_length", max_length=max_len, return_tensors="pt")
+        uncond_embeds = base.text_encoder(uncond_inputs.input_ids.to(device)).last_hidden_state
+        cond_embeds = torch.cat([uncond_embeds, text_embeds], dim=0)
+
+        # 2) Latent init
+        height = kwargs.pop("height", 512); width = kwargs.pop("width", 512)
+        latents = torch.randn(
+            (1, base.unet.in_channels, height // 8, width // 8),
+            device=device, generator=generator, dtype=dtype
+        )
+
+        base.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = base.scheduler.timesteps
+
+        # 3) Load classifier (fp32) from this repo
+        weights_path = classifier_weights or pick_weights_for_pipe(base)
+        classifier = load_classifier_1280(weights_path, device=device, dtype=torch.float32).eval()
+
+        # 4) Hook mid-block features
+        mid_cache = {}
+        def mid_hook(module, inputs, output):
+            mid_cache["feat"] = output[0] if isinstance(output, tuple) else output
+        h = base.unet.mid_block.register_forward_hook(mid_hook)
+
+        base_alpha = 1e-3  # small step
+
+        with torch.no_grad():
+            for i, t in enumerate(timesteps):
+                latent_in = base.scheduler.scale_model_input(latents, t)
+                latent_pair = torch.cat([latent_in, latent_in], dim=0)
+
+                do_guide = (i / len(timesteps)) <= mid_fraction and safety_scale > 0
+                if do_guide:
+                    with torch.enable_grad():
+                        latents_g = latents.detach().clone().requires_grad_(True)
+                        latent_in_g = base.scheduler.scale_model_input(latents_g, t)
+                        latent_pair_g = torch.cat([latent_in_g, latent_in_g], dim=0)
+
+                        _ = base.unet(latent_pair_g, t, encoder_hidden_states=cond_embeds).sample
+                        feat = mid_cache["feat"].detach().float()
+                        logits = classifier(feat)
+                        probs  = torch.softmax(logits, dim=-1)
+                        unsafe_prob = 1.0 - probs[:, safe_class_index].mean()
+
+                        loss = safety_scale * unsafe_prob
+                        loss.backward()
+
+                        alpha = base_alpha
+                        if hasattr(base.scheduler, "sigmas"):
+                            step_idx = min(i, len(base.scheduler.sigmas) - 1)
+                            alpha = base_alpha * float(base.scheduler.sigmas[step_idx])
+                        latents = (latents_g - alpha * latents_g.grad).detach()
+
+                    latent_in = base.scheduler.scale_model_input(latents, t)
+                    latent_pair = torch.cat([latent_in, latent_in], dim=0)
+                    noise_pred = base.unet(latent_pair, t, encoder_hidden_states=cond_embeds).sample
+                else:
+                    noise_pred = base.unet(latent_pair, t, encoder_hidden_states=cond_embeds).sample
+
+                noise_uncond, noise_text = noise_pred.chunk(2)
+                noise = noise_uncond + guidance_scale * (noise_text - noise_uncond)
+                latents = base.scheduler.step(noise, t, latents).prev_sample
+
+        h.remove()
+        with torch.no_grad():
+            image = base.decode_latents(latents)
+            image = base.image_processor.postprocess(image, output_type="pil")[0]
+        return SDGOutput(images=[image])

scripts/push_to_hub.py

@@ -0,0 +1,18 @@
+from huggingface_hub import create_repo, upload_folder
+import os
+
+ORG_REPO = "basimazam/safe-diffusion-guidance"  
+
+def main():
+    create_repo(ORG_REPO, repo_type="model", exist_ok=True)
+    upload_folder(
+        repo_id=ORG_REPO,
+        folder_path=".",  # current folder
+        repo_type="model",
+        ignore_patterns=[".git/*", "outputs/*", "tests/*__pycache__*"]
+    )
+    print(f"Uploaded to https://huggingface.co/{ORG_REPO}")
+
+if __name__ == "__main__":
+    assert os.path.exists("safe_diffusion_guidance.py")
+    main()

tests/test_classifier.py

@@ -0,0 +1,13 @@
+import torch
+from utils.adaptive_classifiers import SafetyClassifier1280
+
+def test_forward_shape():
+    model = SafetyClassifier1280().eval()
+    x = torch.randn(2, 1280, 8, 8)  # fake mid features
+    with torch.no_grad():
+        y = model(x)
+    assert y.shape == (2, 5), f"Expected (2,5), got {tuple(y.shape)}"
+
+if __name__ == "__main__":
+    test_forward_shape()
+    print("OK: classifier forward shape is (B,5)")

tests/test_pipeline_integration.py

@@ -0,0 +1,24 @@
+import torch
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+
+def test_sdg_minimal():
+    sdg = DiffusionPipeline.from_pretrained(
+        "your-org/safe-diffusion-guidance",
+        custom_pipeline="safe_diffusion_guidance",
+        torch_dtype=torch.float16
+    )
+    sdg = sdg.to("cuda" if torch.cuda.is_available() else "cpu")
+    out = sdg(
+        prompt="test scene",
+        base_model_id="runwayml/stable-diffusion-v1-5",
+        num_inference_steps=2,  # small for CI
+        guidance_scale=5.0,
+        safety_scale=2.0,
+        mid_fraction=0.5,
+        safe_class_index=3
+    )
+    assert len(out.images) == 1
+    print("OK: pipeline end-to-end")
+    
+if __name__ == "__main__":
+    test_sdg_minimal()

utils/__init__.py

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

utils/adaptiv_classifiers.py

@@ -0,0 +1,61 @@
+# utils/adaptive_classifiers.py
+import torch
+import torch.nn as nn
+from typing import Optional
+
+CLASS_NAMES = ['gore', 'hate', 'medical', 'safe', 'sexual']
+
+class SafetyClassifier1280(nn.Module):
+    """
+    Unified safety classifier for mid-UNet features of shape (B, 1280, H, W).
+    Robust to variable 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),    # 512 x 4 x 4
+            nn.Conv2d(512, 256, 3, padding=1),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),    # 256 x 2 x 2
+            nn.AdaptiveAvgPool2d(1), nn.Flatten(),                          # 256
+            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.3),
+            nn.Linear(128, num_classes)
+        )
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            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)  # (B, 1280, 8, 8)
+        return self.net(x)
+
+def load_classifier_1280(
+    weights_path: str,
+    device: Optional[torch.device] = None,
+    dtype: torch.dtype = torch.float32
+) -> SafetyClassifier1280:
+    model = SafetyClassifier1280().to(device or "cpu", dtype=dtype)
+    state = torch.load(weights_path, map_location="cpu")
+    if isinstance(state, dict) and "model_state_dict" in state:
+        state = state["model_state_dict"]
+    model.load_state_dict(state, strict=True)
+    model.eval()
+    return model
+
+def pick_weights_for_pipe(pipe) -> str:
+    """
+    Optional helper: return a default weights file based on the base SD pipeline id.
+    You can also use a single shared file 'classifiers/safety_classifier_1280.pth'.
+    """
+    name = str(getattr(pipe, "_internal_dict", {}).get("_name_or_path", "")).lower()
+    # Adjust logic as you like — default to a single shared file:
+    return "classifiers/safety_classifier_1280.pth"

utils/compose.py

@@ -0,0 +1,23 @@
+# utils/compose.py
+import torch
+from diffusers import StableDiffusionPipeline
+from safetensors.torch import load_file
+
+def load_and_patch_sd_pipeline(repo_id, unet_weights_path, dtype=torch.float16, device="cuda"):
+    """
+    Load a base SD pipeline and patch its UNet with ESD/UCE weights.
+    """
+    pipe = StableDiffusionPipeline.from_pretrained(
+        repo_id, torch_dtype=dtype, safety_checker=None
+    ).to(device)
+
+    # Load patch state dict
+    if unet_weights_path.endswith(".safetensors"):
+        patch = load_file(unet_weights_path)
+    else:
+        patch = torch.load(unet_weights_path, map_location="cpu")
+
+    sd = pipe.unet.state_dict()
+    sd.update(patch)
+    pipe.unet.load_state_dict(sd, strict=True)
+    return pipe