Update app.py

app.py

@@ -8,25 +8,23 @@ import gradio as gr
 base_model = "stabilityai/stable-diffusion-xl-base-1.0"
 lightning_repo = "ByteDance/SDXL-Lightning"
 checkpoint_filename = "sdxl_lightning_2step_unet.safetensors"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"  # Force CPU
 
 # --- Download and Load Lightning UNet ---
 print("Downloading Lightning checkpoint...")
 ckpt_path = cached_file(lightning_repo, checkpoint_filename, cache_dir=".cache")
 
-print("Loading UNet (FP16)...")
+print("Loading UNet (CPU)...")
 unet = UNet2DConditionModel.from_config(base_model, subfolder="unet")
 unet.load_state_dict(load_file(ckpt_path))
-unet.to(device, dtype=torch.float16)
+unet.to(device)
 unet.eval()
 
-# --- Load Pipeline with FP16 ---
+# --- Load Pipeline without FP16 ---
 print("Loading Stable Diffusion XL Pipeline...")
 pipe = StableDiffusionXLPipeline.from_pretrained(
     base_model,
-    unet=unet,
-    torch_dtype=torch.float16,
-    variant="fp16"
+    unet=unet
 )
 pipe.to(device)
 
@@ -34,8 +32,7 @@ pipe.to(device)
 def generate_image(prompt):
     if not prompt:
         return "Prompt cannot be empty!"
-    with torch.autocast(device_type="cuda"):
-        image = pipe(prompt, num_inference_steps=2, guidance_scale=0).images[0]
+    image = pipe(prompt, num_inference_steps=2, guidance_scale=0).images[0]
     return image
 
 # --- Example Prompts ---
@@ -52,11 +49,11 @@ demo = gr.Interface(
     fn=generate_image,
     inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g., A castle floating in the clouds"),
     outputs=gr.Image(type="pil"),
-    title="SDXL-Lightning (2-Step, FP16) Image Generator",
-    description="Fast & optimized image generation using ByteDance's SDXL-Lightning 4-step model with FP16 on CUDA.",
+    title="SDXL-Lightning (2-Step, CPU) Image Generator",
+    description="Fast image generation using ByteDance's SDXL-Lightning 2-step model on CPU (no optimization).",
     examples=examples,
-    cache_examples=False  # Prevent FileNotFoundError
+    cache_examples=False
 )
 
-# --- Launch Public Server ---
+# --- Launch Interface ---
 demo.launch()