update

InfiniDepth/utils/hf_demo_utils.py

@@ -2,7 +2,7 @@ import os
 import tempfile
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import cv2
 import numpy as np
@@ -129,6 +129,11 @@ def prepare_runtime_assets() -> None:
     resolve_moge2_pretrained()
 
 
+def _report_stage(stage_callback: Optional[Callable[[str], None]], stage: str) -> None:
+    if stage_callback is not None:
+        stage_callback(stage)
+
+
 @dataclass
 class DemoResult:
     depth_vis: np.ndarray
@@ -310,8 +315,10 @@ def run_single_image_demo(
     cx_org: Optional[float] = None,
     cy_org: Optional[float] = None,
     model_cache: Optional[ModelCache] = None,
+    stage_callback: Optional[Callable[[str], None]] = None,
 ) -> DemoResult:
     image_shape = tuple(int(dim) for dim in image_np.shape) if image_np is not None else None
+    _report_stage(stage_callback, "demo:start")
     Log.info(
         "run_single_image_demo start: "
         f"model_type={model_type}, input_size={input_size_text}, output_resolution_mode={output_resolution_mode}, "
@@ -323,6 +330,7 @@ def run_single_image_demo(
             "No CUDA GPU is available. If using Hugging Face ZeroGPU, "
             "decorate the Gradio inference function with @spaces.GPU and enable queue()."
         )
+    _report_stage(stage_callback, "demo:cuda_ready")
 
     input_size = _parse_image_size(input_size_text)
     if upsample_ratio < 1 or upsample_ratio > 8:
@@ -334,6 +342,7 @@ def run_single_image_demo(
 
     device = torch.device("cuda")
     image, org_h, org_w = _prepare_image_tensor(image_np, input_size, device)
+    _report_stage(stage_callback, "demo:image_prepared")
     h_in, w_in = input_size
     h_out, w_out = resolve_output_size_from_mode(
         output_resolution_mode=output_resolution_mode,
@@ -354,16 +363,20 @@ def run_single_image_demo(
         image=image,
         device=device,
     )
+    _report_stage(stage_callback, f"demo:depth_inputs_ready source={depth_source_label}")
     Log.info(f"Depth source resolved: {depth_source_label}")
     gt = depth_to_disparity(gt_depth)
     prompt = depth_to_disparity(prompt_depth)
     prompt_mask = prompt > 0
 
     ckpt_path = resolve_checkpoint_path(model_type)
+    _report_stage(stage_callback, f"demo:checkpoint_resolved path={ckpt_path}")
     model_cache = model_cache or ModelCache()
     model = model_cache.get(model_type=model_type, model_path=ckpt_path)
+    _report_stage(stage_callback, "demo:model_loaded")
 
     query_2d_uniform_coord = make_2d_uniform_coord((h_out, w_out)).unsqueeze(0).to(device)
+    _report_stage(stage_callback, "demo:inference_started")
     pred_depth, _ = model.inference(
         image=image,
         query_coord=query_2d_uniform_coord,
@@ -372,10 +385,12 @@ def run_single_image_demo(
         prompt_depth=prompt,
         prompt_mask=prompt_mask,
     )
+    _report_stage(stage_callback, "demo:inference_finished")
     Log.info(f"Model inference completed: output_size={h_out}x{w_out}")
 
     pred_depthmap = pred_depth.permute(0, 2, 1).reshape(1, 1, h_out, w_out)
     depth_vis = _colorize_predicted_depth(pred_depthmap)
+    _report_stage(stage_callback, "demo:depth_colorized")
 
     fx, fy, cx, cy = resolve_camera_intrinsics(fx_org, fy_org, cx_org, cy_org, org_h, org_w)
     fx_out, fy_out, cx_out, cy_out, _ = build_scaled_intrinsics_matrix(
@@ -406,6 +421,7 @@ def run_single_image_demo(
         output_path=ply_path,
         max_points_preview=int(max_points_preview),
     )
+    _report_stage(stage_callback, "demo:pointcloud_saved")
     Log.info(
         f"Artifacts saved: depth_npy_path={depth_npy_path}, ply_path={ply_path}, preview_points={xyz.shape[0]}"
     )

app.py

@@ -34,6 +34,7 @@ except ImportError:
 
 MODEL_CACHE = ModelCache()
 OUTPUT_ROOT = Path(tempfile.gettempdir()) / "infinidepth_hf_demo"
+TRACE_ROOT = OUTPUT_ROOT / "trace"
 EXAMPLE_DATA_ROOT = Path(__file__).resolve().parent / "example_data"
 EXAMPLE_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
 EXAMPLE_DEPTH_EXTENSIONS = {".png", ".npy", ".npz", ".h5", ".hdf5", ".exr"}
@@ -202,6 +203,22 @@ def _prepare_output_dir() -> Path:
     return output_dir
 
 
+def _append_trace(trace_path: str, stage: str) -> None:
+    trace_file = Path(trace_path)
+    trace_file.parent.mkdir(parents=True, exist_ok=True)
+    with trace_file.open("a", encoding="utf-8") as handle:
+        handle.write(f"{stage}\n")
+        handle.flush()
+
+
+def _read_trace(trace_path: str, max_lines: int = 12) -> str:
+    trace_file = Path(trace_path)
+    if not trace_file.exists():
+        return ""
+    lines = trace_file.read_text(encoding="utf-8").splitlines()
+    return "\n".join(lines[-max_lines:])
+
+
 def _export_glb_from_points(xyz: np.ndarray, rgb: np.ndarray, output_path: Path) -> None:
     if trimesh is None:
         raise RuntimeError("`trimesh` is required to export .glb for the 3D viewer")
@@ -226,7 +243,9 @@ def _run_demo_impl(
     fy_org: Optional[float],
     cx_org: Optional[float],
     cy_org: Optional[float],
+    trace_path: str,
 ):
+    _append_trace(trace_path, "worker:entered _run_demo_impl")
     if image is None:
         raise ValueError("Input RGB image is required")
 
@@ -234,6 +253,7 @@ def _run_demo_impl(
     if depth_file is not None:
         depth_path = depth_file if isinstance(depth_file, str) else depth_file.name
 
+    _append_trace(trace_path, f"worker:run_single_image_demo depth_path={depth_path}")
     result = run_single_image_demo(
         image_np=image,
         depth_path=depth_path,
@@ -247,12 +267,15 @@ def _run_demo_impl(
         cx_org=_none_if_invalid(cx_org),
         cy_org=_none_if_invalid(cy_org),
         model_cache=MODEL_CACHE,
+        stage_callback=lambda stage: _append_trace(trace_path, stage),
     )
+    _append_trace(trace_path, "worker:prepare_output_dir")
     output_dir = _prepare_output_dir()
     glb_path = output_dir / "pointcloud.glb"
     ply_path = output_dir / "pointcloud.ply"
     depth_vis_path = output_dir / "depth_colorized.png"
 
+    _append_trace(trace_path, "worker:export_glb")
     _export_glb_from_points(result.xyz, result.rgb, glb_path)
     if os.path.exists(result.ply_path):
         shutil.copy2(result.ply_path, ply_path)
@@ -275,6 +298,7 @@ def _run_demo_impl(
         f"Preview points: {result.xyz.shape[0]}. "
         f"Camera intrinsics: {'default' if result.used_default_intrinsics else 'custom'}."
     )
+    _append_trace(trace_path, "worker:completed")
     return result.depth_vis, glb_path.as_posix(), download_files, status
 
 
@@ -291,6 +315,7 @@ def run_demo_gpu(
     fy_org: Optional[float],
     cx_org: Optional[float],
     cy_org: Optional[float],
+    trace_path: str,
 ):
     return _run_demo_impl(
         image=image,
@@ -304,6 +329,7 @@ def run_demo_gpu(
         fy_org=fy_org,
         cx_org=cx_org,
         cy_org=cy_org,
+        trace_path=trace_path,
     )
 
 
@@ -321,10 +347,12 @@ def run_demo(
     cy_org: Optional[float],
 ):
     request_id = uuid.uuid4().hex[:8]
+    trace_path = (TRACE_ROOT / f"{request_id}.log").as_posix()
     depth_path = None
     if depth_file is not None:
         depth_path = depth_file if isinstance(depth_file, str) else depth_file.name
     image_shape = tuple(int(dim) for dim in image.shape) if image is not None else None
+    _append_trace(trace_path, "ui:entered run_demo")
     Log.info(
         f"[{request_id}] run_demo start: model_type={model_type}, "
         f"input_size={input_size}, output_resolution_mode={output_resolution_mode}, "
@@ -344,14 +372,18 @@ def run_demo(
             fy_org=fy_org,
             cx_org=cx_org,
             cy_org=cy_org,
+            trace_path=trace_path,
         )
         status_text = result[3] if isinstance(result, tuple) and len(result) == 4 else "Done."
         Log.info(f"[{request_id}] run_demo success: {status_text}")
         return result
     except Exception as exc:
         error_trace = traceback.format_exc()
+        trace_summary = _read_trace(trace_path)
         Log.exception(f"[{request_id}] run_demo failed")
         error_message = f"Error [{request_id}]: {exc}"
+        if trace_summary:
+            error_message = f"{error_message}\n\nLast worker stages:\n{trace_summary}"
         if os.getenv("INFINIDEPTH_SHOW_TRACEBACK", "0") == "1":
             error_message = f"{error_message}\n\n{error_trace}"
         return None, None, [], error_message