Update app.py

app.py

@@ -9,7 +9,7 @@ from PIL import Image
 
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
 
-# Constants
+# Constants for Starfleet Command theme
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 checkpoints = {
@@ -26,13 +26,12 @@ if torch.cuda.is_available():
 
 if SAFETY_CHECKER:
     from safety_checker import StableDiffusionSafetyChecker
-    from transformers import CLIPFeatureExtractor, CLIPProcessor
+    from transformers import CLIPFeatureExtractor
 
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
     ).to("cuda")
-    # Updated to use CLIPProcessor as CLIPFeatureExtractor is deprecated
-    feature_extractor = CLIPProcessor.from_pretrained(
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
         "openai/clip-vit-base-patch32"
     )
 
@@ -45,11 +44,10 @@ if SAFETY_CHECKER:
 
         return images, has_nsfw_concepts
 
-# Function
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, ckpt):
     global loaded
-    print(prompt, ckpt)
+    print("🌌 Starfleet Command: ", prompt, ckpt)
 
     checkpoint = checkpoints[ckpt][0]
     num_inference_steps = checkpoints[ckpt][1]
@@ -58,20 +56,21 @@ 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=0)
+        
+    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):
-            gr.Warning("NSFW content detected.")
-            return Image.new("RGB", (512, 512))
+            gr.Alert("🚨 NSFW content detected. Displaying a safe placeholder instead.").show()
+            return Image.new("RGB", (512, 512), "black")
         return images[0]
     return results.images[0]
 
-# Gradio Interface
 description = """
-Welcome aboard the Starship SDXL Enterprise! Our mission: to explore strange new AI generations, to seek out new visual frontiers and computational boundaries, to boldly generate images like no one has seen before. Utilizing the cutting-edge SDXL-Lightning model, we're at the forefront of text-to-image technology, ready to transform your imaginative prompts into visual spectacles. Whether you're navigating the uncharted territories of outer space or delving into the realms of fantasy, your adventure begins now. Model powered by the pioneering intellects at ByteDance. Journey safely through the stars!
+🌌 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.
 """
 
 with gr.Blocks(css="style.css") as demo:
@@ -81,10 +80,11 @@ with gr.Blocks(css="style.css") as demo:
         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(scale=1, variant='primary')
+            submit = gr.Button(text="Engage", scale=1, variant='primary')
     img = gr.Image(label='The Universe, As You Envision It')
 
     prompt.submit(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
     submit.click(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
 
 demo.queue().launch()
+