Update app.py

app.py

@@ -13,13 +13,14 @@ 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],
+    "1-Step" : ["sdxl_lightning_1step_unet_x0.safetensors", 1],
+    "2-Step" : ["sdxl_lightning_2step_unet.safetensors", 2],
+    "4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
+    "8-Step" : ["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")
@@ -38,32 +39,39 @@ if SAFETY_CHECKER:
     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")
+        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.bool().tolist()
+        return images, has_nsfw_concepts
 
+# Function 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, ckpt):
     global loaded
-    checkpoint, num_inference_steps = checkpoints[ckpt]
+    print(prompt, ckpt)
+
+    checkpoint = checkpoints[ckpt][0]
+    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=7.5)
+        
+    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):
-            return Image.new("RGB", (512, 512), "black")
+            gr.Warning("NSFW content detected.")
+            return Image.new("RGB", (512, 512))
+        return images[0]
     return results.images[0]
 
+# Gradio Interface
 description = """
 🌌 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)
 """