Update app.py

app.py

@@ -3,6 +3,8 @@ import torch
 from finetuning import FineTunedModel
 from StableDiffuser import StableDiffuser
 from train import train
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
 import os
 model_map = {'Van Gogh' : 'models/vangogh.pt', 
              'Pablo Picasso': 'models/pablopicasso.pt',
@@ -229,7 +231,10 @@ class Demo:
             n_steps=50,
             generator=generator
         )
-        images = self.diffuser.run_safety_checker(images=images, device='cuda')
+        images, has_nsfw_concept = StableDiffusionSafetyChecker(
+                images=images
+            )
+        
         orig_image = images[0][0]
 
         torch.cuda.empty_cache()
@@ -243,7 +248,9 @@ class Demo:
                 n_steps=50,
                 generator=generator
             )
-        images = self.diffuser.run_safety_checker(images=images, device='cuda')
+        images, has_nsfw_concept = StableDiffusionSafetyChecker(
+                images=images
+            )
         edited_image = images[0][0]
 
         del finetuner