scripts/test_rtx_upscale_integration.py

"""Standalone sanity checks for the hidden half-res output restore.

Runs without ComfyUI or torch: minimal stubs exercise the gating logic of
utils/rtx_vsr.py, the compose-at-native-size logic in utils/h3_project.py,
and the restore wiring in the combine nodes.
"""

from __future__ import annotations

import ast
import importlib
import py_compile
import sys
import types
from pathlib import Path

PACK = Path(__file__).resolve().parents[1]

failures: list[str] = []


def check(name: str, condition: bool) -> None:
    status = "PASS" if condition else "FAIL"
    print(f"[{status}] {name}")
    if not condition:
        failures.append(name)


# --- stub module surface ----------------------------------------------------
class FakeTensor:
    def __init__(self, shape):
        self.shape = tuple(shape)
        self.ndim = len(self.shape)
        self.device = "cpu"
        self.dtype = "float32"

    def __getitem__(self, item):
        return FakeTensor(self.shape)

    def movedim(self, *args, **kwargs):
        return FakeTensor(self.shape)


torch_stub = types.ModuleType("torch")
torch_stub.Tensor = FakeTensor
torch_stub.empty = lambda *a, **k: FakeTensor(a[0])
sys.modules["torch"] = torch_stub

upscale_calls: list[tuple] = []
comfy_stub = types.ModuleType("comfy")
comfy_utils_stub = types.ModuleType("comfy.utils")


def _fake_common_upscale(samples, width, height, method, crop):
    upscale_calls.append((width, height, method, crop))
    return samples


comfy_utils_stub.common_upscale = _fake_common_upscale
comfy_stub.utils = comfy_utils_stub
sys.modules["comfy"] = comfy_stub
sys.modules["comfy.utils"] = comfy_utils_stub

# comfy_api.latest stub (VideoFromComponents / Types.VideoComponents)
comfy_api_pkg = types.ModuleType("comfy_api")
comfy_api_latest = types.ModuleType("comfy_api.latest")


class FakeVideoComponents:
    def __init__(self, images=None, audio=None, frame_rate=None):
        self.images = images
        self.audio = audio
        self.frame_rate = frame_rate


class _Types:
    VideoComponents = FakeVideoComponents


built_videos: list = []


class _InputImpl:
    @classmethod
    def VideoFromComponents(cls, components):
        video = ("components-video", components)
        built_videos.append(video)
        return video


comfy_api_latest.InputImpl = _InputImpl
comfy_api_latest.Types = _Types
comfy_api_pkg.latest = comfy_api_latest
sys.modules["comfy_api"] = comfy_api_pkg
sys.modules["comfy_api.latest"] = comfy_api_latest

# utils.video stub (ffprobe only; merge is never called here)
video_stub = types.ModuleType("utils.video")
PROBED: dict = {}
PROBE_RAISE = {"raise": False}


def _fake_ffprobe_info(path):
    if PROBE_RAISE["raise"]:
        raise RuntimeError("probe failed")
    return dict(PROBED.get(path, {"width": 1920, "height": 1080}))


video_stub.ffprobe_info = _fake_ffprobe_info
sys.modules["utils.video"] = video_stub

MANIFEST: dict = {"width": 0, "height": 0}
h3_stub = types.ModuleType("utils.h3_project")
h3_stub.safe_h3_project_name = lambda value: str(value or "default")


def _fake_load(project_name):
    if MANIFEST.get("missing"):
        raise FileNotFoundError("no project")
    return dict(MANIFEST)


h3_stub.load_h3_project_data = _fake_load

utils_pkg = types.ModuleType("utils")
utils_pkg.__path__ = [str(PACK / "utils")]
sys.modules["utils"] = utils_pkg
sys.modules["utils.h3_project"] = h3_stub

rtx_vsr = importlib.import_module("utils.rtx_vsr")

# --- rtx_vsr.upscale_images_to_project_resolution ---------------------------
# 1. half-res frames are restored to the exact manifest size (lanczos here,
#    since nvidia-vfx is absent in the sandbox)
MANIFEST.update({"width": 1920, "height": 1080})
upscale_calls.clear()
rtx_vsr.upscale_images_to_project_resolution(FakeTensor((4, 528, 960, 3)), "proj")
check(
    "half-res frames restored to manifest size",
    upscale_calls == [(1920, 1080, "lanczos", "disabled")],
)

# 2. frames already at the project resolution pass through untouched
upscale_calls.clear()
rtx_vsr.upscale_images_to_project_resolution(FakeTensor((4, 1080, 1920, 3)), "proj")
check("full-res frames pass through", upscale_calls == [])

# 3. missing manifest passes through
MANIFEST.update({"missing": True})
upscale_calls.clear()
rtx_vsr.upscale_images_to_project_resolution(FakeTensor((4, 528, 960, 3)), "proj")
check("missing manifest passes through", upscale_calls == [])
MANIFEST.pop("missing")

# 4. zero and audio-only sentinel manifests pass through
for width, height in ((0, 0), (32, 32)):
    MANIFEST.update({"width": width, "height": height})
    upscale_calls.clear()
    rtx_vsr.upscale_images_to_project_resolution(FakeTensor((4, 528, 960, 3)), "proj")
    check(f"manifest {width}x{height} passes through", upscale_calls == [])
MANIFEST.update({"width": 1920, "height": 1080})

# 5. non-tensor and empty-batch inputs pass through
check(
    "non-tensor passthrough",
    rtx_vsr.upscale_images_to_project_resolution("nope", "proj") == "nope",
)
check(
    "empty batch passthrough",
    rtx_vsr.upscale_images_to_project_resolution(
        FakeTensor((0, 528, 960, 3)), "proj"
    ).shape[0] == 0,
)

# 6. rtx_vsr_upscale identity early-returns (no upscale machinery touched)
same = FakeTensor((2, 1080, 1920, 3))
check("vsr same-size identity", rtx_vsr.rtx_vsr_upscale(same, 1920, 1080) is same)
zero = FakeTensor((0, 528, 960, 3))
check("vsr empty identity", rtx_vsr.rtx_vsr_upscale(zero, 1920, 1080) is zero)

# --- rtx_vsr.restore_project_video_resolution -------------------------------
class FakeVideo:
    def __init__(self, components=None, error=None):
        self._components = components
        self._error = error
        self.decode_calls = 0

    def get_components(self):
        self.decode_calls += 1
        if self._error is not None:
            raise self._error
        return self._components


# 7. composed file already at manifest size -> untouched, zero decode
upscale_calls.clear()
full_video = FakeVideo()
result = rtx_vsr.restore_project_video_resolution(
    full_video, "proj", source_path="composed.mp4"
)
check(
    "full-res composed file untouched without decode",
    result is full_video
    and full_video.decode_calls == 0
    and upscale_calls == [],
)

# 8. half-res composed file -> decode once, upscale, rebuild the video
upscale_calls.clear()
built_videos.clear()
half_images = FakeTensor((4, 528, 960, 3))
half_video = FakeVideo(FakeVideoComponents(images=half_images, audio="a", frame_rate=24))
PROBED["composed.mp4"] = {"width": 960, "height": 528}
result = rtx_vsr.restore_project_video_resolution(
    half_video, "proj", source_path="composed.mp4"
)
check(
    "half-res composed video rebuilt at manifest size",
    result == built_videos[-1]
    and half_video.decode_calls == 1
    and result[1].images is not half_images
    and result[1].audio == "a"
    and result[1].frame_rate == 24
    and upscale_calls == [(1920, 1080, "lanczos", "disabled")],
)

# 9. failed decode keeps the original video (rendering never breaks)
upscale_calls.clear()
broken_video = FakeVideo(error=RuntimeError("decode failed"))
PROBED["composed.mp4"] = {"width": 960, "height": 528}
result = rtx_vsr.restore_project_video_resolution(
    broken_video, "proj", source_path="composed.mp4"
)
check("failed decode keeps original video", result is broken_video)

# 10. missing manifest keeps the original video without decoding
MANIFEST.update({"missing": True})
quiet_video = FakeVideo()
result = rtx_vsr.restore_project_video_resolution(
    quiet_video, "proj", source_path="composed.mp4"
)
check(
    "missing manifest keeps video without decode",
    result is quiet_video and quiet_video.decode_calls == 0,
)
MANIFEST.pop("missing")

# 11. no source path -> decode path decides via the frame gate
upscale_calls.clear()
built_videos.clear()
probeless_video = FakeVideo(FakeVideoComponents(images=FakeTensor((4, 528, 960, 3))))
result = rtx_vsr.restore_project_video_resolution(probeless_video, "proj")
check(
    "no source path still restores",
    result == built_videos[-1]
    and upscale_calls == [(1920, 1080, "lanczos", "disabled")],
)
PROBED.clear()

# --- h3_project._uniform_probed_size (exec-extracted: module imports torch) --
h3_source = (PACK / "utils" / "h3_project.py").read_text(encoding="utf-8")
h3_tree = ast.parse(h3_source)
probe_fn = next(
    node
    for node in h3_tree.body
    if isinstance(node, ast.FunctionDef) and node.name == "_uniform_probed_size"
)
probe_ns: dict = {}
exec(  # noqa: S102 - trusted pack source, isolated namespace
    compile(ast.Module(body=[probe_fn], type_ignores=[]), "<probe>", "exec"),
    probe_ns,
)
uniform_probed_size = probe_ns["_uniform_probed_size"]

check(
    "uniform segment sizes compose at native size",
    uniform_probed_size(
        ["a.mp4", "b.mp4"], lambda p: {"width": 960, "height": 528}
    )
    == (960, 528),
)
check(
    "mixed segment sizes fall back",
    uniform_probed_size(
        ["a.mp4", "b.mp4"],
        lambda p: {"width": 1920, "height": 1080}
        if p == "a.mp4"
        else {"width": 960, "height": 528},
    )
    is None,
)
check(
    "odd segment sizes fall back",
    uniform_probed_size(["a.mp4"], lambda p: {"width": 961, "height": 528}) is None,
)
check(
    "unknowable segment sizes fall back",
    uniform_probed_size(["a.mp4"], lambda p: {}) is None,
)
check(
    "raising probe falls back",
    uniform_probed_size(["a.mp4"], lambda p: (_ for _ in ()).throw(RuntimeError("x")))
    is None,
)
check(
    "manifest-sized uniform timeline keeps manifest dims",
    uniform_probed_size(
        ["a.mp4", "b.mp4"], lambda p: {"width": 1920, "height": 1080}
    )
    == (1920, 1080),
)

# --- node wiring -------------------------------------------------------------
# 12. nodes/project.py: base combine restores before ANY output leaves the node
project_source = (PACK / "nodes" / "project.py").read_text(encoding="utf-8")
check(
    "single-pass half-res gate present",
    "editor_width // 2" in project_source
    and 'sampling_mode == "single"' in project_source,
)
check(
    "no 16px snap policy left behind",
    "half_up" not in project_source
    and "h3_halfres_sampling_dimensions" not in project_source,
)
check(
    "base combine restores before outputs",
    "restore_project_video_resolution(" in project_source
    and "source_path=str(composed_path)" in project_source
    and "InputImpl.VideoFromFile(str(composed_path))" in project_source,
)
check(
    "compose uses the uniform probed size",
    "_uniform_probed_size(" in h3_source
    and "compose_size[0] if compose_size else width" in h3_source,
)

# 13. moxie_combine relies on the base restore (no duplicated gate)
combine_source = (PACK / "moxie_combine.py").read_text(encoding="utf-8")
check(
    "moxie_combine has no duplicated upscale gate",
    "upscale_images_to_project_resolution" not in combine_source,
)
check(
    "moxie_combine exposes restored components",
    "components.images" in combine_source,
)

# 14. heavy imports stay lazy in rtx_vsr (pack import stays light)
tree = ast.parse((PACK / "utils" / "rtx_vsr.py").read_text(encoding="utf-8"))
module_level: list[str] = []
for node in tree.body:
    if isinstance(node, ast.Import):
        module_level.extend(alias.name for alias in node.names)
    elif isinstance(node, ast.ImportFrom) and node.level == 0:
        module_level.append(node.module or "")
check(
    "no torch/nvvfx/comfy imports at module level",
    not any(
        name.split(".")[0] in {"torch", "nvvfx", "comfy"} for name in module_level
    ),
)

# --- format combo relabel + quiet console -----------------------------------
# 15. the per-run upscale console line is gone (fallback diagnostics stay)
rtx_source = (PACK / "utils" / "rtx_vsr.py").read_text(encoding="utf-8")
check(
    "per-run upscale console line removed",
    "Output upscale:" not in rtx_source
    and "Lanczos fallback" in rtx_source,
)

# 16. format combo is the single internal value again; legacy label kept only
#     as execution-time tolerance for old queued workflows
basic_source = (PACK / "nodes" / "basic.py").read_text(encoding="utf-8")
check(
    "format combo shows MiniMax",
    basic_source.count('options=[FORMAT_INTERNAL_VALUE]') == 2
    and 'options=[FORMAT_DISPLAY_LABEL]' not in basic_source
    and 'FORMAT_INTERNAL_VALUE = "MiniMax"' in basic_source,
)
check(
    "legacy queued values still normalized at both node entries and lazy check",
    basic_source.count("\n        format = _normalize_format(format)") == 2
    and "\n        raw_format = _normalize_format(format)" in basic_source,
)

# 17. shipped templates carry the internal format value on the editor widget
import json as _json

for name in ("Moxie-consolidated.json", "Moxie.json", "Moxie-clean.json"):
    template = _json.loads((PACK / name).read_text(encoding="utf-8"))
    editor = next(
        node
        for node in template["nodes"]
        if node.get("type") == "easy multiTrackEditor"
    )
    check(
        f"{name} format widget is MiniMax",
        editor["widgets_values"][2] == "MiniMax"
        and editor["widgets_values_named"]["format"] == "MiniMax",
    )

# 18. frontend is pristine upstream: no display-label mapping, t2v native
frontend_utils = (
    PACK / "frontend" / "src" / "lib" / "multitrack-utils.ts"
).read_text(encoding="utf-8")
check(
    "frontend has no format mapping and resolves t2v natively",
    "normalizeFormatValue" not in frontend_utils
    and "FORMAT_DISPLAY_LABEL" not in frontend_utils
    and 'collected.format = format' in frontend_utils
    and "return 't2v'" in frontend_utils,
)

# --- MoxieCombine output cleanup (GetVideoComponents-style) ------------------
# 19. moxie_combine exposes only images/audio/fps; the base node keeps VIDEO+prefix
check(
    "moxie_combine outputs are images/audio/fps only",
    'io.Image.Output("images")' in combine_source
    and 'io.Audio.Output("audio")' in combine_source
    and 'io.Float.Output("fps")' in combine_source
    and "io.Video.Output" not in combine_source
    and "FILENAME_PREFIX" not in combine_source,
)
check(
    "moxie_combine returns three components (blocker path included)",
    "return io.NodeOutput(video, video, video)" in combine_source
    and "components.images" in combine_source
    and "float(components.frame_rate)" in combine_source,
)
check(
    "base combine keeps VIDEO and FILENAME_PREFIX",
    'io.Video.Output("VIDEO")' in project_source
    and 'io.String.Output("FILENAME_PREFIX")' in project_source,
)

# 20. the consolidated template is rewired to the slimmed node
consolidated = _json.loads((PACK / "Moxie-consolidated.json").read_text(encoding="utf-8"))
consolidated_combine = next(
    node for node in consolidated["nodes"] if node.get("type") == "MoxieCombine"
)
consolidated_save = next(
    node for node in consolidated["nodes"] if node.get("type") == "easy saveVideo"
)
check(
    "consolidated template MoxieCombine slots reindexed",
    [o["name"] for o in consolidated_combine["outputs"]] == ["images", "audio", "fps"],
)
consolidated_links = {link[0]: link for link in consolidated["links"]}
check(
    "consolidated template links follow the new slots",
    consolidated_links[105][2] == 2
    and consolidated_links[106][2] == 0
    and 108 not in consolidated_links
    and consolidated_links[107][1] == 102,
)
consolidated_prefix = next(
    port for port in consolidated_save["inputs"] if port["name"] == "filename_prefix"
)
expected_prefix = "easy_media/projects/girl1-test/out/girl1-test"
check(
    "consolidated template filename_prefix is a literal widget",
    consolidated_prefix["link"] is None
    and consolidated_save["widgets_values"][-1] == expected_prefix
    and consolidated_save["widgets_values_named"]["filename_prefix"] == expected_prefix,
)
check(
    "other templates do not reference MoxieCombine",
    "MoxieCombine" not in (PACK / "Moxie.json").read_text(encoding="utf-8")
    and "MoxieCombine" not in (PACK / "Moxie-clean.json").read_text(encoding="utf-8"),
)

# 21. everything touched compiles
for rel in (
    "nodes/project.py",
    "nodes/basic.py",
    "moxie_combine.py",
    "utils/rtx_vsr.py",
    "utils/h3_project.py",
):
    try:
        py_compile.compile(str(PACK / rel), doraise=True)
        check(f"py_compile {rel}", True)
    except py_compile.PyCompileError as error:
        check(f"py_compile {rel}", False)
        print(f"       {error}")

print()
if failures:
    print(f"{len(failures)} CHECK(S) FAILED")
    sys.exit(1)
print("ALL CHECKS PASSED")