Upload folder using huggingface_hub

safe_diffusion_guidance.py

@@ -3,17 +3,69 @@ import torch
 from typing import Optional, List
 from diffusers import DiffusionPipeline, StableDiffusionPipeline
 from diffusers.utils import BaseOutput
-# at the top of safe_diffusion_guidance.py
 
-import os, sys
-FILE_DIR = os.path.dirname(__file__)
-UTILS_DIR = os.path.join(FILE_DIR, "utils")
-if UTILS_DIR not in sys.path:
-    sys.path.insert(0, UTILS_DIR)
+# utils/adaptive_classifiers.py
+import torch
+import torch.nn as nn
+from typing import Optional
+
+CLASS_NAMES = ['gore', 'hate', 'medical', 'safe', 'sexual']
 
-# from adaptive_classifiers import load_classifier_1280, pick_weights_for_pipe
+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"
 
-from utils.adaptive_classifiers import load_classifier_1280, pick_weights_for_pipe
 
 class SDGOutput(BaseOutput):
     images: List