Update app.py

app.py

@@ -4,7 +4,34 @@ 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(self, image, device, dtype):
+
+    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)
+    )
+    return image, has_nsfw_concept
+    
 import os
 model_map = {'Van Gogh' : 'models/vangogh.pt', 
              'Pablo Picasso': 'models/pablopicasso.pt',