| import numpy as np |
| import torch |
| import torch.nn as nn |
| from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel |
|
|
| from ...utils import logging |
|
|
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class IFSafetyChecker(PreTrainedModel): |
| config_class = CLIPConfig |
|
|
| _no_split_modules = ["CLIPEncoderLayer"] |
|
|
| def __init__(self, config: CLIPConfig): |
| super().__init__(config) |
|
|
| self.vision_model = CLIPVisionModelWithProjection(config.vision_config) |
|
|
| self.p_head = nn.Linear(config.vision_config.projection_dim, 1) |
| self.w_head = nn.Linear(config.vision_config.projection_dim, 1) |
|
|
| @torch.no_grad() |
| def forward(self, clip_input, images, p_threshold=0.5, w_threshold=0.5): |
| image_embeds = self.vision_model(clip_input)[0] |
|
|
| nsfw_detected = self.p_head(image_embeds) |
| nsfw_detected = nsfw_detected.flatten() |
| nsfw_detected = nsfw_detected > p_threshold |
| nsfw_detected = nsfw_detected.tolist() |
|
|
| if any(nsfw_detected): |
| logger.warning( |
| "Potential NSFW content was detected in one or more images. A black image will be returned instead." |
| " Try again with a different prompt and/or seed." |
| ) |
|
|
| for idx, nsfw_detected_ in enumerate(nsfw_detected): |
| if nsfw_detected_: |
| images[idx] = np.zeros(images[idx].shape) |
|
|
| watermark_detected = self.w_head(image_embeds) |
| watermark_detected = watermark_detected.flatten() |
| watermark_detected = watermark_detected > w_threshold |
| watermark_detected = watermark_detected.tolist() |
|
|
| if any(watermark_detected): |
| logger.warning( |
| "Potential watermarked content was detected in one or more images. A black image will be returned instead." |
| " Try again with a different prompt and/or seed." |
| ) |
|
|
| for idx, watermark_detected_ in enumerate(watermark_detected): |
| if watermark_detected_: |
| images[idx] = np.zeros(images[idx].shape) |
|
|
| return images, nsfw_detected, watermark_detected |
|
|
