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

"""
ComfyUI node wrappers around the pure-math operations in perturbations.py.

Three nodes:

- GlitchSigmas — perturbed sigma schedule (MODEL → SIGMAS)
- AttentionDropoutPatcher — value-dropout patches on attention (MODEL → MODEL)
- LatentFrameInterleave — temporal reordering for video latents (LATENT → LATENT)

The earlier `GlitchSampler` node was removed (v0.2): its `extra_noise`
parameter was equivalent to `s_noise` on ComfyUI's built-in
`SamplerEulerAncestral`. Use that instead — it also exposes `eta`, which the
custom version didn't.
"""

from __future__ import annotations

from .perturbations import (
    apply_sigma_jitter,
    call_index_from_context,
    interleave_frames,
    value_dropout_patch,
)

# ComfyUI imports are deferred to call sites so this module can still be
# imported (e.g. by editors / language servers) outside a ComfyUI runtime.


SCHEDULER_NAMES = ["karras", "exponential", "normal", "simple", "ddim_uniform"]


class GlitchSigmas:
    """Perturbed sigma schedule. Outputs a SIGMAS tensor."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "scheduler": (SCHEDULER_NAMES, {"default": "karras"}),
                "steps": ("INT", {"default": 20, "min": 1, "max": 1000}),
                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "jitter": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "glitch-lab"
    FUNCTION = "build"

    def build(self, model, scheduler, steps, denoise, jitter, seed):
        import comfy.samplers

        total = steps if denoise >= 1.0 else int(steps / max(denoise, 1e-6))
        sigmas = comfy.samplers.calculate_sigmas(
            model.get_model_object("model_sampling"), scheduler, total
        )
        if denoise < 1.0:
            sigmas = sigmas[-(steps + 1):]
        sigmas = apply_sigma_jitter(sigmas, jitter, seed)
        return (sigmas,)


class AttentionDropoutPatcher:
    """Clones the model and installs value-dropout patches on attention layers.

    The patches derive their per-call deterministic seed from the (block, sigma)
    context ComfyUI passes in `extra_options`, not from a closure counter. The
    counter approach was broken — see `call_index_from_context` for the
    correction (credit: blepping, code review).
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "self_attn_dropout": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01}),
                "cross_attn_dropout": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("MODEL",)
    CATEGORY = "glitch-lab"
    FUNCTION = "patch"

    def patch(self, model, self_attn_dropout, cross_attn_dropout, seed):
        m = model.clone()

        def _sigma_from(extra_options):
            t = extra_options.get("transformer_options", {}) or {}
            s = t.get("sigmas")
            if s is None:
                return None
            try:
                return float(s[0]) if hasattr(s, "__len__") else float(s)
            except (TypeError, ValueError):
                return None

        def self_patch(q, k, v, extra_options):
            block = extra_options.get("block", ("?", 0))
            sigma = _sigma_from(extra_options)
            idx = call_index_from_context(seed, ("self", block), sigma)
            return value_dropout_patch(q, k, v, self_attn_dropout, seed, idx)

        def cross_patch(q, k, v, extra_options):
            block = extra_options.get("block", ("?", 0))
            sigma = _sigma_from(extra_options)
            idx = call_index_from_context(seed + 1, ("cross", block), sigma)
            return value_dropout_patch(q, k, v, cross_attn_dropout, seed + 1, idx)

        if self_attn_dropout > 0.0:
            m.set_model_attn1_patch(self_patch)
        if cross_attn_dropout > 0.0:
            m.set_model_attn2_patch(cross_patch)
        return (m,)


class LatentFrameInterleave:
    """Temporal reordering for 5-D (B, C, T, H, W) latents. 4-D latents pass through."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "intensity": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("LATENT",)
    CATEGORY = "glitch-lab"
    FUNCTION = "run"

    def run(self, latent, intensity, seed):
        samples = latent["samples"]
        out = interleave_frames(samples, intensity, seed)
        new_latent = {k: v for k, v in latent.items()}
        new_latent["samples"] = out
        return (new_latent,)


NODE_CLASS_MAPPINGS = {
    "GlitchSigmas": GlitchSigmas,
    "AttentionDropoutPatcher": AttentionDropoutPatcher,
    "LatentFrameInterleave": LatentFrameInterleave,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "GlitchSigmas": "Glitch Sigmas",
    "AttentionDropoutPatcher": "Attention Dropout Patcher",
    "LatentFrameInterleave": "Latent Frame Interleave",
}