Refactor HF Spaces detection and GPU handling; streamline process_images function for better clarity and maintainability

app.py

@@ -32,49 +32,28 @@ NUM_INFERENCE_STEPS = 30
 
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"}
 
-# NOTE:
-# - This is the quantized 4-bit (bitsandbytes) model, which REQUIRES GPU at load time.
-# - On HF Spaces ZeroGPU, you MUST only load it inside a @spaces.GPU function.
+# Quantized 4-bit model (requires GPU)
 MODEL_ID = "diffusers/FLUX.2-dev-bnb-4bit"
 TORCH_DTYPE = torch.bfloat16
 
 
 # =========================
-# HF Spaces env detection + safe "spaces" decorator
+# HF Spaces detection (robust)
 # =========================
-def _is_hf_space_env() -> bool:
-    """
-    Detect whether we're running on Hugging Face Spaces runtime.
-    Common env vars present in Spaces:
-      - SPACE_ID
-      - HF_SPACE_ID
-      - SYSTEM=spaces
-    """
-    return any(os.getenv(k) for k in ("SPACE_ID", "HF_SPACE_ID")) or (
-        os.getenv("SYSTEM", "").lower() == "spaces"
-    )
-
-
-IS_HF_SPACES = _is_hf_space_env()
-
 try:
-    import spaces  # available on HF Spaces
+    import spaces  # Only available on Hugging Face Spaces
+
+    SPACES_AVAILABLE = True
 except Exception:
     spaces = None
+    SPACES_AVAILABLE = False
 
-
-def gpu_decorator(duration: int = 180):
-    """
-    If on HF Spaces and spaces.GPU exists -> use it.
-    Else -> no-op decorator (local runs normally, using local GPU if available).
-    """
-    if IS_HF_SPACES and (spaces is not None) and hasattr(spaces, "GPU"):
-        return spaces.GPU(duration=duration)
-
-    def _noop(fn):
-        return fn
-
-    return _noop
+# Extra friendly label for UI (not used for logic)
+IS_HF_SPACES = (
+    SPACES_AVAILABLE
+    or (os.getenv("SYSTEM", "").lower() == "spaces")
+    or bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID"))
+)
 
 
 # =========================
@@ -134,15 +113,17 @@ def find_images(directory: str):
 def _get_pipe():
     """
     Lazy-load the pipeline.
-    - On HF Spaces ZeroGPU: MUST be called inside a @spaces.GPU runtime (gpu_decorator handles this).
-    - Locally: will use local GPU.
+    - On HF ZeroGPU: this MUST be called inside a @spaces.GPU function runtime.
+    - Locally: uses local CUDA GPU.
     """
     global _pipe
 
     if _pipe is None:
         if not torch.cuda.is_available():
             raise RuntimeError(
-                "No CUDA GPU detected. This 4-bit bnb model requires GPU to load/run."
+                "CUDA GPU is not available in the current runtime. "
+                "This bnb-4bit model requires GPU. "
+                "On HF ZeroGPU, ensure the function is decorated with @spaces.GPU."
             )
 
         _pipe = Flux2Pipeline.from_pretrained(
@@ -154,14 +135,15 @@ def _get_pipe():
     return _pipe
 
 
-@gpu_decorator(duration=180)
-def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progress()):
+def _process_images_impl(
+    files, prompt, guidance_scale, num_steps, progress=gr.Progress()
+):
     """
-    Process uploaded files (images or archives) and return:
-    - gallery preview (first 10 pairs)
-    - files download (when only images uploaded)
-    - zip download (when archive uploaded)
-    - summary text
+    Shared implementation used by both:
+      - HF Spaces GPU wrapper
+      - Local runtime
+    Returns 4 outputs:
+      gallery, files_download, zip_download, summary
     """
     if not files:
         return (
@@ -177,12 +159,11 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
     guidance_scale = float(guidance_scale)
     num_steps = int(num_steps)
 
-    # Temp for extraction / staging input
+    # Temp for extraction/staging input
     temp_dir = tempfile.mkdtemp(prefix="flux_in_")
 
     # IMPORTANT:
-    # Result files MUST remain on disk for Gradio download.
-    # So we create a persistent temp dir and DO NOT delete it in finally.
+    # Output files MUST remain on disk for Gradio downloads.
     run_id = uuid.uuid4().hex[:10]
     output_dir = tempfile.mkdtemp(prefix=f"flux_results_{run_id}_")
 
@@ -222,8 +203,7 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
         total_images = len(all_images)
         progress(0.10, desc=f"Found {total_images} images. Loading model...")
 
-        # Load pipeline (inside GPU runtime on Spaces; local GPU otherwise)
-        pipe = _get_pipe()
+        pipe = _get_pipe()  # IMPORTANT: must be inside GPU runtime on Spaces
 
         results = []
         metrics_lines = []
@@ -249,7 +229,6 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
 
             # Preserve structure if from archive
             output_rel_path = rel_path
-
             out_path = os.path.join(output_dir, output_rel_path)
             os.makedirs(os.path.dirname(out_path), exist_ok=True)
 
@@ -287,7 +266,7 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
             "Individual metrics:\n" + "\n".join(metrics_lines)
         )
 
-        # Build gallery preview (first 10 results -> original+enhanced)
+        # Gallery preview
         gallery_images = []
         for r in results[:10]:
             gallery_images.append((r["original"], f"Original: {r['filename']}"))
@@ -298,9 +277,9 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
                 )
             )
 
-        # Decide download output:
-        # - If user uploaded any archive -> provide ZIP
-        # - Else -> provide enhanced files directly
+        # Download behavior:
+        # - If any archive uploaded -> zip
+        # - Else -> direct files list
         if has_archive:
             progress(0.92, desc="Packaging ZIP...")
 
@@ -341,10 +320,30 @@ def process_images(files, prompt, guidance_scale, num_steps, progress=gr.Progres
         )
 
     finally:
-        # Cleanup ONLY input/extraction temp dir.
-        # DO NOT delete output_dir because Gradio downloads need the files to remain.
+        # Clean ONLY input extraction temp.
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir, ignore_errors=True)
+        # Do NOT delete output_dir; downloads need it to exist.
+
+
+# =========================
+# IMPORTANT: Define a real @spaces.GPU function at import-time on Spaces
+# (This fixes: "No @spaces.GPU function detected during startup")
+# =========================
+if SPACES_AVAILABLE:
+
+    @spaces.GPU(duration=180)
+    def process_images(
+        files, prompt, guidance_scale, num_steps, progress=gr.Progress()
+    ):
+        return _process_images_impl(files, prompt, guidance_scale, num_steps, progress)
+
+else:
+
+    def process_images(
+        files, prompt, guidance_scale, num_steps, progress=gr.Progress()
+    ):
+        return _process_images_impl(files, prompt, guidance_scale, num_steps, progress)
 
 
 # =========================
@@ -362,11 +361,12 @@ Upload microscopy images (individual files or compressed archives) for AI-powere
 - Archives: ZIP, 7Z (will process all images inside)
 
 **Download behavior:**
-- If you upload **only images** → you can download the enhanced **image files directly** (`*_flux` suffix)
-- If you upload **a ZIP/7Z** → you can download **one ZIP** (images inside use `*_flux` suffix)
+- Upload **only images** → download enhanced **image files directly** (`*_flux` suffix)
+- Upload **ZIP/7Z** → download **one ZIP** (images inside use `*_flux` suffix)
 
 **Runtime detection:**
 - Detected environment: **{"Hugging Face Spaces" if IS_HF_SPACES else "Local"}**
+- spaces module available: **{SPACES_AVAILABLE}**
 """
     )
 
@@ -443,11 +443,11 @@ Upload microscopy images (individual files or compressed archives) for AI-powere
 - **Inference Steps**: 30 (balanced quality and speed)
 
 ### Quality Metrics
-- **PSNR** (Peak Signal-to-Noise Ratio): Higher is better
-- **SSIM** (Structural Similarity Index): Closer to 1.0 is better
+- **PSNR**: Higher is better
+- **SSIM**: Closer to 1.0 is better
 """
     )
 
 if __name__ == "__main__":
-    demo.queue()  # recommended for Spaces
+    demo.queue()
     demo.launch(share=False)