radames commited on
Commit
0290645
1 Parent(s): 595db11

add safety checker

Browse files
Files changed (2) hide show
  1. app.py +55 -29
  2. safety_checker.py +137 -0
app.py CHANGED
@@ -12,7 +12,7 @@ from PIL import Image
12
  import gradio as gr
13
  import time
14
  from safetensors.torch import load_file
15
- from sfast.compilers.diffusion_pipeline_compiler import compile, CompilationConfig
16
 
17
  # Constants
18
  BASE = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -28,14 +28,16 @@ CHECKPOINT = "sdxl_lightning_2step_unet.safetensors"
28
  # }
29
 
30
 
31
- TORCH_COMPILE = os.environ.get("TORCH_COMPILE", "0") == "1"
 
32
  # check if MPS is available OSX only M1/M2/M3 chips
33
 
34
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
35
  torch_device = device
36
  torch_dtype = torch.float16
37
 
38
- print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 
39
  print(f"device: {device}")
40
 
41
 
@@ -44,34 +46,60 @@ unet = UNet2DConditionModel.from_config(BASE, subfolder="unet").to(
44
  )
45
  unet.load_state_dict(load_file(hf_hub_download(REPO, CHECKPOINT), device="cuda"))
46
  pipe = StableDiffusionXLPipeline.from_pretrained(
47
- BASE, unet=unet, torch_dtype=torch.float16, variant="fp16"
48
  ).to("cuda")
49
 
50
  # Ensure sampler uses "trailing" timesteps.
51
  pipe.scheduler = EulerDiscreteScheduler.from_config(
52
  pipe.scheduler.config, timestep_spacing="trailing"
53
  )
54
-
55
  pipe.set_progress_bar_config(disable=True)
56
- config = CompilationConfig.Default()
57
- try:
58
- import xformers
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59
 
60
- config.enable_xformers = True
61
- except ImportError:
62
- print("xformers not installed, skip")
63
- try:
64
- import triton
65
 
66
- config.enable_triton = True
67
- except ImportError:
68
- print("Triton not installed, skip")
69
- # CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
70
- # But it can increase the amount of GPU memory used.
71
- # For StableVideoDiffusionPipeline it is not needed.
72
- config.enable_cuda_graph = True
73
 
74
- pipe = compile(pipe, config)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
 
76
 
77
  def predict(prompt, seed=1231231):
@@ -87,14 +115,12 @@ def predict(prompt, seed=1231231):
87
  output_type="pil",
88
  )
89
  print(f"Pipe took {time.time() - last_time} seconds")
90
- nsfw_content_detected = (
91
- results.nsfw_content_detected[0]
92
- if "nsfw_content_detected" in results
93
- else False
94
- )
95
- if nsfw_content_detected:
96
- gr.Warning("NSFW content detected.")
97
- return Image.new("RGB", (512, 512))
98
  return results.images[0]
99
 
100
 
 
12
  import gradio as gr
13
  import time
14
  from safetensors.torch import load_file
15
+
16
 
17
  # Constants
18
  BASE = "stabilityai/stable-diffusion-xl-base-1.0"
 
28
  # }
29
 
30
 
31
+ SFAST_COMPILE = os.environ.get("SFAST_COMPILE", "0") == "1"
32
+ SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
33
  # check if MPS is available OSX only M1/M2/M3 chips
34
 
35
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
36
  torch_device = device
37
  torch_dtype = torch.float16
38
 
39
+ print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
40
+ print(f"SFAST_COMPILE: {SFAST_COMPILE}")
41
  print(f"device: {device}")
42
 
43
 
 
46
  )
47
  unet.load_state_dict(load_file(hf_hub_download(REPO, CHECKPOINT), device="cuda"))
48
  pipe = StableDiffusionXLPipeline.from_pretrained(
49
+ BASE, unet=unet, torch_dtype=torch.float16, variant="fp16", safety_checker=False
50
  ).to("cuda")
51
 
52
  # Ensure sampler uses "trailing" timesteps.
53
  pipe.scheduler = EulerDiscreteScheduler.from_config(
54
  pipe.scheduler.config, timestep_spacing="trailing"
55
  )
 
56
  pipe.set_progress_bar_config(disable=True)
57
+ if SAFETY_CHECKER:
58
+ from safety_checker import StableDiffusionSafetyChecker
59
+ from transformers import CLIPFeatureExtractor
60
+
61
+ safety_checker = StableDiffusionSafetyChecker.from_pretrained(
62
+ "CompVis/stable-diffusion-safety-checker"
63
+ ).to(device)
64
+ feature_extractor = CLIPFeatureExtractor.from_pretrained(
65
+ "openai/clip-vit-base-patch32"
66
+ )
67
+
68
+ def check_nsfw_images(
69
+ images: list[Image.Image],
70
+ ) -> tuple[list[Image.Image], list[bool]]:
71
+ safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
72
+ has_nsfw_concepts = safety_checker(
73
+ images=[images],
74
+ clip_input=safety_checker_input.pixel_values.to(torch_device),
75
+ )
76
 
77
+ return images, has_nsfw_concepts
 
 
 
 
78
 
 
 
 
 
 
 
 
79
 
80
+ if SFAST_COMPILE:
81
+ from sfast.compilers.diffusion_pipeline_compiler import compile, CompilationConfig
82
+
83
+ # sfast compilation
84
+ config = CompilationConfig.Default()
85
+ try:
86
+ import xformers
87
+
88
+ config.enable_xformers = True
89
+ except ImportError:
90
+ print("xformers not installed, skip")
91
+ try:
92
+ import triton
93
+
94
+ config.enable_triton = True
95
+ except ImportError:
96
+ print("Triton not installed, skip")
97
+ # CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
98
+ # But it can increase the amount of GPU memory used.
99
+ # For StableVideoDiffusionPipeline it is not needed.
100
+ config.enable_cuda_graph = True
101
+
102
+ pipe = compile(pipe, config)
103
 
104
 
105
  def predict(prompt, seed=1231231):
 
115
  output_type="pil",
116
  )
117
  print(f"Pipe took {time.time() - last_time} seconds")
118
+ if SAFETY_CHECKER:
119
+ images, has_nsfw_concepts = check_nsfw_images(results.images)
120
+ if any(has_nsfw_concepts):
121
+ gr.Warning("NSFW content detected.")
122
+ return Image.new("RGB", (512, 512))
123
+ return images[0]
 
 
124
  return results.images[0]
125
 
126
 
safety_checker.py ADDED
@@ -0,0 +1,137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
21
+ def cosine_distance(image_embeds, text_embeds):
22
+ normalized_image_embeds = nn.functional.normalize(image_embeds)
23
+ normalized_text_embeds = nn.functional.normalize(text_embeds)
24
+ return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
25
+
26
+
27
+ class StableDiffusionSafetyChecker(PreTrainedModel):
28
+ config_class = CLIPConfig
29
+
30
+ _no_split_modules = ["CLIPEncoderLayer"]
31
+
32
+ def __init__(self, config: CLIPConfig):
33
+ super().__init__(config)
34
+
35
+ self.vision_model = CLIPVisionModel(config.vision_config)
36
+ self.visual_projection = nn.Linear(
37
+ config.vision_config.hidden_size, config.projection_dim, bias=False
38
+ )
39
+
40
+ self.concept_embeds = nn.Parameter(
41
+ torch.ones(17, config.projection_dim), requires_grad=False
42
+ )
43
+ self.special_care_embeds = nn.Parameter(
44
+ torch.ones(3, config.projection_dim), requires_grad=False
45
+ )
46
+
47
+ self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
48
+ self.special_care_embeds_weights = nn.Parameter(
49
+ torch.ones(3), requires_grad=False
50
+ )
51
+
52
+ @torch.no_grad()
53
+ def forward(self, clip_input, images):
54
+ pooled_output = self.vision_model(clip_input)[1] # pooled_output
55
+ image_embeds = self.visual_projection(pooled_output)
56
+
57
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
58
+ special_cos_dist = (
59
+ cosine_distance(image_embeds, self.special_care_embeds)
60
+ .cpu()
61
+ .float()
62
+ .numpy()
63
+ )
64
+ cos_dist = (
65
+ cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()
66
+ )
67
+
68
+ result = []
69
+ batch_size = image_embeds.shape[0]
70
+ for i in range(batch_size):
71
+ result_img = {
72
+ "special_scores": {},
73
+ "special_care": [],
74
+ "concept_scores": {},
75
+ "bad_concepts": [],
76
+ }
77
+
78
+ # increase this value to create a stronger `nfsw` filter
79
+ # at the cost of increasing the possibility of filtering benign images
80
+ adjustment = 0.0
81
+
82
+ for concept_idx in range(len(special_cos_dist[0])):
83
+ concept_cos = special_cos_dist[i][concept_idx]
84
+ concept_threshold = self.special_care_embeds_weights[concept_idx].item()
85
+ result_img["special_scores"][concept_idx] = round(
86
+ concept_cos - concept_threshold + adjustment, 3
87
+ )
88
+ if result_img["special_scores"][concept_idx] > 0:
89
+ result_img["special_care"].append(
90
+ {concept_idx, result_img["special_scores"][concept_idx]}
91
+ )
92
+ adjustment = 0.01
93
+
94
+ for concept_idx in range(len(cos_dist[0])):
95
+ concept_cos = cos_dist[i][concept_idx]
96
+ concept_threshold = self.concept_embeds_weights[concept_idx].item()
97
+ result_img["concept_scores"][concept_idx] = round(
98
+ concept_cos - concept_threshold + adjustment, 3
99
+ )
100
+ if result_img["concept_scores"][concept_idx] > 0:
101
+ result_img["bad_concepts"].append(concept_idx)
102
+
103
+ result.append(result_img)
104
+
105
+ has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]
106
+
107
+ return has_nsfw_concepts
108
+
109
+ @torch.no_grad()
110
+ def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor):
111
+ pooled_output = self.vision_model(clip_input)[1] # pooled_output
112
+ image_embeds = self.visual_projection(pooled_output)
113
+
114
+ special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds)
115
+ cos_dist = cosine_distance(image_embeds, self.concept_embeds)
116
+
117
+ # increase this value to create a stronger `nsfw` filter
118
+ # at the cost of increasing the possibility of filtering benign images
119
+ adjustment = 0.0
120
+
121
+ special_scores = (
122
+ special_cos_dist - self.special_care_embeds_weights + adjustment
123
+ )
124
+ # special_scores = special_scores.round(decimals=3)
125
+ special_care = torch.any(special_scores > 0, dim=1)
126
+ special_adjustment = special_care * 0.01
127
+ special_adjustment = special_adjustment.unsqueeze(1).expand(
128
+ -1, cos_dist.shape[1]
129
+ )
130
+
131
+ concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment
132
+ # concept_scores = concept_scores.round(decimals=3)
133
+ has_nsfw_concepts = torch.any(concept_scores > 0, dim=1)
134
+
135
+ images[has_nsfw_concepts] = 0.0 # black image
136
+
137
+ return images, has_nsfw_concepts