Update app.py

app.py

@@ -22,13 +22,13 @@ def numpy_to_pil(images):
 
     return pil_images
     
-def run_safety_checker(self, image, device, dtype):
+def run_safety_checker(image):
 
     feature_extractor = CLIPFeatureExtractor()
     safety_checker = StableDiffusionSafetyChecker()
     safety_checker_input = feature_extractor(numpy_to_pil(image), return_tensors="pt").to('cuda')
     image, has_nsfw_concept = safety_checker(
-        images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+        images=image, clip_input=safety_checker_input.pixel_values
     )
     return image, has_nsfw_concept
     
@@ -258,9 +258,7 @@ class Demo:
             n_steps=50,
             generator=generator
         )
-        images, has_nsfw_concept = StableDiffusionSafetyChecker(
-                images=images
-            )
+        images, has_nsfw_concept = run_safety_checker(images)
         
         orig_image = images[0][0]
 
@@ -275,9 +273,7 @@ class Demo:
                 n_steps=50,
                 generator=generator
             )
-        images, has_nsfw_concept = StableDiffusionSafetyChecker(
-                images=images
-            )
+        images, has_nsfw_concept = run_safety_checker(images)
         edited_image = images[0][0]
 
         del finetuner