Upload SDG pipeline + classifier weights

safe_diffusion_guidance.py

@@ -48,14 +48,48 @@ class SafetyClassifier1280(nn.Module):
 
 def load_classifier_1280(
     weights_path: str,
-    device: Optional[torch.device] = None,
+    device=None,
     dtype: torch.dtype = torch.float32
 ) -> SafetyClassifier1280:
+    """
+    Robustly load classifier weights saved either as a pure state dict or as a
+    training checkpoint (with 'model_state_dict'), handling key prefixes and
+    PyTorch 2.6's weights_only behavior.
+    """
     model = SafetyClassifier1280().to(device or "cpu", dtype=dtype)
-    state = torch.load(weights_path, map_location="cpu")
+
+    # --- load robustly (PyTorch 2.6 & older) ---
+    try:
+        state = torch.load(weights_path, map_location="cpu", weights_only=False)
+    except TypeError:
+        state = torch.load(weights_path, map_location="cpu")
+    except Exception:
+        # Allowlist common numpy globals if needed for old pickles
+        try:
+            import numpy as np
+            from torch.serialization import add_safe_globals
+            add_safe_globals([np.dtype, np.number])
+        except Exception:
+            pass
+        state = torch.load(weights_path, map_location="cpu", weights_only=False)
+
+    # --- unwrap training checkpoint ---
     if isinstance(state, dict) and "model_state_dict" in state:
         state = state["model_state_dict"]
-    model.load_state_dict(state, strict=True)
+
+    # --- normalize keys: 'module.' -> '', 'model.' -> 'net.' ---
+    norm_state = {}
+    for k, v in state.items():
+        k = k.replace("module.", "")
+        k = k.replace("model.", "net.")
+        norm_state[k] = v
+
+    # --- load tolerating tiny diffs (e.g., extra keys from older heads) ---
+    missing, unexpected = model.load_state_dict(norm_state, strict=False)
+    if missing or unexpected:
+        print(f"[SDG] Loaded classifier with relaxed strictness. "
+              f"Missing: {len(missing)}, Unexpected: {len(unexpected)}")
+
     model.eval()
     return model