Implement robust @spaces.GPU handling for ZeroGPU support; streamline HF Spaces detection and improve process_images function clarity

app.py

@@ -18,6 +18,35 @@ from skimage.metrics import peak_signal_noise_ratio, structural_similarity
 from skimage.util import img_as_float
 
 
+# =========================
+# Make @spaces.GPU ALWAYS exist (critical for ZeroGPU startup scan)
+# =========================
+try:
+    import spaces  # Hugging Face Spaces provides this
+except Exception:
+
+    class _DummySpaces:
+        @staticmethod
+        def GPU(duration: int = 180, **kwargs):
+            # No-op decorator for local runtime
+            def _decorator(fn):
+                return fn
+
+            return _decorator
+
+    spaces = _DummySpaces()
+
+
+def _is_hf_spaces_env() -> bool:
+    # Best-effort env detection (informational only)
+    return (os.getenv("SYSTEM", "").lower() == "spaces") or bool(
+        os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID")
+    )
+
+
+IS_HF_SPACES = _is_hf_spaces_env()
+
+
 # =========================
 # Config
 # =========================
@@ -36,29 +65,7 @@ IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"}
 MODEL_ID = "diffusers/FLUX.2-dev-bnb-4bit"
 TORCH_DTYPE = torch.bfloat16
 
-
-# =========================
-# HF Spaces detection (robust)
-# =========================
-try:
-    import spaces  # Only available on Hugging Face Spaces
-
-    SPACES_AVAILABLE = True
-except Exception:
-    spaces = None
-    SPACES_AVAILABLE = False
-
-# 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"))
-)
-
-
-# =========================
 # Global cached pipeline
-# =========================
 _pipe = None
 
 
@@ -89,7 +96,6 @@ def calculate_psnr_ssim(original: Image.Image, enhanced: Image.Image):
 def extract_archive(archive_path: str, extract_to: str):
     """Extract zip or 7z archive."""
     file_ext = Path(archive_path).suffix.lower()
-
     if file_ext == ".zip":
         with zipfile.ZipFile(archive_path, "r") as z:
             z.extractall(extract_to)
@@ -113,8 +119,9 @@ def find_images(directory: str):
 def _get_pipe():
     """
     Lazy-load the pipeline.
-    - On HF ZeroGPU: this MUST be called inside a @spaces.GPU function runtime.
-    - Locally: uses local CUDA GPU.
+    IMPORTANT:
+      - On HF ZeroGPU: must be called inside a @spaces.GPU-decorated function runtime.
+      - Locally: uses local CUDA.
     """
     global _pipe
 
@@ -123,7 +130,7 @@ def _get_pipe():
             raise RuntimeError(
                 "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."
+                "On HF ZeroGPU, ensure inference is inside a @spaces.GPU function."
             )
 
         _pipe = Flux2Pipeline.from_pretrained(
@@ -139,9 +146,7 @@ def _process_images_impl(
     files, prompt, guidance_scale, num_steps, progress=gr.Progress()
 ):
     """
-    Shared implementation used by both:
-      - HF Spaces GPU wrapper
-      - Local runtime
+    Shared implementation.
     Returns 4 outputs:
       gallery, files_download, zip_download, summary
     """
@@ -162,8 +167,7 @@ def _process_images_impl(
     # Temp for extraction/staging input
     temp_dir = tempfile.mkdtemp(prefix="flux_in_")
 
-    # IMPORTANT:
-    # Output files MUST remain on disk for Gradio downloads.
+    # Output MUST remain for Gradio downloads
     run_id = uuid.uuid4().hex[:10]
     output_dir = tempfile.mkdtemp(prefix=f"flux_results_{run_id}_")
 
@@ -172,7 +176,7 @@ def _process_images_impl(
     try:
         progress(0.0, desc="Preparing files...")
 
-        all_images = []  # list of tuples: (img_path, rel_path, base_dir_for_rel)
+        all_images = []  # (img_path, rel_path)
         for file_obj in files:
             file_path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
             file_ext = Path(file_path).suffix.lower()
@@ -180,6 +184,7 @@ def _process_images_impl(
             if file_ext in [".zip", ".7z"]:
                 has_archive = True
                 progress(0.05, desc=f"Extracting: {Path(file_path).name} ...")
+
                 extract_dir = os.path.join(temp_dir, Path(file_path).stem)
                 os.makedirs(extract_dir, exist_ok=True)
                 extract_archive(file_path, extract_dir)
@@ -187,10 +192,10 @@ def _process_images_impl(
                 images = find_images(extract_dir)
                 for img_path in images:
                     rel_path = os.path.relpath(img_path, extract_dir)
-                    all_images.append((img_path, rel_path, extract_dir))
+                    all_images.append((img_path, rel_path))
 
             elif file_ext in IMAGE_EXTENSIONS:
-                all_images.append((file_path, Path(file_path).name, None))
+                all_images.append((file_path, Path(file_path).name))
 
         if not all_images:
             return (
@@ -203,14 +208,14 @@ def _process_images_impl(
         total_images = len(all_images)
         progress(0.10, desc=f"Found {total_images} images. Loading model...")
 
-        pipe = _get_pipe()  # IMPORTANT: must be inside GPU runtime on Spaces
+        pipe = _get_pipe()
 
         results = []
         metrics_lines = []
 
         progress(0.15, desc="Enhancing images...")
 
-        for idx, (img_path, rel_path, base_dir) in enumerate(all_images):
+        for idx, (img_path, rel_path) in enumerate(all_images):
             progress(
                 0.15 + 0.75 * (idx / max(1, total_images)),
                 desc=f"Processing {idx+1}/{total_images}: {Path(img_path).name}",
@@ -227,12 +232,10 @@ def _process_images_impl(
 
             psnr, ssim = calculate_psnr_ssim(input_image, enhanced_image)
 
-            # Preserve structure if from archive
-            output_rel_path = rel_path
-            out_path = os.path.join(output_dir, output_rel_path)
+            out_path = os.path.join(output_dir, rel_path)
             os.makedirs(os.path.dirname(out_path), exist_ok=True)
 
-            # Add _flux suffix
+            # add _flux suffix
             out_name = Path(out_path).stem + "_flux" + Path(out_path).suffix
             out_path = os.path.join(os.path.dirname(out_path), out_name)
 
@@ -242,23 +245,21 @@ def _process_images_impl(
                 {
                     "original": input_image,
                     "enhanced": enhanced_image,
-                    "filename": output_rel_path,
+                    "filename": rel_path,
                     "output_path": out_path,
                     "psnr": psnr,
                     "ssim": ssim,
                 }
             )
 
-            metrics_lines.append(
-                f"{output_rel_path}: PSNR={psnr:.2f} dB, SSIM={ssim:.4f}"
-            )
+            metrics_lines.append(f"{rel_path}: PSNR={psnr:.2f} dB, SSIM={ssim:.4f}")
 
         avg_psnr = float(np.mean([r["psnr"] for r in results])) if results else 0.0
         avg_ssim = float(np.mean([r["ssim"] for r in results])) if results else 0.0
 
         summary = (
             "✅ Processing completed!\n\n"
-            f"Environment: {'Hugging Face Spaces' if IS_HF_SPACES else 'Local'}\n"
+            f"Environment: {'Hugging Face Spaces' if IS_HF_SPACES else 'Local/Unknown'}\n"
             f"GPU available: {torch.cuda.is_available()}\n\n"
             f"Total images processed: {total_images}\n"
             f"Average PSNR: {avg_psnr:.2f} dB\n"
@@ -277,16 +278,12 @@ def _process_images_impl(
                 )
             )
 
-        # Download behavior:
-        # - If any archive uploaded -> zip
-        # - Else -> direct files list
         if has_archive:
             progress(0.92, desc="Packaging ZIP...")
 
             output_zip_path = os.path.join(
                 tempfile.gettempdir(), f"enhanced_images_flux_{run_id}.zip"
             )
-
             with zipfile.ZipFile(output_zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
                 for root, _, fs in os.walk(output_dir):
                     for f in fs:
@@ -297,8 +294,8 @@ def _process_images_impl(
             progress(1.0, desc="Done!")
             return (
                 gallery_images,
-                gr.update(value=None, visible=False),  # files hidden
-                gr.update(value=output_zip_path, visible=True),  # zip shown
+                gr.update(value=None, visible=False),
+                gr.update(value=output_zip_path, visible=True),
                 summary,
             )
         else:
@@ -306,8 +303,8 @@ def _process_images_impl(
             progress(1.0, desc="Done!")
             return (
                 gallery_images,
-                gr.update(value=enhanced_paths, visible=True),  # files shown
-                gr.update(value=None, visible=False),  # zip hidden
+                gr.update(value=enhanced_paths, visible=True),
+                gr.update(value=None, visible=False),
                 summary,
             )
 
@@ -320,30 +317,19 @@ def _process_images_impl(
         )
 
     finally:
-        # Clean ONLY input extraction temp.
+        # Cleanup input temp only
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir, ignore_errors=True)
-        # Do NOT delete output_dir; downloads need it to exist.
+        # DO NOT delete output_dir; needed for downloads
 
 
 # =========================
-# IMPORTANT: Define a real @spaces.GPU function at import-time on Spaces
-# (This fixes: "No @spaces.GPU function detected during startup")
+# CRITICAL: Always define a @spaces.GPU function at top-level
+# (ZeroGPU startup scanner will now ALWAYS find it)
 # =========================
-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)
+@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)
 
 
 # =========================
@@ -364,9 +350,9 @@ Upload microscopy images (individual files or compressed archives) for AI-powere
 - 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}**
+**Runtime info (informational):**
+- Detected environment: **{"Hugging Face Spaces" if IS_HF_SPACES else "Local/Unknown"}**
+- CUDA visible at startup: **{torch.cuda.is_available()}**
 """
     )