Update app.py

app.py

@@ -35,12 +35,15 @@ if SAFETY_CHECKER:
         "openai/clip-vit-base-patch32"
     )
 
-    def check_nsfw_images(images: list[Image.Image]) -> tuple[list[Image.Image], list[bool]]:
+    def check_nsfw_images(
+        images: list[Image.Image],
+    ) -> tuple[list[Image.Image], list[bool]]:
         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.bool().tolist()
 
 @spaces.GPU(enable_queue=True)
@@ -52,18 +55,17 @@ def generate_image(prompt, ckpt):
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps == 1 else "epsilon")
         pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
         loaded = num_inference_steps
-        
+
     results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=7.5)
 
     if SAFETY_CHECKER:
         images, has_nsfw_concepts = check_nsfw_images(results.images)
         if any(has_nsfw_concepts):
-            return [Image.new("RGB", (512, 512), "black")]
+            return Image.new("RGB", (512, 512), "black")
     return results.images[0]
 
 description = """
-🌌 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.
+🌌 Engage in the exploration of galaxies with the advanced SDXL-Lightning model, a creation of ByteDance capable of transforming your textual descriptions into vivid images at warp speed. This is a joint venture initiated by Starfleet, enabling creative minds to visualize the uncharted territories of space. 🚀 Link to model: [ByteDance/SDXL-Lightning](https://huggingface.co/ByteDance/SDXL-Lightning)
 """
 
 with gr.Blocks(css="style.css") as demo:
@@ -71,10 +73,10 @@ 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)', 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 = gr.Textbox(label='Input your cosmic prompt (English)', placeholder="Describe the celestial phenomena...", scale=8)
+            ckpt = gr.Dropdown(label='Choose your Warp Factor', choices=['Warp 1', 'Warp 2', 'Warp 4', 'Warp 8'], value='Warp 4', interactive=True)
+            submit = gr.Button("Initiate Image Generation", scale=1, variant='primary')
+    img = gr.Image(label='Visual Manifestation of the Cosmos')
 
     prompt.submit(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
     submit.click(fn=generate_image, inputs=[prompt, ckpt], outputs=img)