Update inpainting_module.py

inpainting_module.py

@@ -1,5 +1,6 @@
 import gc
 import logging
+import os
 import time
 import traceback
 from dataclasses import dataclass, field
@@ -899,21 +900,37 @@ class InpaintingModule:
             self._last_seed = seed
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
+            # Check if running on Hugging Face Spaces
+            is_spaces = os.getenv('SPACE_ID') is not None
+
             # Stage 1: Preview generation
-            if progress_callback:
-                progress_callback("Generating preview...", 30)
+            # On Spaces, skip preview to save time (300s hard limit)
+            preview_result = None
 
-            preview_result = self._generate_inpaint(
-                image=image,
-                mask=processed_mask,
-                control_image=control_image,
-                prompt=enhanced_prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=self.config.preview_steps,
-                guidance_scale=self.config.preview_guidance_scale,
-                controlnet_conditioning_scale=conditioning_scale,
-                generator=generator
-            )
+            if preview_only or not is_spaces:
+                if progress_callback:
+                    progress_callback("Generating preview...", 30)
+
+                # Optimize preview steps for Hugging Face Spaces
+                preview_steps = self.config.preview_steps
+                if is_spaces:
+                    # On Spaces, use minimal preview steps
+                    preview_steps = min(preview_steps, 8)
+                    logger.info(f"Spaces environment - using {preview_steps} preview steps")
+
+                preview_result = self._generate_inpaint(
+                    image=image,
+                    mask=processed_mask,
+                    control_image=control_image,
+                    prompt=enhanced_prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=preview_steps,
+                    guidance_scale=self.config.preview_guidance_scale,
+                    controlnet_conditioning_scale=conditioning_scale,
+                    generator=generator
+                )
+            else:
+                logger.info("Spaces environment - skipping preview to fit 300s limit")
 
             if preview_only:
                 generation_time = time.time() - start_time
@@ -942,6 +959,24 @@ class InpaintingModule:
             num_steps = kwargs.get('num_inference_steps', self.config.num_inference_steps)
             guidance = kwargs.get('guidance_scale', self.config.guidance_scale)
 
+            # Optimize for Hugging Face Spaces ZeroGPU (stateless, 300s hard limit)
+            if is_spaces:
+                # ZeroGPU timing breakdown (stateless architecture):
+                # - Model loading: ~90-120s (SDXL + ControlNet transfer to GPU)
+                # - Inference: ~15s/step (slower on shared H200 vs dedicated L4)
+                # - Platform limit: 300s hard limit (Pro tier)
+                #
+                # Strategy: Skip preview stage, use all time for single-stage generation
+                # Time budget: 300s - 120s (loading) - 15s (overhead) = 165s available
+                # Safe steps: 165s / 15s = 11 steps maximum
+                spaces_max_steps = 11  # Balanced quality vs. time constraint
+
+                if num_steps > spaces_max_steps:
+                    original_steps = num_steps
+                    num_steps = spaces_max_steps
+                    logger.info(f"Spaces ZeroGPU optimization - reduced steps: {original_steps} → {num_steps}")
+                    logger.info(f"Estimated total time: ~{120 + num_steps * 15}s (within 300s limit)")
+
             full_result = self._generate_inpaint(
                 image=image,
                 mask=processed_mask,
@@ -1308,4 +1343,4 @@ class InpaintingModule:
 
         status["memory"] = self._check_memory_status()
 
-        return status
+        return status