Update app.py

app.py

@@ -20,33 +20,10 @@ checkpoints = {
 }
 loaded = None
 
-
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
-if SAFETY_CHECKER:
-    from safety_checker import StableDiffusionSafetyChecker
-    from transformers import CLIPFeatureExtractor
-
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-        "CompVis/stable-diffusion-safety-checker"
-    ).to("cuda")
-    feature_extractor = CLIPFeatureExtractor.from_pretrained(
-        "openai/clip-vit-base-patch32"
-    )
-
-    def check_nsfw_images(
-        images: list[Image.Image],
-    ) -> tuple[list[Image.Image], list[bool]]:
-        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
-        has_nsfw_concepts = safety_checker(
-            images=[images],
-            clip_input=safety_checker_input.pixel_values.to("cuda")
-        )
-
-        return images, has_nsfw_concepts
-
 # Function 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, ckpt):
@@ -63,16 +40,8 @@ def generate_image(prompt, ckpt):
         
     results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
 
-    if SAFETY_CHECKER:
-        images, has_nsfw_concepts = check_nsfw_images(results.images)
-        if any(has_nsfw_concepts):
-            gr.Warning("NSFW content detected.")
-            return Image.new("RGB", (512, 512))
-        return images[0]
     return results.images[0]
 
-
-
 # Gradio Interface
 description = """
 This demo utilizes the SDXL-Lightning model by ByteDance, which is a lightning-fast text-to-image generative model capable of producing high-quality images in 4 steps.
@@ -98,4 +67,4 @@ with gr.Blocks(css="style.css") as demo:
                  outputs=img,
                  )
     
-demo.queue().launch()
+demo.queue().launch()