Update app.py

app.py

@@ -3,27 +3,24 @@ import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
 
-# Load the pipeline once at startup
+# Load the pipeline once at startup — on CPU only
 print("Booting BRAINDANCE FACTORY neural core...")
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
-pipe.to("cuda")
-
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
 
 print("Neural core online. Synaptic links established.")
 
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
     """Jack in. Generate."""
+    pipe.to("cuda")  # ✅ GPU is only available inside @spaces.GPU functions
+
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    
+
     generator = torch.Generator("cuda").manual_seed(int(seed))
     image = pipe(
         prompt=prompt,
@@ -33,7 +30,9 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
     ).images[0]
-    
+
+    pipe.to("cpu")  # ✅ Optional but good practice: release GPU between ZeroGPU calls
+
     return image, seed
 
 ci_look = """