Update app.py from anycoder

app.py

@@ -4,19 +4,45 @@ A Gradio 6 application for image-to-image editing using the GLM-Image model.
 
 This app allows users to upload an image and provide a prompt to transform
 the image using the GLM-Image diffusion model.
+Features ZeroGPU support for dynamic GPU allocation on Hugging Face Spaces.
 """
 
 import gradio as gr
 import torch
 from diffusers.pipelines.glm_image import GlmImagePipeline
 from PIL import Image
+from gradio import spaces
+import time
 
-# Initialize the model at startup
-pipe = GlmImagePipeline.from_pretrained(
-    "zai-org/GLM-Image",
-    torch_dtype=torch.bfloat16,
-    device_map="cuda"
-)
+# Global pipeline variable
+pipe = None
+
+def load_model():
+    """Load the GLM-Image model with bfloat16 precision."""
+    global pipe
+    if pipe is None:
+        pipe = GlmImagePipeline.from_pretrained(
+            "zai-org/GLM-Image",
+            torch_dtype=torch.bfloat16,
+            device_map="auto"
+        )
+    return pipe
+
+def estimate_duration(num_inference_steps: int) -> str:
+    """
+    Estimate the processing duration based on inference steps.
+    Returns a human-readable time estimate.
+    """
+    base_time = 30  # Base processing time in seconds
+    step_factor = 0.8  # Seconds per inference step
+    estimated_seconds = base_time + (num_inference_steps * step_factor)
+    
+    if estimated_seconds < 60:
+        return f"~{int(estimated_seconds)}s"
+    else:
+        minutes = estimated_seconds // 60
+        seconds = estimated_seconds % 60
+        return f"~{int(minutes)}m {int(seconds)}s"
 
 def validate_dimensions(height: int, width: int) -> tuple:
     """
@@ -31,6 +57,7 @@ def get_image_dimensions(image: Image.Image) -> tuple:
     """Get the dimensions of an uploaded PIL image."""
     return image.size[1], image.size[0]  # height, width
 
+@spaces.GPU(environment="HF_SPACE", memory=8, timeout=1200, queue=True)
 def process_image(
     image: Image.Image,
     prompt: str,
@@ -43,6 +70,7 @@ def process_image(
 ) -> tuple:
     """
     Process the image through the GLM-Image pipeline.
+    Uses ZeroGPU for dynamic GPU allocation.
     
     Args:
         image: Input PIL Image
@@ -69,13 +97,16 @@ def process_image(
         if adjusted_height != height or adjusted_width != width:
             height, width = adjusted_height, adjusted_width
         
-        progress(0.1, desc="Processing image...")
+        progress(0.1, desc="Loading model...")
+        pipeline = load_model()
+        
+        progress(0.2, desc="Preparing image...")
         input_image = image.convert("RGB")
         
         generator = torch.Generator(device="cuda").manual_seed(seed)
         
-        progress(0.3, desc="Generating image...")
-        result = pipe(
+        progress(0.4, desc="Generating image...", visible=True)
+        result = pipeline(
             prompt=prompt,
             image=[input_image],
             height=height,
@@ -110,6 +141,10 @@ def generate_random_seed() -> int:
     import random
     return random.randint(0, 2**32 - 1)
 
+def update_time_estimate(num_steps: int) -> str:
+    """Update the estimated processing time display."""
+    return f"Estimated time: {estimate_duration(num_steps)}"
+
 custom_theme = gr.themes.Soft(
     primary_hue="indigo",
     secondary_hue="blue",
@@ -215,6 +250,11 @@ with gr.Blocks(fill_height=True) as demo:
                         size="sm",
                         variant="secondary"
                     )
+                
+                time_estimate = gr.Markdown(
+                    value=f"**Estimated time:** {estimate_duration(50)}",
+                    elem_classes=["time-estimate"]
+                )
             
             with gr.Row():
                 generate_btn = gr.Button(
@@ -243,7 +283,7 @@ with gr.Blocks(fill_height=True) as demo:
             
             status = gr.Textbox(
                 label="Status",
-                value="Ready to generate!",
+                value="Ready to generate! GPU will be allocated automatically.",
                 interactive=False,
                 show_label=True
             )
@@ -299,6 +339,13 @@ with gr.Blocks(fill_height=True) as demo:
         api_visibility="private"
     )
     
+    num_inference_steps.change(
+        fn=update_time_estimate,
+        inputs=num_inference_steps,
+        outputs=time_estimate,
+        api_visibility="private"
+    )
+    
     generate_btn.click(
         fn=process_image,
         inputs=[
@@ -331,7 +378,7 @@ with gr.Blocks(fill_height=True) as demo:
             input_image: None,
             prompt: "",
             output_image: None,
-            status: "Ready to generate!",
+            status: "Ready to generate! GPU will be allocated automatically.",
             download_btn: gr.DownloadButton(interactive=False)
         }
     
@@ -369,6 +416,14 @@ demo.launch(
         color: #ffd700 !important;
         text-decoration: underline;
     }
+    .time-estimate {
+        font-size: 0.9em;
+        color: var(--neutral-600);
+        padding: 0.5rem;
+        background: var(--neutral-100);
+        border-radius: var(--radius-sm);
+        margin-top: 0.5rem;
+    }
     #input-image, #output-image {
         border: 2px dashed var(--neutral-300);
         border-radius: var(--radius-lg);
@@ -380,6 +435,7 @@ demo.launch(
     footer_links=[
         {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
         {"label": "GLM-Image Model", "url": "https://huggingface.co/zai-org/GLM-Image"},
+        {"label": "ZeroGPU", "url": "https://huggingface.co/docs/spaces/spaces-sdks/gradio-zerogpu"},
         {"label": "Diffusers Library", "url": "https://github.com/huggingface/diffusers"}
     ],
     server_name="0.0.0.0",