updating with zero gpu

app.py

@@ -6,6 +6,7 @@ import io
 import gradio as gr
 from PIL import Image
 import os
+import spaces  # Import spaces for Zero GPU
 
 # Configuration
 repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
@@ -44,7 +45,7 @@ try:
         repo_id,
         text_encoder=None,
         torch_dtype=torch_dtype,
-        device_map="balanced"
+        device_map="cuda"
     )
     if not torch.cuda.is_available():
         pipe = pipe.to(device)
@@ -53,6 +54,13 @@ except Exception as e:
     print(f"Error loading pipeline: {e}")
     raise
 
+def get_duration(num_inference_steps: int, input_image: Image.Image = None):
+    """Calculate dynamic GPU duration based on inference steps and input image."""
+    num_images = 0 if input_image is None else 1
+    step_duration = 1 + 0.7 * num_images
+    return max(65, num_inference_steps * step_duration + 10)
+
+@spaces.GPU(duration=get_duration)  # Dynamic GPU allocation
 def generate_image(
     prompt: str,
     input_image: Image.Image = None,
@@ -74,9 +82,6 @@ def generate_image(
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter a prompt!")
     
-    if not torch.cuda.is_available():
-        raise gr.Error("This Space requires a GPU to run. Please try a different Space or upgrade your hardware.")
-    
     progress(0, desc="Encoding prompt...")
     
     try:
@@ -132,7 +137,7 @@ with gr.Blocks(
         
         Supports both **text-to-image** and **image-to-image** generation.
         
-        ⚠️ **Note**: This Space requires a GPU to run.
+        ⚡ **Powered by Hugging Face Zero GPU** - Automatic GPU allocation on demand!
         """
     )
     
@@ -197,6 +202,7 @@ with gr.Blocks(
                 - Start with 28 steps for a good balance of quality and speed
                 - Higher guidance scale follows your prompt more strictly
                 - Use the same seed to reproduce results
+                - First generation may take longer as the model loads
                 """
             )
         
@@ -261,4 +267,4 @@ with gr.Blocks(
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()