Upload app.py with huggingface_hub

app.py

@@ -1105,6 +1105,7 @@ def do_polish_transform_prompt(prompt: str, style: str, do_polish: bool) -> Tupl
 def generate_with_polish(prompt: str, style: str, do_polish: bool, ratio: str, steps: int, seed: int, randomize: bool):
     """Unified generate with progress feedback using generator.
     Yields intermediate status updates with timer so user knows what's happening.
+    Includes automatic retry for ZeroGPU allocation failures.
     """
     logger.info(f"generate_with_polish: do_polish={do_polish}, style={style}, prompt_len={len(prompt) if prompt else 0}")
     
@@ -1133,8 +1134,33 @@ def generate_with_polish(prompt: str, style: str, do_polish: bool, ratio: str, s
     # Show status before GPU generation with the prompt that will be used
     yield None, create_status_html("Generating image", timer.format()), seed
     
-    # GPU generation
-    image, used_seed = generate(full_prompt, polished_display, ratio, steps, seed, randomize)
+    # GPU generation with automatic retry for ZeroGPU failures
+    max_retries = 3
+    image = None
+    used_seed = seed
+    last_error = None
+    
+    for attempt in range(max_retries):
+        try:
+            image, used_seed = generate(full_prompt, polished_display, ratio, steps, seed, randomize)
+            if image is not None:
+                break  # Success
+        except RuntimeError as e:
+            last_error = e
+            error_msg = str(e).lower()
+            if "cuda" in error_msg or "gpu" in error_msg or "driver" in error_msg:
+                logger.warning(f"GPU allocation failed (attempt {attempt + 1}/{max_retries}): {e}")
+                if attempt < max_retries - 1:
+                    import time
+                    time.sleep(1)  # Brief pause before retry
+                    yield None, create_status_html(f"GPU busy, retrying ({attempt + 2}/{max_retries})", timer.format()), seed
+                    continue
+            raise  # Re-raise non-GPU errors
+    
+    if image is None and last_error:
+        timer.stop()
+        yield None, create_status_html(f"GPU unavailable after {max_retries} attempts. Please try again.", timer.format(), is_generating=False).replace("✅", "❌"), seed
+        return
     
     # Stop timer and show final result
     timer.stop()
@@ -1145,6 +1171,7 @@ def generate_with_polish(prompt: str, style: str, do_polish: bool, ratio: str, s
 def transform_with_polish(input_image: Optional[Image.Image], prompt: str, style: str, do_polish: bool, strength: float, steps: int, seed: int, randomize: bool):
     """Unified transform with progress feedback using generator.
     Yields intermediate status updates with timer so user knows what's happening.
+    Includes automatic retry for ZeroGPU allocation failures.
     
     Style-only transformation: When a style is selected without a prompt, the style 
     suffix contains all the information needed to guide the transformation.
@@ -1181,8 +1208,33 @@ def transform_with_polish(input_image: Optional[Image.Image], prompt: str, style
     # Show status before GPU transform
     yield None, create_status_html("Transforming image", timer.format()), 0
     
-    # GPU transform
-    image, used_seed = transform(input_image, full_prompt, polished_display, strength, steps, seed, randomize)
+    # GPU transform with automatic retry for ZeroGPU failures
+    max_retries = 3
+    image = None
+    used_seed = 0
+    last_error = None
+    
+    for attempt in range(max_retries):
+        try:
+            image, used_seed = transform(input_image, full_prompt, polished_display, strength, steps, seed, randomize)
+            if image is not None:
+                break  # Success
+        except RuntimeError as e:
+            last_error = e
+            error_msg = str(e).lower()
+            if "cuda" in error_msg or "gpu" in error_msg or "driver" in error_msg:
+                logger.warning(f"GPU allocation failed (attempt {attempt + 1}/{max_retries}): {e}")
+                if attempt < max_retries - 1:
+                    import time
+                    time.sleep(1)  # Brief pause before retry
+                    yield None, create_status_html(f"GPU busy, retrying ({attempt + 2}/{max_retries})", timer.format()), 0
+                    continue
+            raise  # Re-raise non-GPU errors
+    
+    if image is None and last_error:
+        timer.stop()
+        yield None, create_status_html(f"GPU unavailable after {max_retries} attempts. Please try again.", timer.format(), is_generating=False).replace("✅", "❌"), 0
+        return
     
     # Stop timer and show final result
     timer.stop()