Update app.py

app.py

@@ -13,41 +13,39 @@ SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 checkpoints = {
-    "Warp 1" : ["sdxl_lightning_1step_unet_x0.safetensors", 1],
-    "Warp 2" : ["sdxl_lightning_2step_unet.safetensors", 2],
-    "Warp 4" : ["sdxl_lightning_4step_unet.safetensors", 4],
-    "Warp 8" : ["sdxl_lightning_8step_unet.safetensors", 8],
+    "Warp 1": ["sdxl_lightning_1step_unet_x0.safetensors", 1],
+    "Warp 2": ["sdxl_lightning_2step_unet.safetensors", 2],
+    "Warp 4": ["sdxl_lightning_4step_unet.safetensors", 4],
+    "Warp 8": ["sdxl_lightning_8step_unet.safetensors", 8],
 }
 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
+    from transformers import CLIPFeatureExtractor, CLIPProcessor
 
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
     ).to("cuda")
-    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+    # Updated to use CLIPProcessor as CLIPFeatureExtractor is deprecated
+    feature_extractor = CLIPProcessor.from_pretrained(
         "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, return_tensors="pt").to("cuda")
         has_nsfw_concepts = safety_checker(
-            images=[images],
-            clip_input=safety_checker_input.pixel_values.to("cuda")
+            images=images,
+            clip_input=safety_checker_input.pixel_values.to("cuda"),
         )
 
         return images, has_nsfw_concepts
 
-# Function 
+# Function
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, ckpt):
     global loaded
@@ -57,10 +55,10 @@ def generate_image(prompt, ckpt):
     num_inference_steps = checkpoints[ckpt][1]
 
     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")
+        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)
 
     if SAFETY_CHECKER:
@@ -71,8 +69,11 @@ def generate_image(prompt, ckpt):
         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!
+"""
 
-# Update the dropdown choices to match the corrected keys in the 'checkpoints' dictionary
 with gr.Blocks(css="style.css") as demo:
     gr.HTML("<h1><center>🌌 Starfleet Command: Text-to-Image Warp Drive - SDXL-Lightning ⚡</center></h1>")
     gr.Markdown(description)
@@ -83,13 +84,7 @@ with gr.Blocks(css="style.css") as demo:
             submit = gr.Button(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,
-                 )
-    
+    prompt.submit(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
+    submit.click(fn=generate_image, inputs=[prompt, ckpt], outputs=img)
+
 demo.queue().launch()