Update app.py

app.py

@@ -6,36 +6,44 @@ from diffusers import FluxPipeline
 from huggingface_hub import login
 
 # --- AUTHENTICATION ---
-# Ensure your secret is named HF_TOKEN in Space Settings
 token = os.getenv("HF_TOKEN")
 if token:
     login(token=token)
 
-# --- MODEL LOADING WITH GATED CHECK ---
+# --- MODEL LOADING ---
 model_id = "black-forest-labs/FLUX.1-schnell"
 
+# Load without .to("cuda") here. Let ZeroGPU handle it.
 pipe = FluxPipeline.from_pretrained(
     model_id, 
     torch_dtype=torch.bfloat16,
-    token=token,           # Explicitly passing token for gated access
-    use_auth_token=token   # Backup parameter for older diffusers versions
+    token=token
 )
 
 @spaces.GPU(duration=60)
 def generate_image(prompt):
-    # Hidden settings for a clean UI
-    generator = torch.Generator("cuda").manual_seed(42)
-        
-    image = pipe(
-        prompt=prompt,
-        width=1024,
-        height=1024,
-        num_inference_steps=4, 
-        generator=generator,
-        guidance_scale=0.0
-    ).images[0]
+    if not prompt:
+        return None
+    
+    try:
+        # ZeroGPU specific: ensure the pipe is on the right device
+        pipe.to("cuda") 
         
-    return image
+        generator = torch.Generator("cuda").manual_seed(42)
+            
+        image = pipe(
+            prompt=prompt,
+            width=1024,
+            height=1024,
+            num_inference_steps=4, 
+            generator=generator,
+            guidance_scale=0.0
+        ).images[0]
+            
+        return image
+    except Exception as e:
+        # This will show the actual error message in the Gradio UI
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 # --- MINIMALIST UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -54,4 +62,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     btn.click(fn=generate_image, inputs=[prompt], outputs=output_img)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)