unified safety_checker

server/main.py

@@ -4,7 +4,8 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi import Request
 import markdown2
-
+from pipelines.utils.safety_checker import SafetyChecker
+from PIL import Image
 import logging
 from config import config, Args
 from connection_manager import ConnectionManager, ServerFullException
@@ -28,6 +29,8 @@ class App:
         self.pipeline = pipeline
         self.app = FastAPI()
         self.conn_manager = ConnectionManager()
+        if self.args.safety_checker:
+            self.safety_checker = SafetyChecker(device=device.type)
         self.init_app()
 
     def init_app(self):
@@ -113,8 +116,14 @@ class App:
                         if params.__dict__ == last_params.__dict__ or params is None:
                             await asyncio.sleep(THROTTLE)
                             continue
-                        last_params = params
+                        last_params: SimpleNamespace = params
                         image = pipeline.predict(params)
+
+                        if self.args.safety_checker:
+                            image, has_nsfw_concept = self.safety_checker(image)
+                            if has_nsfw_concept:
+                                image = None
+
                         if image is None:
                             continue
                         frame = pil_to_frame(image)

server/pipelines/IPcompositionHyperSD15.py

@@ -101,22 +101,13 @@ class Pipeline:
             torch_dtype=torch.float16,
         ).to(device)
 
-        if args.safety_checker:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                model_id,
-                # vae=vae,
-                torch_dtype=torch_dtype,
-                image_encoder=image_encoder,
-                variant="fp16",
-            )
-        else:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                torch_dtype=torch_dtype,
-                image_encoder=image_encoder,
-                variant="fp16",
-            )
+        self.pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            torch_dtype=torch_dtype,
+            image_encoder=image_encoder,
+            variant="fp16",
+        )
 
         self.pipe.load_ip_adapter(
             ip_adapter_model,
@@ -199,13 +190,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/IPcompositionHyperSDXL.py

@@ -106,23 +106,14 @@ class Pipeline:
             torch_dtype=torch.float16,
         ).to(device)
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLPipeline.from_pretrained(
-                model_id,
-                # vae=vae,
-                torch_dtype=torch_dtype,
-                image_encoder=image_encoder,
-                variant="fp16",
-            )
-        else:
-            self.pipe = StableDiffusionXLPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                torch_dtype=torch_dtype,
-                vae=vae,
-                image_encoder=image_encoder,
-                variant="fp16",
-            )
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            torch_dtype=torch_dtype,
+            vae=vae,
+            image_encoder=image_encoder,
+            variant="fp16",
+        )
         self.pipe.load_ip_adapter(
             ip_adapter_model,
             subfolder="",
@@ -214,14 +205,4 @@ class Pipeline:
             ip_adapter_image=[params.image],
             output_type="pil",
         )
-
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/controlnet.py

@@ -159,16 +159,12 @@ class Pipeline:
         controlnet_canny = ControlNetModel.from_pretrained(
             controlnet_model, torch_dtype=torch_dtype
         ).to(device)
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                base_model, controlnet=controlnet_canny
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                base_model,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-            )
+
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            base_model,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -254,14 +250,6 @@ class Pipeline:
             control_guidance_start=params.controlnet_start,
             control_guidance_end=params.controlnet_end,
         )
-
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetDepthFlashSD.py

@@ -146,17 +146,12 @@ class Pipeline:
             device=device,
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_depth, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_depth,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_depth,
+            torch_dtype=torch_dtype,
+        )
 
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
@@ -258,14 +253,6 @@ class Pipeline:
             control_guidance_start=params.controlnet_start,
             control_guidance_end=params.controlnet_end,
         )
-
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_depth:
             # paste control_image on top of result_image

server/pipelines/controlnetDepthHyperSD.py

@@ -148,17 +148,12 @@ class Pipeline:
             device=device,
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_depth, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_depth,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_depth,
+            torch_dtype=torch_dtype,
+        )
 
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
@@ -259,13 +254,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_depth:
             # paste control_image on top of result_image

server/pipelines/controlnetDepthHyperSDXL.py

@@ -151,18 +151,13 @@ class Pipeline:
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_depth, vae=vae, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_depth,
-                vae=vae,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_depth,
+            vae=vae,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.load_lora_weights(
             hf_hub_download("ByteDance/Hyper-SD", "Hyper-SDXL-1step-lora.safetensors")
@@ -258,13 +253,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_depth:
             # paste control_image on top of result_image

server/pipelines/controlnetFlashSD.py

@@ -138,17 +138,12 @@ class Pipeline:
             controlnet_model, torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_canny, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.scheduler = LCMScheduler.from_pretrained(
             model_id,
@@ -252,13 +247,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetFlashSDXL.py

@@ -143,18 +143,13 @@ class Pipeline:
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_canny, vae=vae, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                vae=vae,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            vae=vae,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.scheduler = LCMScheduler.from_pretrained(
             model_id,
@@ -253,13 +248,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetHyperSD.py

@@ -160,17 +160,12 @@ class Pipeline:
             controlnet_model, torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_canny, vae=vae, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.load_lora_weights(
             hf_hub_download("ByteDance/Hyper-SD", "Hyper-SD15-1step-lora.safetensors")
@@ -269,13 +264,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetHyperSDXL.py

@@ -164,18 +164,13 @@ class Pipeline:
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_canny, vae=vae, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                vae=vae,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            vae=vae,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.load_lora_weights(
             hf_hub_download("ByteDance/Hyper-SD", "Hyper-SDXL-1step-lora.safetensors")
@@ -274,13 +269,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetLoraSD15.py

@@ -174,21 +174,13 @@ class Pipeline:
 
         self.pipes = {}
 
-        if args.safety_checker:
-            for base_model_id in base_models.keys():
-                pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                    base_model_id,
-                    controlnet=controlnet_canny,
-                )
-                self.pipes[base_model_id] = pipe
-        else:
-            for base_model_id in base_models.keys():
-                pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                    base_model_id,
-                    safety_checker=None,
-                    controlnet=controlnet_canny,
-                )
-                self.pipes[base_model_id] = pipe
+        for base_model_id in base_models.keys():
+            pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                base_model_id,
+                safety_checker=None,
+                controlnet=controlnet_canny,
+            )
+            self.pipes[base_model_id] = pipe
 
         self.canny_torch = SobelOperator(device=device)
 
@@ -262,13 +254,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetLoraSD15QRCode.py

@@ -154,11 +154,9 @@ class Pipeline:
                 controlnet=controlnet_qrcode,
             )
 
-        self.control_image = Image.open(
-            "qr-code.png").convert("RGB").resize((512, 512))
+        self.control_image = Image.open("qr-code.png").convert("RGB").resize((512, 512))
 
-        self.pipe.scheduler = LCMScheduler.from_config(
-            self.pipe.scheduler.config)
+        self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.set_progress_bar_config(disable=True)
         if device.type != "mps":
             self.pipe.unet.to(memory_format=torch.channels_last)
@@ -206,9 +204,7 @@ class Pipeline:
             steps = math.ceil(1 / max(0.10, strength))
 
         blend_qr_image = Image.blend(
-            params.image,
-            self.control_image,
-            alpha=params.blend
+            params.image, self.control_image, alpha=params.blend
         )
         results = self.pipe(
             image=blend_qr_image,
@@ -227,13 +223,4 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/controlnetLoraSDXL-Lightning.py

@@ -9,7 +9,6 @@ from diffusers import (
 from compel import Compel, ReturnedEmbeddingsType
 import torch
 from pipelines.utils.canny_gpu import SobelOperator
-from pipelines.utils.safety_checker import SafetyChecker
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
@@ -169,10 +168,6 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        self.safety_checker = None
-        if args.safety_checker:
-            self.safety_checker = SafetyChecker(device=device.type)
-
         if args.taesd:
             vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -290,13 +285,7 @@ class Pipeline:
             control_guidance_start=params.controlnet_start,
             control_guidance_end=params.controlnet_end,
         )
-        images = results.images
-        if self.safety_checker:
-            images, has_nsfw_concepts = self.safety_checker(images)
-            if any(has_nsfw_concepts):
-                return None
-
-        result_image = images[0]
+        result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image
             w0, h0 = (200, 200)

server/pipelines/controlnetLoraSDXL.py

@@ -173,19 +173,12 @@ class Pipeline:
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                controlnet=controlnet_canny,
-                vae=vae,
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                vae=vae,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            vae=vae,
+        )
         self.canny_torch = SobelOperator(device=device)
         # Load LCM LoRA
         self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm")
@@ -196,8 +189,7 @@ class Pipeline:
         )
         self.pipe.set_adapters(["lcm", "toy"], adapter_weights=[1.0, 0.8])
 
-        self.pipe.scheduler = LCMScheduler.from_config(
-            self.pipe.scheduler.config)
+        self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype).to(device)
 
@@ -219,8 +211,7 @@ class Pipeline:
         if args.compel:
             self.pipe.compel_proc = Compel(
                 tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2],
-                text_encoder=[self.pipe.text_encoder,
-                              self.pipe.text_encoder_2],
+                text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2],
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
                 requires_pooled=[False, True],
             )
@@ -292,13 +283,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetMistoLineHyperSDXL.py

@@ -166,18 +166,13 @@ class Pipeline:
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id, controlnet=controlnet_canny, vae=vae, torch_dtype=torch_dtype
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                vae=vae,
-                torch_dtype=torch_dtype,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            vae=vae,
+            torch_dtype=torch_dtype,
+        )
 
         self.pipe.load_lora_weights(
             hf_hub_download("ByteDance/Hyper-SD", "Hyper-SDXL-1step-lora.safetensors")
@@ -282,13 +277,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetPCMSD15.py

@@ -140,17 +140,11 @@ class Pipeline:
             controlnet_model, torch_dtype=torch_dtype
         ).to(device)
 
-        if args.safety_checker:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                base_model_id,
-                controlnet=controlnet_canny,
-            )
-        else:
-            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                base_model_id,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-            )
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            base_model_id,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+        )
 
         self.canny_torch = SobelOperator(device=device)
 
@@ -238,13 +232,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/controlnetSDTurbo.py

@@ -7,7 +7,6 @@ from diffusers import (
 from compel import Compel
 import torch
 from pipelines.utils.canny_gpu import SobelOperator
-from pipelines.utils.safety_checker import SafetyChecker
 
 try:
     import intel_extension_for_pytorch as ipex  # type: ignore
@@ -162,10 +161,6 @@ class Pipeline:
         )
         self.pipes = {}
 
-        self.safety_checker = None
-        if args.safety_checker:
-            self.safety_checker = SafetyChecker(device=device.type)
-
         self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
             base_model,
             controlnet=controlnet_canny,
@@ -267,13 +262,7 @@ class Pipeline:
             control_guidance_start=params.controlnet_start,
             control_guidance_end=params.controlnet_end,
         )
-        images = results.images
-        if self.safety_checker:
-            images, has_nsfw_concepts = self.safety_checker(images)
-            if any(has_nsfw_concepts):
-                return None
-
-        result_image = images[0]
+        result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image
             w0, h0 = (200, 200)

server/pipelines/controlnetSDXLTurbo.py

@@ -7,7 +7,6 @@ from diffusers import (
 from compel import Compel, ReturnedEmbeddingsType
 import torch
 from pipelines.utils.canny_gpu import SobelOperator
-from pipelines.utils.safety_checker import SafetyChecker
 
 try:
     import intel_extension_for_pytorch as ipex  # type: ignore
@@ -171,9 +170,6 @@ class Pipeline:
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
-        self.safety_checker = None
-        if args.safety_checker:
-            self.safety_checker = SafetyChecker(device=device.type)
 
         self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
             model_id,
@@ -274,14 +270,7 @@ class Pipeline:
             control_guidance_start=params.controlnet_start,
             control_guidance_end=params.controlnet_end,
         )
-
-        images = results.images
-        if self.safety_checker:
-            images, has_nsfw_concepts = self.safety_checker(images)
-            if any(has_nsfw_concepts):
-                return None
-
-        result_image = images[0]
+        result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image
             w0, h0 = (200, 200)

server/pipelines/controlnetSegmindVegaRT.py

@@ -173,19 +173,12 @@ class Pipeline:
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
-        if args.safety_checker:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                base_model,
-                controlnet=controlnet_canny,
-                vae=vae,
-            )
-        else:
-            self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-                base_model,
-                safety_checker=None,
-                controlnet=controlnet_canny,
-                vae=vae,
-            )
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            base_model,
+            safety_checker=None,
+            controlnet=controlnet_canny,
+            vae=vae,
+        )
         self.canny_torch = SobelOperator(device=device)
 
         self.pipe.load_lora_weights(lora_model)
@@ -285,13 +278,6 @@ class Pipeline:
             control_guidance_end=params.controlnet_end,
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
         result_image = results.images[0]
         if params.debug_canny:
             # paste control_image on top of result_image

server/pipelines/img2img.py

@@ -95,13 +95,10 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(base_model)
-        else:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                safety_checker=None,
-            )
+        self.pipe = AutoPipelineForImage2Image.from_pretrained(
+            base_model,
+            safety_checker=None,
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -171,13 +168,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/img2imgSDTurbo.py

@@ -93,13 +93,10 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(base_model)
-        else:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                safety_checker=None,
-            )
+        self.pipe = AutoPipelineForImage2Image.from_pretrained(
+            base_model,
+            safety_checker=None,
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -187,13 +184,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/img2imgSDXL-Lightning.py

@@ -110,7 +110,6 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-
         if args.taesd:
             vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -214,13 +213,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/img2imgSDXLTurbo.py

@@ -103,13 +103,10 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(base_model)
-        else:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                safety_checker=None,
-            )
+        self.pipe = AutoPipelineForImage2Image.from_pretrained(
+            base_model,
+            safety_checker=None,
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -194,13 +191,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/img2imgSDXS512.py

@@ -92,13 +92,10 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(base_model)
-        else:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                safety_checker=None,
-            )
+        self.pipe = AutoPipelineForImage2Image.from_pretrained(
+            base_model,
+            safety_checker=None,
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -162,14 +159,4 @@ class Pipeline:
             height=params.height,
             output_type="pil",
         )
-
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/img2imgSegmindVegaRT.py

@@ -105,17 +105,11 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                variant="fp16",
-            )
-        else:
-            self.pipe = AutoPipelineForImage2Image.from_pretrained(
-                base_model,
-                safety_checker=None,
-                variant="fp16",
-            )
+        self.pipe = AutoPipelineForImage2Image.from_pretrained(
+            base_model,
+            safety_checker=None,
+            variant="fp16",
+        )
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -205,13 +199,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/txt2img.py

@@ -79,12 +79,7 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = DiffusionPipeline.from_pretrained(base_model)
-        else:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                base_model, safety_checker=None
-            )
+        self.pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None)
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -142,11 +137,5 @@ class Pipeline:
             height=params.height,
             output_type="pil",
         )
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
+
         return results.images[0]

server/pipelines/txt2imgLora.py

@@ -86,12 +86,7 @@ class Pipeline:
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
-        if args.safety_checker:
-            self.pipe = DiffusionPipeline.from_pretrained(base_model)
-        else:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                base_model, safety_checker=None
-            )
+        self.pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None)
         if args.taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
@@ -152,11 +147,5 @@ class Pipeline:
             height=params.height,
             output_type="pil",
         )
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
+
         return results.images[0]

server/pipelines/txt2imgLoraSDXL.py

@@ -95,17 +95,12 @@ class Pipeline:
         vae = AutoencoderKL.from_pretrained(
             "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
         )
-        if args.safety_checker:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                model_id,
-                vae=vae,
-            )
-        else:
-            self.pipe = DiffusionPipeline.from_pretrained(
-                model_id,
-                safety_checker=None,
-                vae=vae,
-            )
+
+        self.pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            safety_checker=None,
+            vae=vae,
+        )
         # Load LCM LoRA
         self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm")
         self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
@@ -184,13 +179,4 @@ class Pipeline:
             output_type="pil",
         )
 
-        nsfw_content_detected = (
-            results.nsfw_content_detected[0]
-            if "nsfw_content_detected" in results
-            else False
-        )
-        if nsfw_content_detected:
-            return None
-        result_image = results.images[0]
-
-        return result_image
+        return results.images[0]

server/pipelines/utils/safety_checker.py

@@ -150,14 +150,20 @@ class SafetyChecker:
         )
 
     def __call__(
-        self, images: list[Image.Image]
-    ) -> tuple[list[Image.Image], list[bool]]:
-        safety_checker_input = self.feature_extractor(images, return_tensors="pt").to(
-            self.device
-        )
+        self, images: list[Image.Image] | Image.Image
+    ) -> tuple[list[Image.Image], list[bool]] | tuple[Image.Image, bool]:
+        images_list = [images] if isinstance(images, Image.Image) else images
+
+        safety_checker_input = self.feature_extractor(
+            images_list, return_tensors="pt"
+        ).to(self.device)
+
         has_nsfw_concepts = self.safety_checker(
-            images=[images],
+            images=[images_list],
             clip_input=safety_checker_input.pixel_values.to(self.device),
         )
 
+        if isinstance(images, Image.Image):
+            return images, has_nsfw_concepts[0]
+
         return images, has_nsfw_concepts