Update app.py

app.py

@@ -9,7 +9,7 @@ from PIL import Image
 
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
 
-# Constants for Starfleet Command theme
+# Constants
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 checkpoints = {
@@ -36,21 +36,17 @@ if SAFETY_CHECKER:
     )
 
     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")
+        safety_checker_input = feature_extractor(images=[image.convert("RGB") for image in 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
+        return images, has_nsfw_concepts.bool().tolist()
 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, ckpt):
     global loaded
-    print("🌌 Starfleet Command: ", prompt, ckpt)
-
-    checkpoint = checkpoints[ckpt][0]
-    num_inference_steps = checkpoints[ckpt][1]
+    checkpoint, num_inference_steps = checkpoints[ckpt]
 
     if loaded != num_inference_steps:
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps == 1 else "epsilon")
@@ -62,15 +58,12 @@ def generate_image(prompt, ckpt):
     if SAFETY_CHECKER:
         images, has_nsfw_concepts = check_nsfw_images(results.images)
         if any(has_nsfw_concepts):
-            gr.Alert("🚨 NSFW content detected. Displaying a safe placeholder instead.").show()
-            return Image.new("RGB", (512, 512), "black")
-        return images[0]
+            return [Image.new("RGB", (512, 512), "black")]
     return results.images[0]
 
 description = """
-🌌 Welcome to Starfleet Command's Text-to-Image Warp Drive - SDXL-Lightning ⚡.
-Utilizing ByteDance's SDXL-Lightning model, this interface allows you to generate high-quality cosmic images from textual prompts. 
-Embark on your voyage of imagination and see the universe through the lens of advanced AI.
+🌌 Welcome to the Starfleet Command's Advanced Image Generation Console. Utilizing the cutting-edge SDXL-Lightning model, this terminal allows Starfleet officers and Federation citizens alike to materialize visual representations from textual descriptions at various warp speeds.
+🖖 Boldly go where no one has gone before - create images as vast as the universe with your imagination.
 """
 
 with gr.Blocks(css="style.css") as demo:
@@ -78,13 +71,12 @@ with gr.Blocks(css="style.css") as demo:
     gr.Markdown(description)
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label='Your cosmic prompt (English)', scale=8)
-            ckpt = gr.Dropdown(label='Warp Factor (Inference Speed)', choices=['Warp 1', 'Warp 2', 'Warp 4', 'Warp 8'], value='Warp 4', interactive=True)
-            submit = gr.Button(text="Engage", scale=1, variant='primary')
-    img = gr.Image(label='The Universe, As You Envision It')
+            prompt = gr.Textbox(label='Your cosmic prompt (English)', placeholder="Describe the universe you wish to visualize...", scale=8)
+            ckpt = gr.Dropdown(label='Warp Factor (Inference Speed)', choices=list(checkpoints.keys()), value='Warp 4', interactive=True)
+            submit = gr.Button("Engage", scale=1, variant='primary')
+    img = gr.Image(label='Your Personalized Universe')
 
     prompt.submit(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
     submit.click(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
 
 demo.queue().launch()
-