Shuang59/Composable-Diffusion

1"""2    file copy from diffusion library from Huggingface: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py3"""4import numpy as np5import torch6import torch.nn as nn7 8from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel9 10from diffusers.utils import logging11 12 13logger = logging.get_logger(__name__)14 15 16def cosine_distance(image_embeds, text_embeds):17    normalized_image_embeds = nn.functional.normalize(image_embeds)18    normalized_text_embeds = nn.functional.normalize(text_embeds)19    return torch.mm(normalized_image_embeds, normalized_text_embeds.T)20 21 22class StableDiffusionSafetyChecker(PreTrainedModel):23    config_class = CLIPConfig24 25    def __init__(self, config: CLIPConfig):26        super().__init__(config)27 28        self.vision_model = CLIPVisionModel(config.vision_config)29        self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False)30 31        self.concept_embeds = nn.Parameter(torch.ones(17, config.projection_dim), requires_grad=False)32        self.special_care_embeds = nn.Parameter(torch.ones(3, config.projection_dim), requires_grad=False)33 34        self.register_buffer("concept_embeds_weights", torch.ones(17))35        self.register_buffer("special_care_embeds_weights", torch.ones(3))36 37    @torch.no_grad()38    def forward(self, clip_input, images):39        pooled_output = self.vision_model(clip_input)[1]  # pooled_output40        image_embeds = self.visual_projection(pooled_output)41 42        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().numpy()43        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().numpy()44 45        result = []46        batch_size = image_embeds.shape[0]47        for i in range(batch_size):48            result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []}49 50            # increase this value to create a stronger `nfsw` filter51            # at the cost of increasing the possibility of filtering benign images52            adjustment = 0.053 54            for concet_idx in range(len(special_cos_dist[0])):55                concept_cos = special_cos_dist[i][concet_idx]56                concept_threshold = self.special_care_embeds_weights[concet_idx].item()57                result_img["special_scores"][concet_idx] = round(concept_cos - concept_threshold + adjustment, 3)58                if result_img["special_scores"][concet_idx] > 0:59                    result_img["special_care"].append({concet_idx, result_img["special_scores"][concet_idx]})60                    adjustment = 0.0161 62            for concet_idx in range(len(cos_dist[0])):63                concept_cos = cos_dist[i][concet_idx]64                concept_threshold = self.concept_embeds_weights[concet_idx].item()65                result_img["concept_scores"][concet_idx] = round(concept_cos - concept_threshold + adjustment, 3)66                if result_img["concept_scores"][concet_idx] > 0:67                    result_img["bad_concepts"].append(concet_idx)68 69            result.append(result_img)70 71        has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]72 73        for idx, has_nsfw_concept in enumerate(has_nsfw_concepts):74            if has_nsfw_concept:75                images[idx] = np.zeros(images[idx].shape)  # black image76 77        if any(has_nsfw_concepts):78            logger.warning(79                "Potential NSFW content was detected in one or more images. A black image will be returned instead."80                " Try again with a different prompt and/or seed."81            )82 83        return images, has_nsfw_concepts84