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

"""second_pass — a general two-pass split on the Chain Sampler.

A scene can be finished differently from how it was started. The invariants that matter
here are the ones a split can silently get wrong:

1. BOTH schedules run in full, exactly as written. Pass 1 is not cut short and pass 2 is
   not derived from it — total steps are simply the two added up. (An earlier version cut
   pass 1 at pass 2's first sigma; it silently shortened the run and distorted the result.)
2. A hand-typed schedule is the one malformed-able input, and both ways it can be wrong
   are silent in the OUTPUT rather than loud at runtime, so they are refused up front.
3. Nothing is ever a silent no-op: every refusal comes back with a reason naming the fix.
4. The i2v anchor stays PINNED for pass 2 — without this pass 2 re-denoises the reference
   frame and the scene drifts off the image it was supposed to start from.
"""
import sys
import types
from pathlib import Path

import torch

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

for _name in (
    "comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling",
    "comfy.model_sampling", "comfy.nested_tensor", "comfy.sample",
    "comfy.samplers", "comfy.utils", "comfy.model_management",
):
    sys.modules.setdefault(_name, types.ModuleType(_name))
sys.modules["comfy.nested_tensor"].NestedTensor = object
# _run_upsampler (detailing.py) moves tensors to the torch device and reads the VAE's
# per-channel statistics; on CPU with vae=None both are no-ops, which is what we want here.
sys.modules["comfy.model_management"].get_torch_device = lambda: torch.device("cpu")
sys.modules["comfy"].model_management = sys.modules["comfy.model_management"]

import samplers  # noqa: E402


# The user's own 9-step schedule.
SCHEDULE = torch.tensor([1.0, 0.955, 0.893, 0.812, 0.715, 0.603, 0.482, 0.241, 0.121, 0.0])


def _node():
    return samplers.FunPackLTXAVSceneChainSampler()


# ── the pass-2 schedule ─────────────────────────────────────────────────────

def _sched(alt):
    return _node()._second_pass_schedule(alt)


def test_a_well_formed_schedule_is_taken_exactly_as_written():
    """Nothing is derived, cut or continued. Pass 1 runs the main schedule in full and pass 2
    runs this one in full, so on the user's 9-step main a 4-step second pass is 13 steps —
    not a re-slicing of the main schedule."""
    alt = torch.tensor([0.812, 0.6, 0.35, 0.15, 0.0])
    out, reason = _sched(alt)
    assert reason is None and out is alt
    assert (SCHEDULE.numel() - 1) + (out.numel() - 1) == 13


def test_an_ascending_schedule_is_refused():
    """Both malformed cases are silent in the OUTPUT rather than loud at runtime, so they
    have to be caught before sampling. An ascending pair walks the trajectory backwards."""
    out, reason = _sched(torch.tensor([0.812, 0.4, 0.55, 0.0]))
    assert out is None and "must descend" in reason


def test_a_schedule_that_stops_short_of_zero_is_refused():
    """Stopping above 0 leaves the clip partially denoised — the noise-artefact symptom."""
    out, reason = _sched(torch.tensor([0.812, 0.5, 0.25]))
    assert out is None and "not 0" in reason


def test_equal_neighbouring_sigmas_are_not_treated_as_ascending():
    out, reason = _sched(torch.tensor([0.812, 0.6, 0.6, 0.0]))
    assert reason is None and out is not None


def test_a_missing_or_degenerate_schedule_is_refused_with_a_reason():
    """second_pass with nothing to run must say so, not silently sample once."""
    for alt in (None, torch.tensor([]), torch.tensor([0.5]), "0.8, 0.0"):
        out, reason = _sched(alt)
        assert out is None and reason


def test_a_schedule_starting_at_zero_is_refused():
    out, reason = _sched(torch.tensor([0.0, 0.0]))
    assert out is None and "nothing for it to denoise" in reason


# ── the pin survives the split ──────────────────────────────────────────────

def test_the_i2v_anchor_is_re_pinned_for_pass_two():
    """_sample_chunk drops noise_mask from what it returns, so without this pass 2 runs
    unpinned and re-denoises the reference frame."""
    frames = 6
    clean = torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1).repeat(1, 4, 1, 2, 2)
    mask = torch.ones_like(clean)
    mask[:, :, :1] = 0.0                       # one pinned anchor frame
    chunk = {"samples": clean, "noise_mask": mask}
    # A "state" where the anchor frame has been polluted (as pass 2's re-noise would).
    state = {"samples": torch.full_like(clean, -99.0)}

    out = _node()._restore_pinned_prefix(state, chunk)
    assert out.get("noise_mask") is not None
    # The pinned frame is back to the clean anchor...
    assert torch.equal(out["samples"][:, :, :1], clean[:, :, :1])
    # ...and nothing else was touched.
    assert torch.equal(out["samples"][:, :, 1:], state["samples"][:, :, 1:])


def test_restore_is_a_no_op_on_a_t2v_scene():
    """No pinned prefix = no anchor to protect; the state must come back untouched."""
    clean = torch.zeros(1, 4, 6, 2, 2)
    state = {"samples": torch.full_like(clean, 3.0)}
    for chunk in ({"samples": clean, "noise_mask": torch.ones_like(clean)},
                  {"samples": clean}):
        assert _node()._restore_pinned_prefix(state, chunk) is state


# ── the optional between-pass operations ────────────────────────────────────

class _FakeUpsampler:
    """Stands in for Lightricks' LatentUpsampler: 2x nearest on the spatial axes."""
    def __call__(self, x):
        return x.repeat_interleave(2, dim=-1).repeat_interleave(2, dim=-2)

    def parameters(self):
        yield torch.zeros(1)

    def to(self, *a, **k):
        return self

    def cpu(self):
        return self


def _video(f=4, h=8, w=8, fill=2.0):
    return {"samples": torch.full((1, 4, f, h, w), fill)}


def test_no_operation_is_the_default_and_returns_the_same_object():
    """The whole feature is opt-in: 'none' must not clone, resample or touch anything."""
    lat = _video()
    for op in ("none", None, ""):
        out, note = _node()._second_pass_operate(lat, op, _FakeUpsampler(), None)
        assert out is lat and note is None


def test_sharpen_comes_back_at_the_original_size():
    """sharpen is upsampler -> downscale: the latent pass 2 receives must be the SAME shape,
    or it is silently an upscale and costs four times as much."""
    out, note = _node()._second_pass_operate(_video(h=8, w=8), "sharpen", _FakeUpsampler(), None)
    assert out["samples"].shape == (1, 4, 4, 8, 8)
    assert "sharpen" in note


def test_upscale_2x_keeps_the_larger_latent():
    out, note = _node()._second_pass_operate(_video(h=8, w=8), "upscale_2x", _FakeUpsampler(), None)
    assert out["samples"].shape == (1, 4, 4, 16, 16)
    assert "4x the pixels" in note


def test_a_missing_upsampler_degrades_instead_of_failing_the_render():
    """A whole montage must not die because one optional sharpen could not load its model."""
    lat = _video()
    out, note = _node()._second_pass_operate(lat, "sharpen", None, None)
    assert out is lat and "no latent upsampler" in note


def test_an_upsampler_that_throws_degrades_too():
    class _Boom(_FakeUpsampler):
        def __call__(self, x):
            raise RuntimeError("out of memory")
    lat = _video()
    out, note = _node()._second_pass_operate(lat, "sharpen", _Boom(), None)
    assert out is lat and "upsampler failed" in note


def test_an_unknown_operation_is_reported_not_guessed():
    lat = _video()
    out, note = _node()._second_pass_operate(lat, "enhance", _FakeUpsampler(), None)
    assert out is lat and "unknown operation" in note


def test_the_pin_survives_an_operation_that_changed_the_resolution():
    """upscale_2x leaves the chunk's mask the wrong size, which used to drop the pin and let
    pass 2 re-denoise the anchor at 4x the pixels — the run most likely to drift.

    The anchor is not re-derived from the source image: the pinned frames in the upscaled
    state ARE the anchor, carried through the same upsampler as every other frame. Only the
    mask has to be brought to the new grid.
    """
    clean = torch.zeros(1, 4, 4, 8, 8)
    mask = torch.ones_like(clean)
    mask[:, :, :1] = 0.0
    chunk = {"samples": clean, "noise_mask": mask}
    anchor = torch.full((1, 4, 1, 16, 16), 7.0)
    upscaled = {"samples": torch.zeros(1, 4, 4, 16, 16)}
    upscaled["samples"][:, :, :1] = anchor

    out = _node()._restore_pinned_prefix(upscaled, chunk)
    assert out is not upscaled
    m = out["noise_mask"]
    assert m.shape[-2:] == (16, 16)
    assert float(m[:, :, :1].max()) == 0.0      # anchor frame still pinned...
    assert float(m[:, :, 1:].min()) == 1.0      # ...and nothing else is
    # The upsampler's anchor is what stays pinned — untouched by the restore.
    assert torch.equal(out["samples"][:, :, :1], anchor)


def test_the_pin_is_dropped_when_the_mask_cannot_be_rebuilt():
    """A mask with no spatial axes to scale is unusable at the new size. Dropping the pin and
    saying so beats pinning against a grid that isn't there."""
    clean = torch.zeros(1, 4, 4, 8, 8)
    mask = torch.ones(1, 4, 4)                  # 3-D: no H/W to rescale
    mask[:, :1] = 0.0
    chunk = {"samples": clean, "noise_mask": mask}
    upscaled = {"samples": torch.zeros(1, 4, 4, 16, 16)}
    assert _node()._restore_pinned_prefix(upscaled, chunk) is upscaled