Update app.py

app.py

@@ -3,34 +3,6 @@ import torch
 from finetuning import FineTunedModel
 from StableDiffuser import StableDiffuser
 from train import train
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-from transformers import CLIPFeatureExtractor
-
-
-def numpy_to_pil(images):
-    """
-    Convert a numpy image or a batch of images to a PIL image.
-    """
-    if images.ndim == 3:
-        images = images[None, ...]
-    images = (images * 255).round().astype("uint8")
-    if images.shape[-1] == 1:
-        # special case for grayscale (single channel) images
-        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
-    else:
-        pil_images = [Image.fromarray(image) for image in images]
-
-    return pil_images
-    
-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
-    )
-    return image, has_nsfw_concept
     
 import os
 model_map = {'Van Gogh' : 'models/vangogh.pt', 
@@ -258,7 +230,7 @@ class Demo:
             n_steps=50,
             generator=generator
         )
-        images, has_nsfw_concept = run_safety_checker(images)
+
         
         orig_image = images[0][0]
 
@@ -273,7 +245,7 @@ class Demo:
                 n_steps=50,
                 generator=generator
             )
-        images, has_nsfw_concept = run_safety_checker(images)
+
         edited_image = images[0][0]
 
         del finetuner