fix bug for output

InfiniDepth/utils/hf_demo_utils.py

@@ -7,6 +7,7 @@ import os
 import re
 import shutil
 import tempfile
+import time
 import uuid
 
 import cv2
@@ -20,6 +21,7 @@ from .vis_utils import colorize_depth_maps
 
 _DEPTH_EXTENSIONS = {".png", ".npz", ".npy", ".h5", ".hdf5", ".exr"}
 _IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+_SAFE_PATH_PART = re.compile(r"[^A-Za-z0-9._-]+")
 
 
 @dataclass(frozen=True)
@@ -98,9 +100,45 @@ def scan_example_cases(example_dir: str | Path) -> list[ExampleCase]:
     return cases
 
 
-def ensure_session_output_dir(app_name: str, session_hash: Optional[str]) -> Path:
-    safe_hash = session_hash or f"session-{uuid.uuid4().hex[:10]}"
-    output_dir = Path(tempfile.gettempdir()) / app_name / safe_hash
+def _sanitize_path_part(value: str, fallback_prefix: str) -> str:
+    sanitized = _SAFE_PATH_PART.sub("-", value).strip(".-")
+    if sanitized:
+        return sanitized
+    return f"{fallback_prefix}-{uuid.uuid4().hex[:10]}"
+
+
+def _prune_stale_run_dirs(session_dir: Path, max_age_seconds: int) -> None:
+    if max_age_seconds <= 0 or not session_dir.exists():
+        return
+
+    now = time.time()
+    for child in session_dir.iterdir():
+        if not child.is_dir():
+            continue
+        try:
+            age_seconds = now - child.stat().st_mtime
+        except FileNotFoundError:
+            continue
+        if age_seconds > max_age_seconds:
+            shutil.rmtree(child, ignore_errors=True)
+
+
+def ensure_session_output_dir(
+    app_name: str,
+    session_hash: Optional[str],
+    run_id: Optional[str] = None,
+    max_age_seconds: int = 60 * 60,
+) -> Path:
+    safe_app_name = _sanitize_path_part(app_name, "app")
+    raw_hash = session_hash or f"session-{uuid.uuid4().hex[:10]}"
+    safe_hash = _sanitize_path_part(raw_hash, "session")
+    session_dir = Path(tempfile.gettempdir()) / safe_app_name / safe_hash
+    session_dir.mkdir(parents=True, exist_ok=True)
+    _prune_stale_run_dirs(session_dir, max_age_seconds=max_age_seconds)
+
+    raw_run_id = run_id or f"run-{time.strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:8]}"
+    safe_run_id = _sanitize_path_part(raw_run_id, "run")
+    output_dir = session_dir / safe_run_id
     shutil.rmtree(output_dir, ignore_errors=True)
     output_dir.mkdir(parents=True, exist_ok=True)
     return output_dir

InfiniDepth/utils/hf_gs_viewer.py

@@ -132,17 +132,25 @@ def ensure_supersplat_viewer_assets() -> Path:
     return _copy_viewer_dist(dist_dir)
 
 
-def _gradio_file_url(path: str | Path) -> str:
-    return "/gradio_api/file=" + quote(str(Path(path).resolve()), safe="/")
-
-
-def build_embedded_viewer_html(ply_path: str | Path, height_px: int = 700) -> str:
+def _gradio_file_url(path: str | Path, cache_key: str | None = None) -> str:
+    url = "/gradio_api/file=" + quote(str(Path(path).resolve()), safe="/")
+    if cache_key:
+        url += f"?v={quote(cache_key, safe='')}"
+    return url
+
+
+def build_embedded_viewer_html(
+    ply_path: str | Path,
+    height_px: int = 700,
+    cache_key: str | None = None,
+) -> str:
     viewer_dir = ensure_supersplat_viewer_assets()
-    viewer_url = _gradio_file_url(viewer_dir / "index.html")
-    splat_url = _gradio_file_url(ply_path)
+    viewer_url = _gradio_file_url(viewer_dir / "index.html", cache_key=cache_key)
+    splat_url = _gradio_file_url(ply_path, cache_key=cache_key)
     filename = Path(ply_path).name
+    query_join = "&" if "?" in viewer_url else "?"
     iframe_src = (
-        f"{viewer_url}?load={quote(splat_url, safe='')}&filename={quote(filename, safe='')}"
+        f"{viewer_url}{query_join}load={quote(splat_url, safe='')}&filename={quote(filename, safe='')}"
     )
     escaped_src = html.escape(iframe_src, quote=True)
     return (
@@ -155,9 +163,9 @@ def build_embedded_viewer_html(ply_path: str | Path, height_px: int = 700) -> st
     )
 
 
-def build_viewer_error_html(message: str, ply_path: str | Path) -> str:
+def build_viewer_error_html(message: str, ply_path: str | Path, cache_key: str | None = None) -> str:
     escaped_message = html.escape(message)
-    ply_url = html.escape(_gradio_file_url(ply_path), quote=True)
+    ply_url = html.escape(_gradio_file_url(ply_path, cache_key=cache_key), quote=True)
     return (
         '<div style="padding:1rem; border:1px solid #d0d7de; border-radius:12px; background:#fff;">'
         '<strong>Embedded SuperSplat viewer is unavailable.</strong>'

app.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from functools import lru_cache
 import os
 from pathlib import Path
+from time import perf_counter
 from typing import Optional
 
 
@@ -63,6 +64,7 @@ from InfiniDepth.utils.hf_gs_viewer import (
     APP_TEMP_ROOT as GS_VIEWER_ROOT,
     build_embedded_viewer_html,
     build_viewer_error_html,
+    ensure_supersplat_viewer_assets,
 )
 from InfiniDepth.utils.inference_utils import (
     apply_sky_mask_to_depth,
@@ -217,7 +219,12 @@ def _resolve_moge2_source() -> str:
 def _preload_repo_assets() -> tuple[str, ...]:
     depth_paths = tuple(_resolve_depth_checkpoint(model_type) for model_type in MODEL_CHOICES)
     gs_paths = tuple(_resolve_gs_checkpoint(model_type) for model_type in MODEL_CHOICES)
-    return depth_paths + gs_paths + (_resolve_moge2_source(), _resolve_skyseg_path())
+    viewer_paths: tuple[str, ...] = ()
+    try:
+        viewer_paths = (str(ensure_supersplat_viewer_assets()),)
+    except Exception as exc:
+        print(f"[Warning] Failed to preload SuperSplat viewer assets: {exc}")
+    return depth_paths + gs_paths + (_resolve_moge2_source(), _resolve_skyseg_path()) + viewer_paths
 
 
 @lru_cache(maxsize=2)
@@ -291,9 +298,10 @@ def _format_gs_status(
     intrinsics_source: str,
     depth_file: Optional[str],
     gaussian_count: int,
+    stage_timings: Optional[list[tuple[str, float]]] = None,
 ) -> str:
     depth_label = Path(depth_file).name if depth_file else "None"
-    return (
+    status = (
         f"Task: `GS`\n\n"
         f"Model: `{model_type}`\n\n"
         f"Input depth: `{depth_label}`\n\n"
@@ -301,6 +309,10 @@ def _format_gs_status(
         f"Camera intrinsics source: `{intrinsics_source}`\n\n"
         f"Exported gaussians: `{gaussian_count}`"
     )
+    if stage_timings:
+        timing_lines = "\n".join(f"- `{name}`: `{seconds:.2f}s`" for name, seconds in stage_timings)
+        status += f"\n\nStage timings:\n{timing_lines}"
+    return status
 
 
 def _load_example_image(example_name: str) -> tuple[np.ndarray, Optional[str], Optional[np.ndarray], str, str]:
@@ -375,6 +387,10 @@ def _normalize_filtered_gaussians(filtered_result):
     return filtered_result
 
 
+def _format_seconds(value: float) -> float:
+    return round(float(value), 2)
+
+
 @torch.no_grad()
 def _run_depth_inference(
     image: np.ndarray,
@@ -535,6 +551,7 @@ def _run_gs_inference(
     request: gr.Request,
 ):
     _ensure_cuda()
+    stage_timings: list[tuple[str, float]] = []
     if image is None:
         raise gr.Error("Upload an image or load an example before running inference.")
     if model_type == "InfiniDepth_DepthSensor" and not depth_file:
@@ -591,6 +608,7 @@ def _run_gs_inference(
         dilate_px=0,
     )
 
+    model_stage_start = perf_counter()
     depthmap, dino_tokens, query_3d_uniform_coord, pred_depth_3d = model.inference_for_gs(
         image=image_tensor,
         intrinsics=intrinsics,
@@ -627,13 +645,16 @@ def _run_gs_inference(
     gaussian_count = int(pixel_gaussians.means.shape[1])
     if gaussian_count == 0:
         raise gr.Error("No valid gaussians remained after filtering.")
+    stage_timings.append(("GS model inference", _format_seconds(perf_counter() - model_stage_start)))
 
     means, harmonics, opacities, scales, rotations = unpack_gaussians_for_export(pixel_gaussians)
 
     session_hash = getattr(request, "session_hash", None)
     output_dir = ensure_session_output_dir(APP_NAME, session_hash)
-    ply_path = output_dir / "gaussians.ply"
+    run_id = output_dir.name
+    ply_path = output_dir / f"{run_id}_gaussians.ply"
 
+    export_stage_start = perf_counter()
     export_ply(
         means=means,
         harmonics=harmonics,
@@ -646,12 +667,15 @@ def _run_gs_inference(
         image_shape=(org_h, org_w),
         extrinsic_matrix=extrinsics[0],
     )
+    stage_timings.append(("PLY export", _format_seconds(perf_counter() - export_stage_start)))
 
+    viewer_stage_start = perf_counter()
     try:
-        gs_viewer_html = build_embedded_viewer_html(ply_path)
+        gs_viewer_html = build_embedded_viewer_html(ply_path, cache_key=run_id)
     except Exception as exc:
         print(f"[Warning] Failed to build embedded GS viewer: {exc}")
-        gs_viewer_html = build_viewer_error_html(str(exc), ply_path)
+        gs_viewer_html = build_viewer_error_html(str(exc), ply_path, cache_key=run_id)
+    stage_timings.append(("Viewer HTML prep", _format_seconds(perf_counter() - viewer_stage_start)))
 
     metric_depth_source = "user depth" if use_gt_depth and depth_file else "MoGe-2"
     status = _format_gs_status(
@@ -660,6 +684,7 @@ def _run_gs_inference(
         intrinsics_source=intrinsics_source,
         depth_file=depth_file,
         gaussian_count=gaussian_count,
+        stage_timings=stage_timings,
     )
     download_files = [str(ply_path)]
 
@@ -837,6 +862,18 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
                         depth_download_files = gr.File(label="Depth Files", type="filepath")
                         gs_download_files = gr.File(label="GS Files", type="filepath")
 
+    result_outputs = [
+        status_output,
+        depth_comparison,
+        color_depth,
+        gray_depth,
+        depth_model_3d,
+        depth_download_files,
+        gs_viewer_html,
+        gs_download_files,
+    ]
+    example_load_outputs = [input_image, input_depth_file, input_depth_preview, model_type, depth_info]
+
     task_type.change(
         fn=_settings_visibility,
         inputs=[task_type, output_resolution_mode],
@@ -857,6 +894,9 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
     example_gallery.select(
         fn=_select_example,
         outputs=[selected_example_name, example_selection],
+    ).then(
+        fn=_clear_outputs,
+        outputs=result_outputs,
     )
 
     input_depth_file.change(
@@ -866,9 +906,12 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
     )
 
     load_example_btn.click(
+        fn=_clear_outputs,
+        outputs=result_outputs,
+    ).then(
         fn=_load_example_image,
         inputs=[selected_example_name],
-        outputs=[input_image, input_depth_file, input_depth_preview, model_type, depth_info],
+        outputs=example_load_outputs,
     )
 
     submit_btn.click(
@@ -877,16 +920,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
         outputs=[primary_view_tabs],
     ).then(
         fn=_clear_outputs,
-        outputs=[
-            status_output,
-            depth_comparison,
-            color_depth,
-            gray_depth,
-            depth_model_3d,
-            depth_download_files,
-            gs_viewer_html,
-            gs_download_files,
-        ],
+        outputs=result_outputs,
     ).then(
         fn=_run_inference,
         inputs=[
@@ -905,16 +939,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
             cx_org,
             cy_org,
         ],
-        outputs=[
-            status_output,
-            depth_comparison,
-            color_depth,
-            gray_depth,
-            depth_model_3d,
-            depth_download_files,
-            gs_viewer_html,
-            gs_download_files,
-        ],
+        outputs=result_outputs,
     )