HitmanLoges-comfyui-glitch-lab-sampler/sampler/test_sampler.py

"""
Standalone tests for sampler/perturbations.py and node shape/registration.

Requires only torch — no ComfyUI runtime. Run from the sampler/ directory:

    python test_sampler.py

Exits non-zero on any failure. No external test framework dependency.
"""

from __future__ import annotations

import io
import sys
import traceback
from pathlib import Path

# Force UTF-8 stdout so unicode arrows in test names don't crash on Windows cp1252.
if hasattr(sys.stdout, "buffer"):
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", line_buffering=True)


def _bootstrap_imports():
    """Allow running this file directly without `python -m sampler.test_sampler`."""
    here = Path(__file__).resolve().parent
    sys.path.insert(0, str(here.parent))


_bootstrap_imports()

import torch  # noqa: E402

from sampler.perturbations import (  # noqa: E402
    apply_sigma_jitter,
    call_index_from_context,
    interleave_frames,
    value_dropout_patch,
)


_failures: list[str] = []


def _check(name: str, cond: bool, detail: str = ""):
    if cond:
        print(f"  ok   {name}")
    else:
        msg = f"  FAIL {name}" + (f" — {detail}" if detail else "")
        print(msg)
        _failures.append(name)


# ---------------------------------------------------------------------------
# apply_sigma_jitter
# ---------------------------------------------------------------------------

def test_sigma_jitter():
    print("apply_sigma_jitter:")
    base = torch.linspace(14.0, 0.0, steps=21)

    out0 = apply_sigma_jitter(base, jitter=0.0, seed=42)
    _check("zero jitter equals base elementwise", torch.equal(out0, base))

    out1 = apply_sigma_jitter(base, jitter=0.3, seed=42)
    _check("jittered length preserved", out1.shape == base.shape)
    _check("jittered terminal sigma unchanged", out1[-1].item() == base[-1].item())
    _check("jittered differs from base", not torch.equal(out1, base))
    _check("jittered intermediates strictly positive", bool((out1[:-1] > 0).all()))

    out2 = apply_sigma_jitter(base, jitter=0.3, seed=42)
    _check("same seed -> same output", torch.equal(out1, out2))

    out3 = apply_sigma_jitter(base, jitter=0.3, seed=43)
    _check("different seed -> different output", not torch.equal(out1, out3))


# ---------------------------------------------------------------------------
# value_dropout_patch
# ---------------------------------------------------------------------------

def test_value_dropout():
    print("value_dropout_patch:")
    q = torch.randn(1, 16, 64)
    k = torch.randn(1, 16, 64)
    v = torch.randn(1, 16, 64)

    qo, ko, vo = value_dropout_patch(q, k, v, dropout=0.0, seed=0, call_idx=0)
    _check("zero dropout: v unchanged", torch.equal(vo, v))
    _check("zero dropout: q passes through", qo is q)
    _check("zero dropout: k passes through", ko is k)

    _, _, v1 = value_dropout_patch(q, k, v, dropout=0.5, seed=7, call_idx=0)
    _check("half dropout: shape preserved", v1.shape == v.shape)
    _check("half dropout: some rows zeroed", bool((v1.abs().sum(dim=-1) == 0).any()))

    _, _, v1b = value_dropout_patch(q, k, v, dropout=0.5, seed=7, call_idx=0)
    _check("determinism: same seed/idx -> same output", torch.equal(v1, v1b))

    _, _, v1c = value_dropout_patch(q, k, v, dropout=0.5, seed=7, call_idx=1)
    _check("call_idx changes pattern", not torch.equal(v1, v1c))


# ---------------------------------------------------------------------------
# call_index_from_context (new — replaces closure counter)
# ---------------------------------------------------------------------------

def test_call_index_from_context():
    print("call_index_from_context:")
    a = call_index_from_context(seed=42, block=("input", 0), sigma=14.6)
    b = call_index_from_context(seed=42, block=("input", 0), sigma=14.6)
    _check("same inputs -> same index", a == b)

    c = call_index_from_context(seed=42, block=("input", 0), sigma=2.1)
    _check("different sigma -> different index", a != c)

    d = call_index_from_context(seed=42, block=("output", 5), sigma=14.6)
    _check("different block -> different index", a != d)

    e = call_index_from_context(seed=43, block=("input", 0), sigma=14.6)
    _check("different seed -> different index", a != e)

    f = call_index_from_context(seed=42, block=("input", 0), sigma=None)
    _check("None sigma is accepted (no crash)", isinstance(f, int))

    _check("index fits in 32 bits", 0 <= a < (1 << 32))


# ---------------------------------------------------------------------------
# interleave_frames — note (B, C, T, H, W) layout, not the previous wrong assumption
# ---------------------------------------------------------------------------

def test_interleave():
    print("interleave_frames:")
    img4d = torch.randn(1, 4, 64, 64)
    out4d = interleave_frames(img4d, intensity=0.7, seed=0)
    _check("4-D input passes through unchanged", torch.equal(out4d, img4d))

    # Video latent in ComfyUI's actual layout: (batch, channels, time, height, width)
    vid5d = torch.randn(1, 4, 8, 32, 32)

    out0 = interleave_frames(vid5d, intensity=0.0, seed=0)
    _check("zero intensity: identity", torch.equal(out0, vid5d))

    out_mid = interleave_frames(vid5d, intensity=0.4, seed=11)
    _check("mid intensity: shape preserved", out_mid.shape == vid5d.shape)
    # Per-frame sums (summing over B, C, H, W and leaving T) should be the same
    # set after permutation, just reordered.
    _check(
        "mid intensity: same set of frames (per-frame sums match after sort)",
        torch.equal(
            out_mid.sum(dim=(0, 1, 3, 4)).sort().values,
            vid5d.sum(dim=(0, 1, 3, 4)).sort().values,
        ),
    )

    out_high = interleave_frames(vid5d, intensity=0.9, seed=11)
    _check("high intensity: shape preserved", out_high.shape == vid5d.shape)
    _check(
        "high intensity: elementwise sum preserved",
        torch.isclose(out_high.sum(), vid5d.sum()),
    )

    # Operates on time, not channels: if you sum over (B, T, H, W) you should
    # get the SAME per-channel totals (unchanged), proving channels were
    # untouched.
    _check(
        "high intensity: per-channel sums unchanged (proves T not C is the permuted axis)",
        torch.allclose(
            out_high.sum(dim=(0, 2, 3, 4)),
            vid5d.sum(dim=(0, 2, 3, 4)),
        ),
    )


# ---------------------------------------------------------------------------
# Node registration (no ComfyUI runtime needed; nodes.py defers comfy imports)
# ---------------------------------------------------------------------------

def test_node_shapes():
    print("node shapes & registration:")
    from sampler.nodes import (
        NODE_CLASS_MAPPINGS,
        NODE_DISPLAY_NAME_MAPPINGS,
        AttentionDropoutPatcher,
        GlitchSigmas,
        LatentFrameInterleave,
    )

    expected = {"GlitchSigmas", "AttentionDropoutPatcher", "LatentFrameInterleave"}
    _check(
        "NODE_CLASS_MAPPINGS has all three keys",
        set(NODE_CLASS_MAPPINGS.keys()) == expected,
        detail=str(set(NODE_CLASS_MAPPINGS.keys())),
    )
    _check(
        "NODE_DISPLAY_NAME_MAPPINGS covers all three",
        set(NODE_DISPLAY_NAME_MAPPINGS.keys()) == expected,
    )
    _check(
        "GlitchSampler is gone (replaced by built-in SamplerEulerAncestral)",
        "GlitchSampler" not in NODE_CLASS_MAPPINGS,
    )

    for cls in (GlitchSigmas, AttentionDropoutPatcher, LatentFrameInterleave):
        it = cls.INPUT_TYPES()
        _check(f"{cls.__name__}.INPUT_TYPES has 'required'", "required" in it)
        _check(f"{cls.__name__} has RETURN_TYPES", hasattr(cls, "RETURN_TYPES"))
        _check(f"{cls.__name__} has CATEGORY=glitch-lab", getattr(cls, "CATEGORY", None) == "glitch-lab")
        _check(f"{cls.__name__} has FUNCTION attr", isinstance(getattr(cls, "FUNCTION", None), str))


def main() -> int:
    suites = [
        test_sigma_jitter,
        test_value_dropout,
        test_call_index_from_context,
        test_interleave,
        test_node_shapes,
    ]
    for suite in suites:
        try:
            suite()
        except Exception:
            print(f"  EXCEPTION in {suite.__name__}:")
            traceback.print_exc()
            _failures.append(f"{suite.__name__} raised")

    print()
    if _failures:
        print(f"{len(_failures)} failure(s):")
        for f in _failures:
            print(f"  - {f}")
        return 1
    print("All checks passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())