Upload SDG pipeline + classifier weights

safe_diffusion_guidance.py

@@ -12,44 +12,25 @@ from typing import Optional
 
 CLASS_NAMES = ['gore', 'hate', 'medical', 'safe', 'sexual']
 
-class AdaptiveClassifier1280(nn.Module):
-    """
-    Same CNN topology you trained (keys start with 'model.*').
-    Input (B,1280,H,W) -> AdaptiveAvgPool2d(8,8) -> conv stack -> head
-    """
+class SafetyClassifier1280(nn.Module):
     def __init__(self, num_classes: int = 5):
         super().__init__()
         self.pre = nn.AdaptiveAvgPool2d((8, 8))
-        # Keep the attribute name 'model' to match the checkpoint keys.
-        self.model = nn.Sequential(
-            nn.Conv2d(1280, 512, kernel_size=3, padding=1),
-            nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(2),  # (512,4,4)
-            nn.Dropout2d(0.1),
-
-            nn.Conv2d(512, 256, kernel_size=3, padding=1),
-            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),  # (256,2,2)
-            nn.Dropout2d(0.1),
-
-            nn.AdaptiveAvgPool2d(1),  # -> (256,1,1)
-            nn.Flatten(),             # -> (256,)
-            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.5),
+        self.model = nn.Sequential(            # <--- use "model" to match checkpoint
+            nn.Conv2d(1280, 512, 3, padding=1),
+            nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(2),
+            nn.Conv2d(512, 256, 3, padding=1),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(2),
+            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
+            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.3),
             nn.Linear(128, num_classes)
         )
-        self.apply(self._init)
-
-    @staticmethod
-    def _init(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):
-        x = self.pre(x)  # (B,1280,8,8)
-        return self.model(x)
+        self.apply(self._init_weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pre(x)
+        return self.model(x)                   # <--- forward through "model"
+
 
 def _find_weights_path() -> str:
     # 1) explicit env; 2) repo root file; 3) classifiers/ subdir
@@ -64,27 +45,12 @@ def _find_weights_path() -> str:
         "or pass `classifier_weights=...` to the pipeline call."
     )
 
-def load_classifier_1280(
-    weights_path: Optional[str],
-    device: torch.device,
-    dtype: torch.dtype = torch.float32
-) -> AdaptiveClassifier1280:
-    path = weights_path or _find_weights_path()
-    ckpt = torch.load(path, map_location="cpu", weights_only=False)
-
-    # Extract actual state dict
-    if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
-        state = ckpt["model_state_dict"]
-    elif isinstance(ckpt, dict) and any(k.startswith("model.") for k in ckpt.keys()):
-        state = ckpt
-    else:
-        # Fallback: allow whole-object saves (only if trusted)
-        state = ckpt
-
-    model = AdaptiveClassifier1280().to(device=device, dtype=torch.float32)  # keep classifier in fp32
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing or unexpected:
-        print(f"[SDG] load_state_dict: missing={missing[:4]}... ({len(missing)}), unexpected={unexpected[:4]}... ({len(unexpected)})")
+def load_classifier_1280(weights_path: str, device=None, dtype=torch.float32):
+    model = SafetyClassifier1280().to(device or "cpu", dtype=dtype)
+    state = torch.load(weights_path, map_location="cpu", weights_only=False)
+    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