amazonaws-la commited on
Commit
e6a2a6b
1 Parent(s): 07c0e5d

Create safety_checker.py

Browse files
Files changed (1) hide show
  1. safety_checker.py +106 -0
safety_checker.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import numpy as np
16
+ import torch
17
+ import torch.nn as nn
18
+ from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel
19
+
20
+ def cosine_distance(image_embeds, text_embeds):
21
+ normalized_image_embeds = nn.functional.normalize(image_embeds)
22
+ normalized_text_embeds = nn.functional.normalize(text_embeds)
23
+ return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
24
+
25
+
26
+ class StableDiffusionSafetyChecker(PreTrainedModel):
27
+ config_class = CLIPConfig
28
+
29
+ _no_split_modules = ["CLIPEncoderLayer"]
30
+
31
+ def __init__(self, config: CLIPConfig):
32
+ super().__init__(config)
33
+
34
+ self.vision_model = CLIPVisionModel(config.vision_config)
35
+ self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False)
36
+
37
+ self.concept_embeds = nn.Parameter(torch.ones(17, config.projection_dim), requires_grad=False)
38
+ self.special_care_embeds = nn.Parameter(torch.ones(3, config.projection_dim), requires_grad=False)
39
+
40
+ self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
41
+ self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
42
+
43
+ @torch.no_grad()
44
+ def forward(self, clip_input, images):
45
+ pooled_output = self.vision_model(clip_input)[1] # pooled_output
46
+ image_embeds = self.visual_projection(pooled_output)
47
+
48
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
49
+ special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().float().numpy()
50
+ cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()
51
+
52
+ result = []
53
+ batch_size = image_embeds.shape[0]
54
+ for i in range(batch_size):
55
+ result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []}
56
+
57
+ # increase this value to create a stronger `nfsw` filter
58
+ # at the cost of increasing the possibility of filtering benign images
59
+ adjustment = 0.0
60
+
61
+ for concept_idx in range(len(special_cos_dist[0])):
62
+ concept_cos = special_cos_dist[i][concept_idx]
63
+ concept_threshold = self.special_care_embeds_weights[concept_idx].item()
64
+ result_img["special_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3)
65
+ if result_img["special_scores"][concept_idx] > 0:
66
+ result_img["special_care"].append({concept_idx, result_img["special_scores"][concept_idx]})
67
+ adjustment = 0.01
68
+
69
+ for concept_idx in range(len(cos_dist[0])):
70
+ concept_cos = cos_dist[i][concept_idx]
71
+ concept_threshold = self.concept_embeds_weights[concept_idx].item()
72
+ result_img["concept_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3)
73
+ if result_img["concept_scores"][concept_idx] > 0:
74
+ result_img["bad_concepts"].append(concept_idx)
75
+
76
+ result.append(result_img)
77
+
78
+ has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]
79
+
80
+ return has_nsfw_concepts
81
+
82
+ @torch.no_grad()
83
+ def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor):
84
+ pooled_output = self.vision_model(clip_input)[1] # pooled_output
85
+ image_embeds = self.visual_projection(pooled_output)
86
+
87
+ special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds)
88
+ cos_dist = cosine_distance(image_embeds, self.concept_embeds)
89
+
90
+ # increase this value to create a stronger `nsfw` filter
91
+ # at the cost of increasing the possibility of filtering benign images
92
+ adjustment = 0.0
93
+
94
+ special_scores = special_cos_dist - self.special_care_embeds_weights + adjustment
95
+ # special_scores = special_scores.round(decimals=3)
96
+ special_care = torch.any(special_scores > 0, dim=1)
97
+ special_adjustment = special_care * 0.01
98
+ special_adjustment = special_adjustment.unsqueeze(1).expand(-1, cos_dist.shape[1])
99
+
100
+ concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment
101
+ # concept_scores = concept_scores.round(decimals=3)
102
+ has_nsfw_concepts = torch.any(concept_scores > 0, dim=1)
103
+
104
+ images[has_nsfw_concepts] = 0.0 # black image
105
+
106
+ return images, has_nsfw_concepts