File size: 2,562 Bytes
b8877ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | # 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"
|