Extend and trace ZeroGPU depth inference

app.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from functools import lru_cache
 import os
 from pathlib import Path
+import time
 from typing import Optional
 
 
@@ -110,7 +111,7 @@ MODEL_CHOICES = [RGB_MODEL_TYPE, DEPTH_SENSOR_MODEL_TYPE]
 OUTPUT_MODE_CHOICES = ["upsample", "original", "specific"]
 GS_SAMPLE_POINT_NUM = 2000000
 GS_COORD_DETERMINISTIC_SAMPLING = True
-DEPTH_GPU_DURATION_SECONDS = 180
+DEPTH_GPU_DURATION_SECONDS = 600
 GS_GPU_DURATION_SECONDS = 240
 
 LOCAL_DEPTH_MODEL_PATHS = {
@@ -476,7 +477,14 @@ def _run_depth_inference(
     cy_org: Optional[float],
     request: gr.Request,
 ):
+    started_at = time.perf_counter()
+
+    def log_stage(stage: str) -> None:
+        elapsed = time.perf_counter() - started_at
+        print(f"[INFO] Depth inference {stage} after {elapsed:.1f}s", flush=True)
+
     _ensure_cuda()
+    log_stage(f"started on {torch.cuda.get_device_name(0)}")
     if image is None:
         raise gr.Error("Upload an image or load an example before running inference.")
     if model_type == DEPTH_SENSOR_MODEL_TYPE and not depth_file:
@@ -488,6 +496,7 @@ def _run_depth_inference(
     device = torch.device("cuda")
     image_tensor = image_tensor.to(device)
     model = _load_model(model_type)
+    log_stage("prepared input and model")
 
     gt_depth, prompt_depth, gt_depth_mask, use_gt_depth, moge2_intrinsics = prepare_metric_depth_inputs(
         input_depth_path=depth_file,
@@ -496,6 +505,7 @@ def _run_depth_inference(
         device=device,
         moge2_pretrained=_resolve_moge2_source(),
     )
+    log_stage("prepared metric depth")
 
     gt_disp = depth_to_disparity(gt_depth)
     prompt_disp = depth_to_disparity(prompt_depth)
@@ -511,6 +521,7 @@ def _run_depth_inference(
         moge2_pretrained=_resolve_moge2_source(),
         moge2_intrinsics=moge2_intrinsics,
     )
+    log_stage("resolved intrinsics")
 
     _, _, h, w = image_tensor.shape
     fx, fy, cx, cy, _ = build_scaled_intrinsics_matrix(
@@ -530,6 +541,7 @@ def _run_depth_inference(
         enable_skyseg_model=enable_skyseg_model,
         sky_model_ckpt_path=skyseg_path or str(LOCAL_SKYSEG_PATH),
     )
+    log_stage("prepared sky mask")
 
     h_out, w_out = resolve_output_size_from_mode(
         output_resolution_mode=output_resolution_mode,
@@ -550,6 +562,7 @@ def _run_depth_inference(
         prompt_depth=prompt_disp,
         prompt_mask=prompt_disp > 0,
     )
+    log_stage("completed model inference")
     pred_depthmap = pred_2d_uniform_depth.permute(0, 2, 1).reshape(1, 1, h_out, w_out)
 
     pred_depthmap, pred_2d_uniform_depth = apply_sky_mask_to_depth(
@@ -566,6 +579,7 @@ def _run_depth_inference(
 
     pred_depth_np = pred_depthmap.squeeze(0).squeeze(0).detach().cpu().numpy().astype(np.float32)
     artifacts = save_demo_artifacts(image_rgb=image_rgb, pred_depth=pred_depth_np, output_dir=output_dir)
+    log_stage("saved depth artifacts")
     ply_path, glb_path = export_point_cloud_assets(
         sampled_coord=query_2d_uniform_coord.squeeze(0).cpu(),
         sampled_depth=pred_2d_uniform_depth.squeeze(0).squeeze(-1).cpu(),
@@ -577,6 +591,7 @@ def _run_depth_inference(
         output_dir=output_dir,
         filter_flying_points=filter_point_cloud,
     )
+    log_stage("exported point cloud")
     artifacts = DemoArtifacts(
         comparison_path=artifacts.comparison_path,
         color_depth_path=artifacts.color_depth_path,