Mitchins/comfyui-junior-safety

ComfyUI Junior Prompt Safety Classifier (v0.7 Draft)

Status: Draft / Provisional Pre-release (v0.7)
This model repository provides the lightweight prompt safety classifier used by ComfyUI Junior.

Model Architecture

The classifier uses a fine-tuned distilbert-base-uncased transformer backbone combined with a set of multi-width linear regression heads (heads.pt) operating directly on the 768-dimensional [CLS] token representation.

Safety Dimensions & Head Widths

Dimension	Width	Scope / Categories
sexual	3	Mild suggestive, overt sexuality, explicit adult content
nudity	2	Partial nudity, full explicit nudity
violence_gore	2	Combat/action violence, severe gore & injury
disturbing	2	Macabre/dark themes, grotesque imagery
substances	1	Illicit substances, drug paraphernalia
fetish	2	Mild fetish, severe fetish/kink

Usage

import torch
import torch.nn as nn
from transformers import AutoTokenizer, DistilBertModel

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model_id = 'Mitchins/comfyui-junior-safety'

# Load tokenizer and encoder
tokenizer = AutoTokenizer.from_pretrained(model_id)
encoder = DistilBertModel.from_pretrained(model_id).to(device).eval()

# Load multi-width heads
heads_checkpoint = torch.hub.load_state_dict_from_url(
    'https://huggingface.co/Mitchins/comfyui-junior-safety/resolve/main/heads.pt',
    map_location=device,
    weights_only=True
)

dims = heads_checkpoint['dims']
widths = heads_checkpoint['widths']
heads = nn.ModuleDict()

for i, (dim, width) in enumerate(zip(dims, widths)):
    lin = nn.Linear(768, width)
    lin.weight.data = heads_checkpoint['heads'][f'{i}.weight'].to(device)
    lin.bias.data = heads_checkpoint['heads'][f'{i}.bias'].to(device)
    heads[dim] = lin.eval()

# Classify prompt
prompt = 'A cute puppy playing in the garden'
inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)

with torch.no_grad():
    cls_rep = encoder(**inputs).last_hidden_state[:, 0]
    scores = {dim: heads[dim](cls_rep).squeeze(0).cpu().tolist() for dim in dims}

print('Safety scores:', scores)

Intended Use

Designed specifically for inline, low-latency (<5ms on GPU) prompt moderation for kid-safe / family image generation appliances.