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',