112 lines
4.5 KiB
Python
112 lines
4.5 KiB
Python
import os
|
|
|
|
import requests
|
|
import torch
|
|
from torch import nn
|
|
from transformers import CLIPImageProcessor, CLIPConfig, CLIPVisionModel, PreTrainedModel
|
|
|
|
from .shori import download_model
|
|
|
|
# 优先使用GPU进行处理
|
|
if torch.cuda.is_available():
|
|
device = 'cuda'
|
|
elif torch.backends.mps.is_available():
|
|
device = 'mps'
|
|
else:
|
|
device = 'cpu'
|
|
|
|
def cal_cos_dist(a, b):
|
|
'''计算余弦相似度'''
|
|
return nn.functional.cosine_similarity(a.unsqueeze(1), b, dim=-1)
|
|
|
|
# CompVis模型存储目录
|
|
compvis_model_dir = os.path.join(os.path.dirname(__file__), 'compvis_model')
|
|
|
|
def download_compvis_model():
|
|
'''下载CompVis模型文件'''
|
|
if not os.path.exists(compvis_model_dir):
|
|
os.makedirs(compvis_model_dir)
|
|
|
|
# 下载模型文件
|
|
root_url = 'https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/refs%2Fpr%2F41/'
|
|
model_file = os.path.join(compvis_model_dir, 'model.safetensors')
|
|
model_url = root_url + 'model.safetensors'
|
|
|
|
download_model(model_url, model_file, 'CompVis')
|
|
|
|
|
|
# 下载配置文件
|
|
for config_json in ['config.json', 'preprocessor_config.json']:
|
|
config_url = root_url + config_json
|
|
config_file = os.path.join(compvis_model_dir, config_json)
|
|
try:
|
|
with requests.get(config_url) as response:
|
|
open(config_file, 'wb').write(response.content)
|
|
except Exception as e:
|
|
raise RuntimeError(f'CompVis配置文件下载失败: {e}') from e
|
|
|
|
class CLIPSafetyChecker(PreTrainedModel):
|
|
'''
|
|
CompVis使用的NSFW检查器
|
|
改编自:
|
|
- https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py
|
|
- https://github.com/Acly/comfyui-tooling-nodes/blob/main/nsfw.py
|
|
'''
|
|
config_class = CLIPConfig
|
|
_no_split_modules = ['CLIPEncoderLayer']
|
|
|
|
def __init__(self, config: CLIPConfig):
|
|
super().__init__(config)
|
|
proj_dim = config.projection_dim
|
|
|
|
self.vision_model = CLIPVisionModel(config.vision_config)
|
|
self.visual_projection = nn.Linear(config.vision_config.hidden_size, proj_dim, bias=False)
|
|
|
|
self.concept_embeds = nn.Parameter(torch.ones(17, proj_dim), requires_grad=False)
|
|
self.special_care_embeds = nn.Parameter(torch.ones(3, proj_dim), requires_grad=False)
|
|
self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
|
|
self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
|
|
|
|
# 在 Transformers 4.57.3 版本之后，模型需要执行 post_init
|
|
if hasattr(self, 'post_init'):
|
|
self.post_init()
|
|
|
|
def forward(self, img):
|
|
'''前向传播计算NSFW分数'''
|
|
with torch.no_grad():
|
|
image_batch = self.vision_model(img.to(device))[1]
|
|
image_embed = self.visual_projection(image_batch)
|
|
sensitivity = 0.14 * 0.5 - 0.1
|
|
|
|
# 计算特殊关注内容的余弦距离
|
|
special_cos_dist = cal_cos_dist(image_embed, self.special_care_embeds)
|
|
special_score = special_cos_dist - self.special_care_embeds_weights.unsqueeze(0) + sensitivity
|
|
|
|
# 计算一般概念的余弦距离
|
|
cos_dist = cal_cos_dist(image_embed, self.concept_embeds)
|
|
concept_score = cos_dist - self.concept_embeds_weights.unsqueeze(0) + sensitivity
|
|
concept_score += torch.any(special_score > 0, dim=1, keepdim=True) * 0.01
|
|
|
|
# 使用sigmoid函数获得0～1的分数值
|
|
sigmoid_score = nn.functional.sigmoid(concept_score * 100)
|
|
# 返回每个图像的最高分数
|
|
return [float(max(cs)) for cs in sigmoid_score]
|
|
|
|
class CompvisDetector:
|
|
'''CompVis NSFW检测模型'''
|
|
def __init__(self):
|
|
# 检查模型文件是否存在，不存在则下载
|
|
for file_name in ['model.safetensors', 'config.json', 'preprocessor_config.json']:
|
|
model_file = os.path.join(compvis_model_dir, file_name)
|
|
if not os.path.exists(model_file):
|
|
download_compvis_model()
|
|
break
|
|
self.feature_extractor = CLIPImageProcessor.from_pretrained(compvis_model_dir)
|
|
self.safety_checker = CLIPSafetyChecker.from_pretrained(compvis_model_dir).to(device)
|
|
|
|
def cal_nsfw_score(self, image):
|
|
'''计算图像的NSFW分数'''
|
|
image = image.movedim(-1, 1) # 调整维度顺序
|
|
clip_input = self.feature_extractor(image, do_rescale=False, return_tensors='pt')
|
|
score = self.safety_checker(clip_input.pixel_values)
|
|
return score |