Update app.py

app.py

@@ -3,29 +3,38 @@ import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
 
+# ====== DEVICE SETUP ======
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=False,
+    torch_dtype=torch.float32,          # CPU-friendly dtype
+    low_cpu_mem_usage=True,
 )
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe.to(device)
-
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
 
+pipe.to(device)
 print("Pipeline loaded!")
 
-@spaces.GPU
-def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+
+def generate_image(
+    prompt,
+    height,
+    width,
+    num_inference_steps,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
     """Generate an image from the given prompt."""
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    
-    generator = torch.Generator("cuda").manual_seed(int(seed))
+
+    gen_device = "cuda" if (device == "cuda") else "cpu"
+    generator = torch.Generator(gen_device).manual_seed(int(seed))
+
     image = pipe(
         prompt=prompt,
         height=int(height),
@@ -34,9 +43,10 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
     ).images[0]
-    
+
     return image, seed
 
+
 # Example prompts
 examples = [
     ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
@@ -258,4 +268,4 @@ if __name__ == "__main__":
             "gradio"
         ],
         mcp_server=True
-    )
+    )