improve

app.py

@@ -9,11 +9,14 @@ from PIL import Image
 import time
 import psutil
 import random
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 
 
 start_time = time.time()
 current_steps = 25
 
+SAFETY_CHECKER = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16)
+
 
 class Model:
     def __init__(self, name, path=""):
@@ -22,12 +25,12 @@ class Model:
 
         if path != "":
             self.pipe_t2i = StableDiffusionPipeline.from_pretrained(
-                path, torch_dtype=torch.float16
+                path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER
             )
             self.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config(
                 self.pipe_t2i.scheduler.config
             )
-            self.pipe_i2i = StableDiffusionImg2ImgPipeline(**self.pipe_t2i.components)
+            self.pipe_i2i = StableDiffusionImg2ImgPipeline(**self.pipe_t2i.components, safety_checker=SAFETY_CHECKER)
         else:
             self.pipe_t2i = None
             self.pipe_i2i = None
@@ -36,9 +39,9 @@ class Model:
 models = [
     Model("2.2", "darkstorm2150/Protogen_v2.2_Official_Release"),
     Model("3.4", "darkstorm2150/Protogen_x3.4_Official_Release"),
-#    Model("5.3", "darkstorm2150/Protogen_v5.3_Official_Release"),
-#    Model("5.8", "darkstorm2150/Protogen_x5.8_Official_Release"),
-#    Model("Dragon", "darkstorm2150/Protogen_Dragon_Official_Release"),
+    Model("5.3", "darkstorm2150/Protogen_v5.3_Official_Release"),
+    Model("5.8", "darkstorm2150/Protogen_x5.8_Official_Release"),
+    Model("Dragon", "darkstorm2150/Protogen_Dragon_Official_Release"),
 ]
 
 MODELS = {m.name: m for m in models}
@@ -46,11 +49,6 @@ MODELS = {m.name: m for m in models}
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
 
-# if torch.cuda.is_available():
-#  pipe = pipe.to("cuda")
-#  pipe.enable_xformers_memory_efficient_attention()
-
-
 def error_str(error, title="Error"):
     return (
         f"""#### {title}