update

InfiniDepth/model/block/pe.py

@@ -7,11 +7,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from typing import Any, Optional, Tuple, Dict
 
-acc_dtype = (
-    torch.bfloat16
-    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
-    else torch.float16
-)
+def _get_acc_dtype() -> torch.dtype:
+    """Determine autocast dtype lazily so ZeroGPU can set up CUDA first."""
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
+        return torch.bfloat16
+    return torch.float16
 
 POS_EMB_REGISTRY = {}
 

InfiniDepth/model/block/prompt_models/rope.py

@@ -3,11 +3,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-acc_dtype = (
-    torch.bfloat16
-    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
-    else torch.float16
-)
+def _get_acc_dtype() -> torch.dtype:
+    """Determine autocast dtype lazily so ZeroGPU can set up CUDA first."""
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
+        return torch.bfloat16
+    return torch.float16
 
 
 class PositionGetter:
@@ -186,7 +186,7 @@ class RotaryPositionEmbedding2D(nn.Module):
         """Positionally encode points that are normalized to [0,1]."""
         # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
         max_position = int(coords.max()) + 1
-        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, coords.device, acc_dtype)
+        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, coords.device, _get_acc_dtype())
         vertical_cos = F.embedding(coords[..., 0], cos_comp)
         vertical_sin = F.embedding(coords[..., 0], sin_comp)
         horizontal_cos = F.embedding(coords[..., 1], cos_comp)
@@ -201,7 +201,7 @@ class RotaryPositionEmbedding2D(nn.Module):
         x_coords = torch.arange(width, device=device) * (self.patch_size * 2) + self.patch_size - 1
         positions = torch.cartesian_prod(y_coords, x_coords)  # h, w
         max_position = int(positions.max()) + 1
-        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, device, acc_dtype)
+        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, device, _get_acc_dtype())
         vertical_cos = F.embedding(positions[..., 0], cos_comp)
         vertical_sin = F.embedding(positions[..., 0], sin_comp)
         horizontal_cos = F.embedding(positions[..., 1], cos_comp)

InfiniDepth/model/block/prompt_models/selfattn.py

@@ -6,11 +6,11 @@ from .rope import RotaryPositionEmbedding2D
 from .utils.pe_utils import PositionEmbeddingRandom
 from torch import Tensor
 
-acc_dtype = (
-    torch.bfloat16
-    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
-    else torch.float16
-)
+def _get_acc_dtype() -> torch.dtype:
+    """Determine autocast dtype lazily so ZeroGPU can set up CUDA first."""
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
+        return torch.bfloat16
+    return torch.float16
 
 
 class SelfAttnPromptModel(nn.Module):
@@ -97,7 +97,7 @@ class SelfAttnPromptModel(nn.Module):
             query_pe = image_pe.reshape(1, -1, image_pe.shape[-1])
             prompt = prompt_embeddings  # + prompt_pe
             query = image_embeddings[b : (b + 1)]  # + query_pe
-            with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
+            with torch.autocast("cuda", enabled=True, dtype=_get_acc_dtype()):
                 for block in self.prompt_blocks:
                     query, prompt = block(query, query_pe, prompt, prompt_pe)
             image_embeddings_list.append(query[..., : image_embeddings.shape[-1]])
@@ -197,7 +197,7 @@ class SelfAttnRopePromptModel(nn.Module):
             query_pe = image_pe.reshape(1, -1, image_pe.shape[-1])
             prompt = prompt_embeddings  # + prompt_pe
             query = image_embeddings[b : (b + 1)]  # + query_pe
-            with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
+            with torch.autocast("cuda", enabled=True, dtype=_get_acc_dtype()):
                 for block in self.prompt_blocks:
                     query, prompt = block(query, query_pe, prompt, prompt_pe)
             image_embeddings_list.append(query[..., : image_embeddings.shape[-1]])

InfiniDepth/model/model.py

@@ -17,11 +17,11 @@ from .block.prompt_models import GeneralPromptModel, SelfAttnPromptModel
 from .block.implicit_decoder import ImplicitHead
 from .block.convolution import BasicEncoder
 
-acc_dtype = (
-    torch.bfloat16
-    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
-    else torch.float16
-)
+def _get_acc_dtype() -> torch.dtype:
+    """Determine autocast dtype lazily so ZeroGPU can set up CUDA first."""
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
+        return torch.bfloat16
+    return torch.float16
 
 
 def _resolve_local_dinov3_repo() -> str:
@@ -96,7 +96,8 @@ class _BaseInfiniDepthModel(nn.Module):
                 raise FileNotFoundError(f"Model file {model_path} not found")
 
         # only for inference
-        self.cuda()
+        if torch.cuda.is_available():
+            self.cuda()
         self.eval()
 
     def _init_variant_modules(self):
@@ -131,7 +132,7 @@ class _BaseInfiniDepthModel(nn.Module):
     ):
         h, w = x.shape[-2:]
         x_dino = (x - self._mean) / self._std
-        with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
+        with torch.autocast("cuda", enabled=True, dtype=_get_acc_dtype()):
             features = self.pretrained.get_intermediate_layers(
                 x_dino,
                 n=self.model_config["layer_idxs"],

app.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 import tempfile
+import time
 import traceback
 import uuid
 from pathlib import Path
@@ -246,7 +247,22 @@ def run_demo_gpu(
     trace_path: str,
 ):
     """GPU-only inference. Returns a GPUInferenceResult with all data on CPU."""
+    import torch
     _append_trace(trace_path, "worker:entered run_demo_gpu")
+
+    # Log GPU info for diagnostics (visible in HF Space Logs tab)
+    if torch.cuda.is_available():
+        dev = torch.cuda.current_device()
+        gpu_name = torch.cuda.get_device_name(dev)
+        gpu_mem_total = torch.cuda.get_device_properties(dev).total_mem / 1e9
+        gpu_mem_alloc = torch.cuda.memory_allocated(dev) / 1e9
+        Log.info(f"[GPU-INFO] device={gpu_name}, total={gpu_mem_total:.1f}GB, allocated={gpu_mem_alloc:.1f}GB")
+        _append_trace(trace_path, f"worker:cuda_ready {gpu_name} {gpu_mem_total:.1f}GB")
+        torch.cuda.empty_cache()
+    else:
+        Log.warning("[GPU-INFO] CUDA not available inside @spaces.GPU!")
+        _append_trace(trace_path, "worker:cuda_NOT_available")
+
     if image is None:
         raise ValueError("Input RGB image is required")
 
@@ -298,19 +314,35 @@ def run_demo(
     )
     try:
         # --- GPU-only inference (consumes ZeroGPU quota) ---
-        gpu_result = run_demo_gpu(
-            image=image,
-            depth_file=depth_file,
-            model_type=model_type,
-            input_size=input_size,
-            output_resolution_mode=output_resolution_mode,
-            upsample_ratio=upsample_ratio,
-            fx_org=fx_org,
-            fy_org=fy_org,
-            cx_org=cx_org,
-            cy_org=cy_org,
-            trace_path=trace_path,
-        )
+        # Retry on transient "GPU task aborted" errors (common on ZeroGPU)
+        max_gpu_retries = 2
+        gpu_result = None
+        for attempt in range(max_gpu_retries + 1):
+            try:
+                gpu_result = run_demo_gpu(
+                    image=image,
+                    depth_file=depth_file,
+                    model_type=model_type,
+                    input_size=input_size,
+                    output_resolution_mode=output_resolution_mode,
+                    upsample_ratio=upsample_ratio,
+                    fx_org=fx_org,
+                    fy_org=fy_org,
+                    cx_org=cx_org,
+                    cy_org=cy_org,
+                    trace_path=trace_path,
+                )
+                break  # success
+            except Exception as gpu_exc:
+                is_aborted = "GPU task aborted" in str(gpu_exc)
+                if is_aborted and attempt < max_gpu_retries:
+                    Log.warning(
+                        f"[{request_id}] GPU task aborted (attempt {attempt + 1}/{max_gpu_retries + 1}), retrying..."
+                    )
+                    _append_trace(trace_path, f"ui:retry_{attempt + 1}_after_gpu_abort")
+                    time.sleep(2)
+                    continue
+                raise  # not retryable or out of retries
         _append_trace(trace_path, "ui:gpu_done, starting cpu postprocess")
 
         # --- CPU post-processing (no GPU quota consumed) ---