Update app.py

app.py

@@ -4,15 +4,24 @@
 # ────────────────────────────────────────────────────────
 import sys, types
 try:
-    # If old path exists, do nothing
     import torchvision.transforms.functional_tensor as _ft  # noqa: F401
 except Exception:
-    # Map to the new API location
     from torchvision.transforms import functional as _F
     _mod = types.ModuleType("torchvision.transforms.functional_tensor")
     _mod.rgb_to_grayscale = _F.rgb_to_grayscale
     sys.modules["torchvision.transforms.functional_tensor"] = _mod
 
+# ────────────────────────────────────────────────────────
+# Spaces ZeroGPU decorator (safe no-op locally)
+# ────────────────────────────────────────────────────────
+try:
+    import spaces
+    GPU = spaces.GPU
+except Exception:
+    def GPU(*args, **kwargs):
+        def _wrap(f): return f
+        return _wrap
+
 # ────────────────────────────────────────────────────────
 # Standard imports
 # ────────────────────────────────────────────────────────
@@ -105,7 +114,9 @@ def model_tip_text(model_name: str) -> str:
 
 # ────────────────────────────────────────────────────────
 # Core upscaling
+# Decorated for Hugging Face Spaces ZeroGPU
 # ────────────────────────────────────────────────────────
+@GPU()  # <-- lets Spaces know this function uses GPU; safe no-op locally
 def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
     if img is None:
         return
@@ -141,13 +152,11 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
     weights_dir = os.path.join(ROOT_DIR, 'weights')
     os.makedirs(weights_dir, exist_ok=True)
 
-    local_paths = []
     for url in file_url:
         fname = os.path.basename(url)
         local_path = os.path.join(weights_dir, fname)
         if not os.path.isfile(local_path):
-            local_path = load_file_from_url(url=url, model_dir=weights_dir, progress=True)
-        local_paths.append(local_path)
+            load_file_from_url(url=url, model_dir=weights_dir, progress=True)
 
     if model_name == 'realesr-general-x4v3':
         base_path = os.path.join(weights_dir, 'realesr-general-x4v3.pth')
@@ -172,7 +181,7 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
         model_path=model_path,
         dni_weight=dni_weight,
         model=model,
-        tile=256,
+        tile=256,       # VRAM-safe default; lower to 128 if OOM
         tile_pad=10,
         pre_pad=10,
         half=bool(use_cuda),
@@ -274,6 +283,7 @@ def main():
         input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
         model_name.change(fn=model_tip_text, inputs=model_name, outputs=model_tips)
 
+        # IMPORTANT: call the @GPU-decorated function here
         restore_btn.click(
             fn=realesrgan,
             inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
@@ -283,7 +293,8 @@ def main():
 
         gr.Markdown("")  # spacer
 
-    demo.launch()
+    # Disable SSR as suggested by the log
+    demo.launch(ssr_mode=False)  # share=True if you want a public link
 
 
 if __name__ == "__main__":