Update app.py

app.py

@@ -3,6 +3,7 @@ import os
 import requests
 import torch
 from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from diffusers.models import AutoencoderKL
 from PIL import Image
 from RealESRGAN import RealESRGAN
@@ -117,6 +118,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 controlnet = ControlNetModel.from_single_file(
     "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
 )
+safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
 
 # Load the Stable Diffusion pipeline with Juggernaut Reborn model
 model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
@@ -124,7 +126,8 @@ pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
     model_path,
     controlnet=controlnet,
     torch_dtype=torch.float16,
-    use_safetensors=True
+    use_safetensors=True,
+    safety_checker=safety_checker
 )
 
 # Load and set VAE
@@ -134,7 +137,6 @@ vae = AutoencoderKL.from_single_file(
 )
 pipe.vae = vae
 
-
 # Load embeddings and LoRA models
 pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
 pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
@@ -248,7 +250,6 @@ def process_image(input_image, prompt, negative_prompt, resolution=2048, num_inf
 
     return result
 
-
 # Simple options
 simple_options = [
     gr.Image(type="pil", label="Input Image"),
@@ -268,5 +269,4 @@ iface = gr.Interface(
     description="Upload an image and adjust the settings to process it using Stable Diffusion."
 )
 
-
 iface.launch()