olivv-cs-ComfyUI-FunPack/tests/test_detailing.py

"""Unit tests for segmented detailing (detailing.py).

Covers the pure pieces (target parsing, tube geometry + area cap, layout-cond
stripping, spatial downscale) and the detail_refine_scene contract: bit-identical
no-op when disabled / nothing detected, and the crop -> refine -> feathered
paste round trip with the model-touching stages stubbed out.
"""
import sys
import types
from pathlib import Path

import pytest
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

sys.modules.setdefault("comfy", types.ModuleType("comfy"))
_nested = types.ModuleType("comfy.nested_tensor")


class _FakeNested:
    is_nested = True

    def __init__(self, tensors):
        self._tensors = list(tensors)

    def unbind(self):
        return list(self._tensors)


_nested.NestedTensor = _FakeNested
sys.modules.setdefault("comfy.nested_tensor", _nested)


@pytest.fixture(autouse=True)
def _nested_stub(monkeypatch):
    """Re-install the usable NestedTensor stub around every test.

    Other test modules stub comfy.nested_tensor.NestedTensor = object on the SAME
    shared sys.modules entry at collection time (whichever file is collected last
    wins), and stubbed submodules never run real import machinery, so the parent
    attribute must be bound by hand too."""
    monkeypatch.setattr(sys.modules["comfy.nested_tensor"], "NestedTensor",
                        _FakeNested, raising=False)
    monkeypatch.setattr(sys.modules["comfy"], "nested_tensor",
                        sys.modules["comfy.nested_tensor"], raising=False)

import detailing  # noqa: E402


# ---------------------------------------------------------------------------
# parse_targets
# ---------------------------------------------------------------------------

def test_parse_targets_splits_and_strips():
    assert detailing.parse_targets(" hands , feet,,") == ["hands", "feet"]
    assert detailing.parse_targets("") == []
    assert detailing.parse_targets(None) == []


# ---------------------------------------------------------------------------
# stage2_sigmas
# ---------------------------------------------------------------------------

def test_stage2_sigmas_default_matches_official_recipe():
    # Must reproduce the exact official two-stage workflow values unless the
    # user raises/lowers detail_denoise — no behavior change for existing runs.
    assert torch.allclose(torch.tensor(detailing.stage2_sigmas(0.85)),
                          torch.tensor(detailing.STAGE2_SIGMAS))
    assert detailing.stage2_sigmas(0.85)[-1] == 0.0

def test_stage2_sigmas_scales_with_renoise_sigma():
    # Raising the re-entry point should raise every step in the tail proportionally
    # (more genuine freedom to reconstruct, not just a bigger first jump).
    low = detailing.stage2_sigmas(0.5)
    high = detailing.stage2_sigmas(0.95)
    assert high[0] > low[0] > 0
    assert high[1] > low[1]
    assert low[-1] == high[-1] == 0.0
    # Shape (ratios between steps) preserved regardless of entry point.
    assert abs(low[1] / low[0] - high[1] / high[0]) < 1e-9


# ---------------------------------------------------------------------------
# find_tube
# ---------------------------------------------------------------------------

def _heat_with_hot_square(h=64, w=64, y=(20, 30), x=(40, 50), value=0.9):
    heat = torch.zeros(h, w)
    heat[y[0]:y[1], x[0]:x[1]] = value
    return heat

def test_find_tube_boxes_hot_region_with_padding():
    heat = _heat_with_hot_square()
    tube = detailing.find_tube(heat, 64, 64, threshold=0.35)
    assert tube is not None
    y0, y1, x0, x1, mask = tube
    # Hot square plus TUBE_PAD on each side (interpolation is identity at same size).
    assert y0 <= 20 - detailing.TUBE_PAD + 1 and y1 >= 30 + detailing.TUBE_PAD - 1
    assert x0 <= 40 - detailing.TUBE_PAD + 1 and x1 >= 50 + detailing.TUBE_PAD - 1
    assert mask.shape == (1, 1, 1, y1 - y0, x1 - x0)
    assert float(mask.max()) <= 1.0 and float(mask.min()) >= 0.0
    # Mask peaks inside the hot region, feathers to ~0 at the padded border.
    assert float(mask[..., (y1 - y0) // 2, (x1 - x0) // 2]) > 0.5
    assert float(mask[..., 0, 0]) < 0.2

def test_find_tube_none_below_threshold():
    heat = torch.full((64, 64), 0.1)
    assert detailing.find_tube(heat, 64, 64, threshold=0.35) is None
    assert detailing.find_tube(None, 64, 64) is None

def test_find_tube_refuses_oversized_region():
    heat = torch.full((64, 64), 0.9)  # whole frame hot -> a re-render, not a detail pass
    assert detailing.find_tube(heat, 64, 64, threshold=0.35) is None

def test_find_tube_diag_reports_max_heat_and_reason():
    diag = {}
    heat = torch.full((64, 64), 0.1)
    assert detailing.find_tube(heat, 64, 64, threshold=0.35, diag=diag) is None
    assert diag["reason"] == "below_threshold"
    assert abs(diag["max_heat"] - 0.1) < 1e-6

    diag2 = {}
    heat2 = torch.full((64, 64), 0.9)  # whole frame -> area cap, not threshold
    assert detailing.find_tube(heat2, 64, 64, threshold=0.35, diag=diag2) is None
    assert diag2["reason"] == "area_cap"
    assert diag2["area"] > detailing.MAX_TUBE_AREA

def test_find_tube_enforces_minimum_edge():
    heat = _heat_with_hot_square(y=(31, 32), x=(31, 32))  # single hot cell
    tube = detailing.find_tube(heat, 64, 64, threshold=0.35)
    assert tube is not None
    y0, y1, x0, x1, _ = tube
    assert (y1 - y0) >= detailing.MIN_TUBE_EDGE
    assert (x1 - x0) >= detailing.MIN_TUBE_EDGE


# ---------------------------------------------------------------------------
# _clipseg_heat
# ---------------------------------------------------------------------------

def test_clipseg_heat_passes_padding_and_truncation(monkeypatch):
    """Regression: a multi-prompt call without padding=True crashes CLIPSeg's
    tokenizer ("excessive nesting... type list where int is expected") the moment
    detail_targets has more than one entry — caught live on a real rental run.
    truncation=True is the same call, cheap to require alongside it."""
    calls = []

    class _FakeInputs(dict):
        def to(self, *a, **k):
            return self

    class _FakeProcessor:
        def __call__(self, text, images, return_tensors, **kwargs):
            calls.append(kwargs)
            return _FakeInputs()

    class _FakeModel:
        def __call__(self, **inputs):
            n = 2
            return types.SimpleNamespace(logits=torch.zeros(n, 8, 8))

        def eval(self):
            return self

    monkeypatch.setattr(detailing, "_get_clipseg", lambda: (_FakeProcessor(), _FakeModel()))
    frames = [(0, torch.zeros(16, 16, 3))]
    detailing._clipseg_heat(frames, ["hands", "feet"])
    assert calls and calls[0].get("padding") is True
    assert calls[0].get("truncation") is True


# ---------------------------------------------------------------------------
# resolve_upsampler_name
# ---------------------------------------------------------------------------

def _patch_folder_paths(monkeypatch, files, folder="/models/latent_upscale_models"):
    fp = sys.modules["folder_paths"]
    monkeypatch.setattr(fp, "get_filename_list", lambda kind: list(files), raising=False)
    monkeypatch.setattr(fp, "get_folder_paths", lambda kind: [folder], raising=False)

def test_resolve_explicit_name_passes_through(monkeypatch):
    _patch_folder_paths(monkeypatch, [])
    assert detailing.resolve_upsampler_name("my_upscaler.safetensors") == "my_upscaler.safetensors"

def test_resolve_auto_prefers_newest_spatial_upscaler(monkeypatch):
    _patch_folder_paths(monkeypatch, [
        "other_model.safetensors",
        "ltx-2.3-spatial-upscaler-x2-1.0.safetensors",
        "ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
    ])
    for alias in ("auto", "None", "", None):
        assert detailing.resolve_upsampler_name(alias) == "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"

def test_resolve_auto_falls_back_to_any_installed(monkeypatch):
    _patch_folder_paths(monkeypatch, ["something_else.safetensors"])
    assert detailing.resolve_upsampler_name("auto") == "something_else.safetensors"

def test_resolve_auto_downloads_when_folder_empty(monkeypatch, tmp_path):
    _patch_folder_paths(monkeypatch, [], folder=str(tmp_path / "lup"))
    calls = {}
    hub = types.ModuleType("huggingface_hub")
    hub.hf_hub_download = lambda repo_id, filename, local_dir: calls.update(
        repo=repo_id, file=filename, dest=local_dir)
    monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
    name = detailing.resolve_upsampler_name("auto")
    assert name == detailing.DEFAULT_UPSAMPLER_FILE
    assert calls["repo"] == detailing.DEFAULT_UPSAMPLER_REPO
    assert calls["file"] == detailing.DEFAULT_UPSAMPLER_FILE

def test_resolve_download_failure_raises_actionable_error(monkeypatch, tmp_path):
    _patch_folder_paths(monkeypatch, [], folder=str(tmp_path / "lup"))
    hub = types.ModuleType("huggingface_hub")
    def _boom(**kw):
        raise OSError("401 gated")
    hub.hf_hub_download = _boom
    monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
    with pytest.raises(RuntimeError, match="manually"):
        detailing.resolve_upsampler_name("auto")


# ---------------------------------------------------------------------------
# strip_layout_conds / _downscale_to
# ---------------------------------------------------------------------------

def test_strip_layout_conds_removes_guide_keys_only():
    cond = [[torch.zeros(1, 4), {"keyframe_idxs": "kf", "guiding_latent": "g",
                                 "pooled_output": "keep"}]]
    out = detailing.strip_layout_conds(cond)
    assert out[0][1] == {"pooled_output": "keep"}
    # Original untouched (the scene's own conds must not lose their guides).
    assert "keyframe_idxs" in cond[0][1]

def test_downscale_to_shape_and_content():
    video = torch.arange(2 * 3 * 4 * 8 * 8, dtype=torch.float32).reshape(2, 3, 4, 8, 8)
    down = detailing._downscale_to(video, 4, 4)
    assert down.shape == (2, 3, 4, 4, 4)
    # The kernel is weight-normalised, so a constant tensor comes back exact.
    const = torch.full((1, 2, 2, 6, 6), 3.5)
    assert torch.allclose(detailing._downscale_to(const, 3, 3), torch.full((1, 2, 2, 3, 3), 3.5))


# ---------------------------------------------------------------------------
# detail_refine_scene
# ---------------------------------------------------------------------------

class _FakeChain:
    """Just enough of FunPackLTXAVSceneChainSampler for detail_refine_scene."""

    def __init__(self):
        self.chunk_calls = []

    def _latent_tensors(self, latent):
        samples = latent["samples"]
        if getattr(samples, "is_nested", False):
            return list(samples.unbind())
        return [samples]

    def _sample_chunk(self, model, sampler, sigmas, seed, cfg, positive, negative, latent,
                      **kwargs):
        self.chunk_calls.append({"sigmas": sigmas, "seed": seed,
                                 "positive": positive, "latent": latent})
        return latent  # identity "refine": returns the crop unchanged


def _scene_latent(nested=False):
    video = torch.randn(1, 8, 5, 16, 24)
    if not nested:
        return {"samples": video}, video
    audio = torch.randn(1, 4, 40)
    return {"samples": _FakeNested([video, audio])}, video

def test_refine_noop_without_targets_or_upsampler():
    chain = _FakeChain()
    latent, _ = _scene_latent()
    out, note = detailing.detail_refine_scene(
        chain, None, None, None, [], [], latent, "", object(), 0, 1.0)
    assert out is latent and note is None  # same object: bit-identical guarantee
    out, note = detailing.detail_refine_scene(
        chain, None, None, None, [], [], latent, "hands", None, 0, 1.0)
    assert out is latent and note is None
    assert chain.chunk_calls == []

def test_refine_noop_when_nothing_decodable(monkeypatch):
    """No keyframe could be decoded: same-object no-op, but WITH a diagnostic note —
    a miss must never look identical to "the feature is off" once it actually ran."""
    chain = _FakeChain()
    latent, _ = _scene_latent()
    monkeypatch.setattr(detailing, "_decode_detection_frames", lambda *a, **k: [])
    monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: None)
    out, note = detailing.detail_refine_scene(
        chain, None, None, None, [], [], latent, "hands", object(), 0, 1.0)
    assert out is latent
    assert note is not None and "no match" in note
    assert chain.chunk_calls == []

def test_refine_reports_max_heat_when_below_threshold(monkeypatch):
    """Regression: a real live run hit CLIPSeg-found-nothing with zero visibility into
    WHY (was it a bad prompt, or just a threshold too strict for CLIPSeg's raw score?).
    The note must carry the actual max score so the fix (lower detail_threshold) is
    obvious without re-running with debug_log."""
    chain = _FakeChain()
    latent, _ = _scene_latent()
    monkeypatch.setattr(detailing, "_decode_detection_frames",
                        lambda *a, **k: [(0, torch.zeros(16, 16, 3))])
    monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: torch.full((16, 16), 0.18))
    out, note = detailing.detail_refine_scene(
        chain, None, None, None, [], [], latent, "hands", object(), 0, 1.0, threshold=0.35)
    assert out is latent
    assert "0.18" in note and "0.35" in note
    assert chain.chunk_calls == []

def test_refine_reports_area_cap_reason(monkeypatch):
    chain = _FakeChain()
    latent, _ = _scene_latent()
    monkeypatch.setattr(detailing, "_decode_detection_frames",
                        lambda *a, **k: [(0, torch.zeros(160, 240, 3))])
    monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: torch.full((160, 240), 0.9))
    out, note = detailing.detail_refine_scene(
        chain, None, None, None, [], [], latent, "hands", object(), 0, 1.0)
    assert out is latent
    assert "cap" in note.lower()
    assert chain.chunk_calls == []

def test_refine_area_cap_is_a_cost_knob_not_a_judgment(monkeypatch):
    """The cap must never override what the user explicitly asked to detail — it's a
    cost guardrail the user can raise, not FunPack deciding a big region 'isn't worth
    detailing'. Same whole-frame region: refused at the default cap, allowed once the
    user raises detail_max_area (mirrors area_cap, threaded from the sampler node)."""
    chain = _FakeChain()
    latent, video = _scene_latent()
    monkeypatch.setattr(detailing, "_decode_detection_frames",
                        lambda *a, **k: [(0, torch.zeros(160, 240, 3))])
    monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: torch.full((160, 240), 0.9))
    monkeypatch.setattr(
        detailing, "_run_upsampler",
        lambda upsampler, crop, vae, debug=False: torch.nn.functional.interpolate(
            crop, scale_factor=(1, 2, 2), mode="nearest"))
    out_default, note_default = detailing.detail_refine_scene(
        chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0)
    assert out_default is latent and "cap" in note_default.lower()
    out_raised, note_raised = detailing.detail_refine_scene(
        chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0, area_cap=1.0)
    assert note_raised is not None and "segmented_detail(" in note_raised and "cap" not in note_raised.lower()
    assert out_raised is not latent

def _patch_detection_and_upsampler(monkeypatch):
    """Hot square in pixel space; upsampler = exact 2x nearest (model-free)."""
    heat = torch.zeros(160, 240)
    heat[40:80, 120:180] = 0.9
    monkeypatch.setattr(detailing, "_decode_detection_frames",
                        lambda *a, **k: [(0, torch.zeros(160, 240, 3))])
    monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: heat)
    monkeypatch.setattr(
        detailing, "_run_upsampler",
        lambda upsampler, crop, vae, debug=False: torch.nn.functional.interpolate(
            crop, scale_factor=(1, 2, 2), mode="nearest"))
    return heat

def test_refine_passes_renoise_sigma_to_the_tail_schedule(monkeypatch):
    chain = _FakeChain()
    latent, _ = _scene_latent()
    _patch_detection_and_upsampler(monkeypatch)
    detailing.detail_refine_scene(
        chain, "model", "vae", "sampler", [], [], latent, "hands", object(),
        seed=0, cfg=1.0, renoise_sigma=0.95)
    sigmas = chain.chunk_calls[0]["sigmas"]
    assert abs(float(sigmas[0]) - 0.95) < 1e-6
    assert torch.allclose(sigmas, torch.tensor(detailing.stage2_sigmas(0.95)))

def test_refine_sharpen_mode_skips_the_diffusion_tail_entirely(monkeypatch):
    """The whole point of sharpen mode: zero video-model forwards. It must use the
    upsampler's own output directly, not merely a fast-looking diffusion call."""
    chain = _FakeChain()
    latent, video = _scene_latent()
    _patch_detection_and_upsampler(monkeypatch)
    out, note = detailing.detail_refine_scene(
        chain, "model", "vae", "sampler", [], [], latent, "hands", object(),
        seed=0, cfg=1.0, mode="sharpen")
    assert note is not None and "sharpen" in note
    assert chain.chunk_calls == []  # no _sample_chunk call at all
    assert out["samples"].shape == video.shape
    diff = (out["samples"] - video).abs().amax(dim=(0, 1, 2))
    assert not diff[:2].any() and not diff[:, :2].any()  # far corners still untouched

def test_refine_repair_mode_is_still_the_default(monkeypatch):
    chain = _FakeChain()
    latent, _ = _scene_latent()
    _patch_detection_and_upsampler(monkeypatch)
    _, note = detailing.detail_refine_scene(
        chain, "model", "vae", "sampler", [], [], latent, "hands", object(), 0, 1.0)
    assert "repair" in note
    assert len(chain.chunk_calls) == 1

def test_refine_round_trip_touches_only_the_tube(monkeypatch):
    chain = _FakeChain()
    latent, video = _scene_latent()
    _patch_detection_and_upsampler(monkeypatch)
    pos = [[torch.zeros(1, 4), {"keyframe_idxs": "kf", "pooled_output": "keep"}]]
    out, note = detailing.detail_refine_scene(
        chain, "model", "vae", "sampler", pos, [], latent, "hands", object(),
        seed=11, cfg=1.0)
    assert note is not None and "segmented_detail" in note
    assert len(chain.chunk_calls) == 1
    call = chain.chunk_calls[0]
    # Official stage-2 tail schedule and a detail-offset seed.
    assert torch.allclose(call["sigmas"], torch.tensor(detailing.STAGE2_SIGMAS))
    assert call["seed"] == 11 + 7777
    # Layout conds stripped for the crop pass.
    assert "keyframe_idxs" not in call["positive"][0][1]
    assert call["positive"][0][1]["pooled_output"] == "keep"
    # The crop passed to the refine is 2x the tube size.
    crop = call["latent"]["samples"]
    out_video = out["samples"]
    assert out_video.shape == video.shape
    # With an identity refine, upsample(2x)->downscale(area) reproduces the crop, so
    # the paste changes (nearly) nothing — but crucially everything OUTSIDE the tube
    # must be exactly untouched regardless of what the refine did.
    diff = (out_video - video).abs().amax(dim=(0, 1, 2))  # [H, W]
    hot = diff > 1e-6
    assert not hot[:2].any() and not hot[:, :2].any()  # far corners untouched

def test_refine_preserves_audio_stream(monkeypatch):
    chain = _FakeChain()
    latent, video = _scene_latent(nested=True)
    audio = chain._latent_tensors(latent)[1]
    _patch_detection_and_upsampler(monkeypatch)
    out, note = detailing.detail_refine_scene(
        chain, "model", "vae", "sampler", [], [], latent, "hands", object(),
        seed=0, cfg=1.0)
    assert note is not None
    out_tensors = chain._latent_tensors(out)
    assert len(out_tensors) == 2
    # Audio comes back as the SAME tensor the scene finished with: protected by
    # construction, not by a lossy round trip.
    assert out_tensors[1] is audio
    # The crop pass carried an audio stream so joint attention saw a legal AV latent.
    crop_samples = chain.chunk_calls[0]["latent"]["samples"]
    assert getattr(crop_samples, "is_nested", False)

def test_refine_strength_scales_paste(monkeypatch):
    chain = _FakeChain()
    latent, video = _scene_latent()
    _patch_detection_and_upsampler(monkeypatch)
    # Non-identity refine: chunk returns a shifted crop so the paste has signal.
    def _shifted_chunk(model, sampler, sigmas, seed, cfg, positive, negative, crop_latent, **kw):
        return {"samples": crop_latent["samples"] + 10.0}
    chain._sample_chunk = _shifted_chunk
    out_full, _ = detailing.detail_refine_scene(
        chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0, strength=1.0)
    out_half, _ = detailing.detail_refine_scene(
        chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0, strength=0.5)
    d_full = (out_full["samples"] - video).abs().sum()
    d_half = (out_half["samples"] - video).abs().sum()
    assert d_full > 0
    assert abs(float(d_half / d_full) - 0.5) < 1e-3