tests/test_hf_runtime.py

import ast
import os
from pathlib import Path

import pytest
import torch
from PIL import Image

from src.demo import hf_runtime, hf_ui
from src.demo.hf_runtime import (
    LOCAL_RGB_CHECKPOINT,
    SPLAT_TRANSFORM_PACKAGE,
    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,
    resolve_rgb_checkpoint,
    save_gaussian_artifact,
)
from src.demo.hf_ui import (
    build_viewer_iframe,
    build_viewer_preloader,
    cleanup_request_directories,
)
from src.demo.infer_single_image import patch_supersplat_html_viewer_bridge
from src.utils.gaussians import Gaussians3D


def _gaussians() -> Gaussians3D:
    return Gaussians3D(
        mean_vectors=torch.tensor([[[1.0, 2.0, 3.0]]]),
        singular_values=torch.tensor([[[0.1, 0.2, 0.3]]]),
        quaternions=torch.tensor([[[1.0, 0.0, 0.0, 0.0]]]),
        colors=torch.tensor([[[0.4, 0.5, 0.6]]]),
        opacities=torch.tensor([[0.9]]),
        covariances=torch.eye(3).reshape(1, 1, 3, 3),
    )


def test_checkpoint_override_has_highest_priority(tmp_path, monkeypatch) -> None:
    override = tmp_path / "override.ckpt"
    override.touch()
    monkeypatch.setenv("INFINISPLAT_CHECKPOINT", str(override))

    assert resolve_rgb_checkpoint() == override


def test_repository_checkpoint_is_reused_before_hub(tmp_path, monkeypatch) -> None:
    checkpoint = tmp_path / "infinisplat_rgb.ckpt"
    checkpoint.touch()
    monkeypatch.delenv("INFINISPLAT_CHECKPOINT", raising=False)
    monkeypatch.setattr(hf_runtime, "LOCAL_RGB_CHECKPOINT", checkpoint)
    monkeypatch.setattr(
        hf_runtime,
        "hf_hub_download",
        lambda **_: pytest.fail("Hub download should not be called"),
    )

    assert resolve_rgb_checkpoint() == checkpoint


def test_checkpoint_falls_back_to_hub(tmp_path, monkeypatch) -> None:
    downloaded = tmp_path / "downloaded.ckpt"
    downloaded.touch()
    monkeypatch.delenv("INFINISPLAT_CHECKPOINT", raising=False)
    monkeypatch.setattr(hf_runtime, "LOCAL_RGB_CHECKPOINT", tmp_path / "missing.ckpt")
    monkeypatch.setattr(hf_runtime, "hf_hub_download", lambda **_: str(downloaded))

    assert resolve_rgb_checkpoint() == downloaded


def test_local_checkpoint_path_matches_repository_layout() -> None:
    assert LOCAL_RGB_CHECKPOINT == Path("checkpoints/infinisplat_rgb.ckpt").resolve()


def test_request_cleanup_removes_only_expired_uuid_directories(tmp_path) -> None:
    old_request = tmp_path / ("a" * 32)
    recent_request = tmp_path / ("b" * 32)
    shared_assets = tmp_path / "_viewer_assets"
    old_request.mkdir()
    recent_request.mkdir()
    shared_assets.mkdir()
    (old_request / "viewer.sog").write_bytes(b"old")
    (recent_request / "viewer.sog").write_bytes(b"recent")
    (shared_assets / "index.js").write_text("shared")
    os.utime(old_request, (100.0, 100.0))
    os.utime(recent_request, (900.0, 900.0))

    removed = cleanup_request_directories(
        output_root=tmp_path,
        max_age_seconds=500,
        now=1000.0,
    )

    assert removed == [old_request]
    assert not old_request.exists()
    assert recent_request.is_dir()
    assert shared_assets.is_dir()


def test_gaussian_artifact_round_trip(tmp_path) -> None:
    path = tmp_path / "artifact.pt"
    expected = GaussianArtifact(
        gaussians=_gaussians(),
        focal_length_px=1200.0,
        image_shape=(1152, 1536),
    )

    save_gaussian_artifact(expected, path)
    actual = load_gaussian_artifact(path)

    assert actual.focal_length_px == expected.focal_length_px
    assert actual.image_shape == expected.image_shape
    for expected_tensor, actual_tensor in zip(expected.gaussians, actual.gaussians):
        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_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[:4] == [
        "npx",
        "--yes",
        "--prefer-offline",
        SPLAT_TRANSFORM_PACKAGE,
    ]
    assert command[4:] == [
        "--quiet",
        "--overwrite",
        "--unbundled",
        "--viewer-settings",
        "viewer.json",
        "template.ply",
        "--filter-harmonics",
        "0",
        "template.html",
    ]


def test_build_viewer_iframe_uses_gradio_file_route(tmp_path) -> None:
    viewer_path = tmp_path / "viewer output.html"

    iframe = build_viewer_iframe(viewer_path)

    assert '<iframe class="splat-frame"' in iframe
    assert 'data-viewer-state="loading"' in iframe
    assert "Initializing 3D viewer" in iframe
    assert "/gradio_api/file=" in iframe
    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(
        "<html><head></head><body>window.firstFrame?.();</body></html>"
    )

    patch_supersplat_html_viewer_bridge(viewer_path)
    patched_once = viewer_path.read_text()
    patch_supersplat_html_viewer_bridge(viewer_path)

    assert "data-infinisplat-viewer-bridge" in patched_once
    assert 'window.firstFrame = () => setViewerState("ready")' in patched_once
    assert "host.dataset.viewerState = state" in patched_once
    assert viewer_path.read_text() == patched_once


def test_supersplat_viewer_bridge_accepts_external_script(tmp_path) -> None:
    viewer_path = tmp_path / "viewer.html"
    script_path = tmp_path / "index.js"
    viewer_path.write_text("<html><head></head><body></body></html>")
    script_path.write_text("window.firstFrame?.();")

    patch_supersplat_html_viewer_bridge(viewer_path, script_path)

    assert "data-infinisplat-viewer-bridge" in viewer_path.read_text()


def test_supersplat_viewer_bridge_rejects_changed_bundle(tmp_path) -> None:
    viewer_path = tmp_path / "viewer.html"
    viewer_path.write_text("<html><head></head><body></body></html>")

    with pytest.raises(RuntimeError, match="viewer bundle changed"):
        patch_supersplat_html_viewer_bridge(viewer_path)


def test_hf_space_examples_exist() -> None:
    discovered = {
        path.relative_to(hf_ui.REPO_ROOT).as_posix()
        for path in hf_ui.RGB_EXAMPLE_DIR.iterdir()
        if path.is_file()
        and path.suffix.lower() in hf_ui.SUPPORTED_EXAMPLE_SUFFIXES
    }

    assert set(hf_ui.RGB_EXAMPLES) == discovered
    assert len(hf_ui.RGB_EXAMPLES) == len(discovered)
    assert hf_ui.RGB_EXAMPLES[0] == "examples/data/rgb_demo/painting_room.jpg"
    assert hf_ui.RGB_EXAMPLES[1] == "examples/data/rgb_demo/summer_room.jpg"
    assert hf_ui.RGB_EXAMPLES[3] == "examples/data/rgb_demo/meerkat.jpg"
    assert hf_ui.RGB_EXAMPLES[8] == "examples/data/rgb_demo/eth3d_courtyard.png"
    assert hf_ui.RGB_EXAMPLES[14] == "examples/data/rgb_demo/pexels-masi.jpg"
    for path in hf_ui.RGB_EXAMPLES:
        assert Path(path).is_file()


def test_hf_space_uses_node_22_for_splat_transform() -> None:
    requirements = Path("requirements.txt").read_text().splitlines()
    packages = Path("packages.txt").read_text().splitlines()

    assert "nodejs-wheel==22.20.0" in requirements
    assert "nodejs" not in packages
    assert "npm" not in packages


def test_select_example_returns_gallery_path() -> None:
    event = type("SelectEvent", (), {"index": 2})()

    assert hf_ui.select_example(event) == hf_ui.RGB_EXAMPLES[2]


def test_example_thumbnails_are_smaller_than_gallery_sources(tmp_path) -> None:
    source = tmp_path / "source.jpg"
    Image.new("RGB", (1200, 900), "white").save(source)

    thumbnails = hf_ui.prepare_example_thumbnails(
        example_paths=[str(source)],
        output_dir=tmp_path / "thumbnails",
    )

    with Image.open(thumbnails[0]) as thumbnail:
        assert thumbnail.size == (427, 320)
        assert thumbnail.format == "WEBP"


def test_viewer_statuses_are_explicit() -> None:
    assert hf_ui.GPU_DURATION_SECONDS == 6
    assert "Reconstructing scene" in hf_ui.show_reconstructing_viewer()
    exporting_status = hf_ui.show_exporting_viewer()
    assert "Preparing scene" in exporting_status
    assert "float" not in exporting_status.lower()
    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 test_hf_space_header_content() -> None:
    assert hf_ui.FULL_TITLE == (
        "Implicit Gaussian Decoding for Large-Baseline Monocular View Synthesis"
    )
    assert hf_ui.GITHUB_URL == "https://github.com/PLUS-WAVE/InfiniSplat-oss"
    assert hf_ui.PROJECT_PAGE_URL == "https://pluswave.top/InfiniSplat-page/"
    assert hf_ui.INPUT_IMAGE_HINT == (
        "Better for indoor scenes due to HyperSim-only training"
    )
    assert ".panel-heading {\n    display: flex;\n    align-items: center;" in (
        hf_ui.APP_CSS
    )
    assert "box-sizing: border-box;" in hf_ui.APP_CSS
    assert "padding: 17px 12px 15px;" in hf_ui.APP_CSS
    assert ".panel-heading .input-hint" in hf_ui.APP_CSS
    assert "white-space: normal;" in hf_ui.APP_CSS
    assert "font-size: clamp(0.58rem, 0.72vw, 0.66rem);" in hf_ui.APP_CSS


def test_hf_space_gallery_has_no_captions_or_thumbnail_frames() -> None:
    assert all(isinstance(path, str) for path in hf_ui.RGB_EXAMPLES)
    assert "#example-gallery .thumbnail-item" in hf_ui.APP_CSS
    assert "#example-gallery .caption-label" in hf_ui.APP_CSS
    assert "grid-template-columns: repeat(28, minmax(0, 1fr))" in hf_ui.APP_CSS
    assert "#example-gallery .gallery-item:nth-child(10)" in hf_ui.APP_CSS
    assert "#example-gallery .gallery-item:nth-child(8)" in hf_ui.APP_CSS
    assert "grid-column: span 7" not in hf_ui.APP_CSS
    assert "grid-column: span 2" in hf_ui.APP_CSS
    assert "#example-gallery .gallery-item:last-child:nth-child(odd)" in (
        hf_ui.APP_CSS
    )


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())};'
        f'const scene = fetch("{scene_sog.name}");'
        "import { main } from './index.js';"
        "</script></body></html>"
    )
    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}}')
    return viewer_html


def test_build_standalone_viewer_embeds_unbundled_assets(tmp_path) -> None:
    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, shared_dir, standalone_html)
    source = standalone_html.read_text()

    assert "canvas { display: block; }" in source
    assert "const main = true;" in source
    assert 'settings: {"camera":{"fov":60}}' in source
    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()
    viewer_html = _write_unbundled_viewer(request_dir)

    shared_dir = install_shared_viewer_assets(viewer_html)
    source = viewer_html.read_text()

    assert shared_dir.parent == tmp_path / "_viewer_assets"
    assert len(shared_dir.name) == 16
    assert (shared_dir / "index.js").read_text() == "const main = true;"
    assert f"../_viewer_assets/{shared_dir.name}/index.js" in source


def test_export_gaussian_artifact_builds_filtered_downloads(
    tmp_path, monkeypatch
) -> None:
    artifact_path = tmp_path / "gaussians.pt"
    expected = _gaussians()
    save_gaussian_artifact(
        GaussianArtifact(
            gaussians=expected,
            focal_length_px=1200.0,
            image_shape=(1152, 1536),
        ),
        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
        path.write_bytes(b"ply")

    def fake_run(command, **_) -> None:
        captured["command"] = command
        Path(command[-1]).write_bytes(b"scene")

    monkeypatch.setattr(hf_runtime, "save_ply", fake_save_ply)
    monkeypatch.setattr(hf_runtime.subprocess, "run", fake_run)
    exported = hf_runtime.export_gaussian_artifact(
        artifact_path=artifact_path,
        output_dir=tmp_path / "output",
        viewer_template=viewer_template,
        command_prefix=["splat-transform"],
    )

    actual = captured["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" 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_ply_result_exposes_ply_before_viewer(tmp_path, monkeypatch) -> None:
    internal_artifact = tmp_path / "gaussians.pt"
    scene_ply = tmp_path / "scene.ply"
    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_browser_viewer",
        lambda **_: BrowserViewerArtifacts(
            scene_sog=scene_sog,
            viewer_html=viewer_html,
        ),
    )

    result = hf_ui.export_viewer_result(str(scene_ply))

    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:
    tree = ast.parse(Path("app.py").read_text())
    imported_modules = []
    for node in tree.body:
        if isinstance(node, ast.Import):
            imported_modules.extend(alias.name for alias in node.names)
        elif isinstance(node, ast.ImportFrom) and node.module:
            imported_modules.append(node.module)

    assert imported_modules.index("spaces") < imported_modules.index("src.demo.hf_runtime")