PeterL1n commited on
Commit
220d1de
1 Parent(s): 09f95d8
Files changed (1) hide show
  1. app.py +3 -5
app.py CHANGED
@@ -35,12 +35,10 @@ from transformers import CLIPFeatureExtractor
35
  safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device)
36
  feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
37
 
38
- def check_nsfw_images(
39
- images: list[Image.Image],
40
- ) -> tuple[list[Image.Image], list[bool]]:
41
  safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
42
  has_nsfw_concepts = safety_checker(images=[images], clip_input=safety_checker_input.pixel_values.to(device))
43
- return images, has_nsfw_concepts
44
 
45
  # Function
46
  @spaces.GPU(enable_queue=True)
@@ -73,7 +71,7 @@ def generate_image(prompt, base, motion, step, progress=gr.Progress()):
73
 
74
  output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step, callback=progress_callback, callback_steps=1)
75
 
76
- images, has_nsfw_concepts = check_nsfw_images([output.frames[0][0]])
77
  if has_nsfw_concepts[0]:
78
  gr.Warning("NSFW content detected.")
79
  return None
 
35
  safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device)
36
  feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
37
 
38
+ def check_nsfw_images(images: list[Image.Image]) -> list[bool]:
 
 
39
  safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
40
  has_nsfw_concepts = safety_checker(images=[images], clip_input=safety_checker_input.pixel_values.to(device))
41
+ return has_nsfw_concepts
42
 
43
  # Function
44
  @spaces.GPU(enable_queue=True)
 
71
 
72
  output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step, callback=progress_callback, callback_steps=1)
73
 
74
+ has_nsfw_concepts = check_nsfw_images([output.frames[0][0]])
75
  if has_nsfw_concepts[0]:
76
  gr.Warning("NSFW content detected.")
77
  return None