Sync Hugging Face runtime tests

tests/test_hf_runtime.py

@@ -7,15 +7,19 @@ import torch
 from src.demo import hf_runtime, hf_ui
 from src.demo.hf_runtime import (
     SPLAT_TRANSFORM_PACKAGE,
-    ExportedArtifacts,
+    BrowserViewerArtifacts,
     GaussianArtifact,
+    ViewerTemplate,
     build_standalone_viewer,
     build_splat_transform_command,
+    build_viewer_template_command,
+    export_filtered_gaussian_ply,
     install_shared_viewer_assets,
     load_gaussian_artifact,
+    prepare_viewer_template,
     save_gaussian_artifact,
 )
-from src.demo.hf_ui import build_viewer_iframe
+from src.demo.hf_ui import build_viewer_iframe, build_viewer_preloader
 from src.demo.infer_single_image import patch_supersplat_html_viewer_bridge
 from src.utils.gaussians import Gaussians3D
 
@@ -48,24 +52,92 @@ def test_gaussian_artifact_round_trip(tmp_path) -> None:
         assert torch.equal(expected_tensor, actual_tensor)
 
 
+def test_filtered_ply_export_removes_spatial_outlier(tmp_path, monkeypatch) -> None:
+    points = torch.cat(
+        [
+            torch.randn(100, 3) * 0.01,
+            torch.tensor([[10.0, 10.0, 10.0]]),
+        ],
+        dim=0,
+    ).unsqueeze(0)
+    count = points.shape[1]
+    gaussians = Gaussians3D(
+        mean_vectors=points,
+        singular_values=torch.ones(1, count, 3),
+        quaternions=torch.ones(1, count, 4),
+        colors=torch.ones(1, count, 3),
+        opacities=torch.ones(1, count),
+    )
+    artifact_path = tmp_path / "gaussians.pt"
+    save_gaussian_artifact(
+        GaussianArtifact(
+            gaussians=gaussians,
+            focal_length_px=1200.0,
+            image_shape=(1152, 1536),
+        ),
+        artifact_path,
+    )
+    captured = {}
+
+    def fake_save_ply(*, gaussians, path, **_) -> None:
+        captured["gaussians"] = gaussians
+        path.write_bytes(b"ply")
+
+    monkeypatch.setattr(hf_runtime, "save_ply", fake_save_ply)
+
+    result = export_filtered_gaussian_ply(artifact_path, tmp_path / "output")
+
+    assert result.read_bytes() == b"ply"
+    assert captured["gaussians"].mean_vectors.shape[1] == count - 1
+    assert float(captured["gaussians"].mean_vectors.abs().max()) < 1.0
+
+
 def test_build_splat_transform_command_is_fixed_and_non_interactive() -> None:
     command = build_splat_transform_command(
         scene_ply=Path("scene.ply"),
-        output_html=Path("scene.html"),
+        output_sog=Path("viewer.sog"),
+    )
+
+    assert command[:4] == [
+        "npx",
+        "--yes",
+        "--prefer-offline",
+        SPLAT_TRANSFORM_PACKAGE,
+    ]
+    assert SPLAT_TRANSFORM_PACKAGE == "@playcanvas/splat-transform@3.1.6"
+    assert command[4:] == [
+        "--quiet",
+        "--overwrite",
+        "scene.ply",
+        "--filter-harmonics",
+        "0",
+        "viewer.sog",
+    ]
+
+
+def test_build_viewer_template_command_is_unbundled() -> None:
+    command = build_viewer_template_command(
+        scene_ply=Path("template.ply"),
+        output_html=Path("template.html"),
         viewer_settings=Path("viewer.json"),
     )
 
-    assert command[:3] == ["npx", "--yes", SPLAT_TRANSFORM_PACKAGE]
-    assert command[3:] == [
+    assert command[:4] == [
+        "npx",
+        "--yes",
+        "--prefer-offline",
+        SPLAT_TRANSFORM_PACKAGE,
+    ]
+    assert command[4:] == [
         "--quiet",
         "--overwrite",
         "--unbundled",
         "--viewer-settings",
         "viewer.json",
-        "scene.ply",
+        "template.ply",
         "--filter-harmonics",
         "0",
-        "scene.html",
+        "template.html",
     ]
 
 
@@ -81,6 +153,15 @@ def test_build_viewer_iframe_uses_gradio_file_route(tmp_path) -> None:
     assert "viewer%20output.html" in iframe
 
 
+def test_build_viewer_preloader_warms_template_in_hidden_iframe(tmp_path) -> None:
+    preload = build_viewer_preloader(tmp_path / "template.html")
+
+    assert "Ready" in preload
+    assert 'class="viewer-preloader"' in preload
+    assert "/gradio_api/file=" in preload
+    assert "data-infinisplat-viewer" not in preload
+
+
 def test_supersplat_viewer_bridge_reports_first_frame(tmp_path) -> None:
     viewer_path = tmp_path / "viewer.html"
     viewer_path.write_text(
@@ -130,23 +211,31 @@ def test_select_example_returns_gallery_path() -> None:
 
 
 def test_viewer_statuses_are_explicit() -> None:
+    assert hf_ui.GPU_DURATION_SECONDS == 6
     assert "Reconstructing scene" in hf_ui.show_reconstructing_viewer()
-    assert "Preparing viewer" in hf_ui.show_exporting_viewer()
+    assert "Filtering floaters" in hf_ui.show_exporting_viewer()
+    assert "PLY ready" in hf_ui.show_ply_ready_viewer()
     assert "Reconstruction stopped" in hf_ui.show_failed_viewer()
+    assert hf_ui.show_failed_html_download()["label"] == "HTML export failed"
 
 
-def _write_unbundled_viewer(output_dir: Path) -> Path:
-    viewer_html = output_dir / "viewer.html"
+def _write_unbundled_viewer(
+    output_dir: Path,
+    viewer_name: str = "viewer.html",
+) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    viewer_html = output_dir / viewer_name
+    scene_sog = viewer_html.with_suffix(".sog")
     viewer_html.write_text(
         '<html><head><link rel="stylesheet" href="./index.css"></head>'
         '<body><script type="module">'
         "const settingsUrl = './settings.json';"
         'const settings = {settings: fetch(settingsUrl).then(response => response.json())};'
-        'const scene = fetch("viewer.sog");'
+        f'const scene = fetch("{scene_sog.name}");'
         "import { main } from './index.js';"
         "</script></body></html>"
     )
-    (output_dir / "viewer.sog").write_bytes(b"scene")
+    scene_sog.write_bytes(b"scene")
     (output_dir / "index.css").write_text("canvas { display: block; }")
     (output_dir / "index.js").write_text("const main = true;")
     (output_dir / "settings.json").write_text('{"camera":{"fov":60}}')
@@ -154,10 +243,12 @@ def _write_unbundled_viewer(output_dir: Path) -> Path:
 
 
 def test_build_standalone_viewer_embeds_unbundled_assets(tmp_path) -> None:
-    viewer_html = _write_unbundled_viewer(tmp_path)
-    standalone_html = tmp_path / "scene.html"
+    request_dir = tmp_path / "request"
+    viewer_html = _write_unbundled_viewer(request_dir)
+    shared_dir = install_shared_viewer_assets(viewer_html)
+    standalone_html = request_dir / "scene.html"
 
-    build_standalone_viewer(viewer_html, standalone_html)
+    build_standalone_viewer(viewer_html, shared_dir, standalone_html)
     source = standalone_html.read_text()
 
     assert "canvas { display: block; }" in source
@@ -166,6 +257,37 @@ def test_build_standalone_viewer_embeds_unbundled_assets(tmp_path) -> None:
     assert "data:application/octet-stream;base64,c2NlbmU=" in source
 
 
+def test_prepare_viewer_template_builds_assets_once(tmp_path, monkeypatch) -> None:
+    captured = {}
+
+    def fake_save_ply(*, path, **_) -> None:
+        path.write_bytes(b"ply")
+
+    def fake_run(command, **_) -> None:
+        captured["command"] = command
+        output_html = Path(command[-1])
+        _write_unbundled_viewer(output_html.parent, output_html.name)
+
+    monkeypatch.setattr(hf_runtime, "save_ply", fake_save_ply)
+    monkeypatch.setattr(hf_runtime.subprocess, "run", fake_run)
+    monkeypatch.setattr(hf_runtime, "patch_supersplat_html_auto_rotate", lambda _: None)
+    monkeypatch.setattr(
+        hf_runtime,
+        "patch_supersplat_html_viewer_bridge",
+        lambda *_, **__: None,
+    )
+
+    template = prepare_viewer_template(
+        output_root=tmp_path,
+        command_prefix=["splat-transform"],
+    )
+
+    assert "--unbundled" in captured["command"]
+    assert template.viewer_html.name == "template.html"
+    assert template.viewer_assets_dir.parent == tmp_path / "_viewer_assets"
+    assert "../_viewer_assets/" in template.viewer_html.read_text()
+
+
 def test_install_shared_viewer_assets_uses_content_hash(tmp_path) -> None:
     request_dir = tmp_path / "request"
     request_dir.mkdir()
@@ -180,7 +302,7 @@ def test_install_shared_viewer_assets_uses_content_hash(tmp_path) -> None:
     assert f"../_viewer_assets/{shared_dir.name}/index.js" in source
 
 
-def test_export_gaussian_artifact_preserves_all_gaussians(
+def test_export_gaussian_artifact_builds_filtered_downloads(
     tmp_path, monkeypatch
 ) -> None:
     artifact_path = tmp_path / "gaussians.pt"
@@ -194,6 +316,14 @@ def test_export_gaussian_artifact_preserves_all_gaussians(
         artifact_path,
     )
     captured = {}
+    template_html = _write_unbundled_viewer(
+        tmp_path / "_viewer_template",
+        "template.html",
+    )
+    viewer_template = ViewerTemplate(
+        viewer_html=template_html,
+        viewer_assets_dir=install_shared_viewer_assets(template_html),
+    )
 
     def fake_save_ply(*, gaussians, path, **_) -> None:
         captured["gaussians"] = gaussians
@@ -201,20 +331,14 @@ def test_export_gaussian_artifact_preserves_all_gaussians(
 
     def fake_run(command, **_) -> None:
         captured["command"] = command
-        _write_unbundled_viewer(Path(command[-1]).parent)
+        Path(command[-1]).write_bytes(b"scene")
 
     monkeypatch.setattr(hf_runtime, "save_ply", fake_save_ply)
     monkeypatch.setattr(hf_runtime.subprocess, "run", fake_run)
-    monkeypatch.setattr(hf_runtime, "patch_supersplat_html_auto_rotate", lambda _: None)
-    monkeypatch.setattr(
-        hf_runtime,
-        "patch_supersplat_html_viewer_bridge",
-        lambda *_, **__: None,
-    )
-
     exported = hf_runtime.export_gaussian_artifact(
         artifact_path=artifact_path,
         output_dir=tmp_path / "output",
+        viewer_template=viewer_template,
         command_prefix=["splat-transform"],
     )
 
@@ -222,35 +346,127 @@ def test_export_gaussian_artifact_preserves_all_gaussians(
     assert actual.mean_vectors.shape == expected.mean_vectors.shape
     for expected_tensor, actual_tensor in zip(expected, actual):
         assert torch.equal(expected_tensor, actual_tensor)
-    assert "--unbundled" in captured["command"]
+    assert "--unbundled" not in captured["command"]
+    assert captured["command"][-1] == str(exported.scene_sog)
+    assert exported.scene_sog.read_bytes() == b"scene"
     assert exported.viewer_html.is_file()
+    assert 'fetch("viewer.sog")' in exported.viewer_html.read_text()
+    assert not (exported.viewer_html.parent / "index.js").exists()
     assert exported.standalone_html.is_file()
+    assert "data:application/octet-stream;base64,c2NlbmU=" in exported.standalone_html.read_text()
 
 
-def test_export_results_removes_only_internal_artifact(tmp_path, monkeypatch) -> None:
+def test_export_ply_result_exposes_ply_before_viewer(tmp_path, monkeypatch) -> None:
     internal_artifact = tmp_path / "gaussians.pt"
     scene_ply = tmp_path / "scene.ply"
-    viewer_html = tmp_path / "viewer.html"
-    standalone_html = tmp_path / "scene.html"
-    for path in (internal_artifact, scene_ply, viewer_html, standalone_html):
+    for path in (internal_artifact, scene_ply):
         path.write_bytes(b"test")
+    monkeypatch.setattr(
+        hf_ui,
+        "export_filtered_gaussian_ply",
+        lambda **_: scene_ply,
+    )
+
+    result = hf_ui.export_ply_result(str(internal_artifact))
+
+    assert not internal_artifact.exists()
+    assert scene_ply.exists()
+    assert "PLY ready" in result[0]
+    assert result[1]["label"] == "Download PLY - Ready"
+    assert result[1]["value"] == str(scene_ply)
+    assert result[1]["interactive"] is True
+    assert result[1]["elem_classes"] == ["artifact-button", "artifact-ready"]
+    assert result[2] == str(scene_ply)
 
+
+def test_export_viewer_result_starts_iframe_before_html(tmp_path, monkeypatch) -> None:
+    scene_ply = tmp_path / "scene.ply"
+    scene_sog = tmp_path / "viewer.sog"
+    viewer_html = tmp_path / "viewer.html"
+    scene_sog.write_bytes(b"sog")
+    viewer_html.write_text("<html></html>")
+    template = ViewerTemplate(
+        viewer_html=tmp_path / "template.html",
+        viewer_assets_dir=tmp_path / "assets",
+    )
+    monkeypatch.setattr(hf_ui, "_viewer_template", template)
     monkeypatch.setattr(
         hf_ui,
-        "export_gaussian_artifact",
-        lambda **_: ExportedArtifacts(
-            scene_ply=scene_ply,
+        "export_browser_viewer",
+        lambda **_: BrowserViewerArtifacts(
+            scene_sog=scene_sog,
             viewer_html=viewer_html,
-            standalone_html=standalone_html,
         ),
     )
 
-    hf_ui.export_results(str(internal_artifact))
+    result = hf_ui.export_viewer_result(str(scene_ply))
 
-    assert not internal_artifact.exists()
-    assert scene_ply.exists()
-    assert viewer_html.exists()
-    assert standalone_html.exists()
+    assert "Initializing 3D viewer" in result[0]
+    assert result[1] == str(viewer_html)
+
+
+def test_export_html_result_returns_native_download(tmp_path, monkeypatch) -> None:
+    viewer_html = tmp_path / "viewer.html"
+    standalone_html = tmp_path / "scene.html"
+    standalone_html.write_text("<html></html>")
+    template = ViewerTemplate(
+        viewer_html=tmp_path / "template.html",
+        viewer_assets_dir=tmp_path / "assets",
+    )
+    monkeypatch.setattr(hf_ui, "_viewer_template", template)
+    monkeypatch.setattr(
+        hf_ui,
+        "export_standalone_viewer",
+        lambda **_: standalone_html,
+    )
+
+    result = hf_ui.export_html_result(str(viewer_html))
+
+    assert result["label"] == "Download HTML viewer - Ready"
+    assert result["value"] == str(standalone_html)
+    assert result["interactive"] is True
+    assert result["elem_classes"] == ["artifact-button", "artifact-ready"]
+
+
+def test_begin_reconstruction_resets_download_states() -> None:
+    result = hf_ui.begin_reconstruction()
+
+    assert "Reconstructing scene" in result[0]
+    assert result[1]["label"] == "Preparing PLY..."
+    assert result[1]["value"] is None
+    assert result[1]["interactive"] is False
+    assert result[2] is None
+    assert result[3]["label"] == "Preparing HTML viewer..."
+    assert result[3]["value"] is None
+    assert result[3]["interactive"] is False
+    assert result[4] is None
+
+
+def test_demo_uses_native_staged_downloads(tmp_path, monkeypatch) -> None:
+    template = ViewerTemplate(
+        viewer_html=tmp_path / "template.html",
+        viewer_assets_dir=tmp_path / "assets",
+    )
+    monkeypatch.setattr(hf_ui, "prepare_viewer_template", lambda _: template)
+    monkeypatch.setattr(hf_ui, "configure_runtime", lambda *_: None)
+
+    demo = hf_ui.create_demo(object())
+    config = demo.get_config_file()
+    download_labels = {
+        component["props"].get("label"): component
+        for component in config["components"]
+        if component["type"] == "downloadbutton"
+    }
+    function_names = [
+        block.fn.__name__
+        for block in demo.fns.values()
+        if block.fn is not None
+    ]
+
+    assert "Download PLY" in download_labels
+    assert "Download HTML viewer" in download_labels
+    assert function_names.index("export_ply_result") < function_names.index("export_viewer_result")
+    assert function_names.index("export_viewer_result") < function_names.index("export_html_result")
 
 
 def test_app_imports_spaces_before_model_modules() -> None: