malkuthro-ComfyUI-Koolook/k_ltx_reference_bind_schedule.py

# SPDX-License-Identifier: GPL-3.0-or-later
#
# ComfyUI-Koolook — LTX Reference Bind Schedule node.
# Copyright (C) 2026 ComfyUI-Koolook contributors (kforgelabs).
"""LTX Reference Bind Schedule (Koolook) — ramp reference identity up at low sigma.

Companion to the LTX Director Ghost Mask reference. The reference image is added
as a guide frame; the core LTX transformer attends to it with a per-guide
**attention strength** (`guide_attention_entries[].strength`) which it turns into
a log-space additive self-attention bias, rebuilt **every denoise step** in
`comfy/ldm/lightricks/model.py::_build_guide_self_attention_mask`. That builder
explicitly supports `strength > 1.0` to *amplify* a guide's attention.

This node schedules that amplification over the denoise so identity locks late
without disturbing motion that forms early:

* **Early steps (high sigma)** — gain ~1.0: the reference attention is neutral,
  big motion resolves freely.
* **Late steps (low sigma)** — gain ramps up to `peak_strength`: the refinement
  conforms harder to the reference's identity (face / mouth).

It's the inverse of the motion curve, and the mirror of `LTXAVBindSchedule`
(which ramps audio->video up at low sigma for lip-sync). It touches ONLY the
attention bias of the trailing `num_references` guide entries (references are
appended after the keyframes) — the keyframe pins and the noise-mask freezing
are untouched.

WORKS WITHOUT AN IC-LoRA. The attention-bias machinery is BASE-model
functionality — `_build_guide_self_attention_mask` builds the mask from
`resolved_guide_entries` regardless of any LoRA. The catch is that
LTXDirectorGuide only *populates* those entries `if is_lora_active`, so in the
no-LoRA Ghost Mask path the list is empty. This node closes that gap: when the
entries are missing it **fabricates them itself** for the trailing reference
tokens (splitting `num_guide_tokens` across the `num_keyframes + num_references`
guide frames and assigning the ramped strength to the reference share). So the
reference gets the same amplified-attention channel an IC-LoRA would use, inside
the Director, with no LoRA, no guide-node vendor, and no graph change. When an
IC-LoRA *is* active and real entries exist, it scales those instead.

Mechanics (same as the A/V node): a `model_function_wrapper` captures the
current sigma each step and stashes the gain in `transformer_options`; an
object-patch on `_build_guide_self_attention_mask` reads it and fabricates (or
scales) the reference entries' `strength` before the mask is built.

NOTE: the `ref_gain` ramp is unit-tested; the model wiring is validated by
rendering against a live LTX 2.3 model (not in CI).
"""
from __future__ import annotations

import types
import logging

log = logging.getLogger(__name__)

_REF_GAIN_KEY = "nghtdrp_ref_gain"  # per-step reference-attention multiplier


def ref_gain(sigma, sigma_max, ramp_start=0.5, ramp_end=0.9, peak=2.0):
    """Reference-attention multiplier for the current ``sigma``.

    Progress runs 0 at the start of denoise (``sigma == sigma_max``) to 1 at the
    end (``sigma -> 0``). Holds at 1.0 (neutral) until ``ramp_start``, then
    smoothsteps up to ``peak`` by ``ramp_end`` (progress fractions in [0, 1]).
    Returning 1.0 means "no change" — the mask builder then skips entirely.
    """
    if sigma_max <= 0.0:
        return float(peak)
    progress = 1.0 - (float(sigma) / float(sigma_max))
    progress = 0.0 if progress < 0.0 else (1.0 if progress > 1.0 else progress)
    if ramp_end <= ramp_start:
        return float(peak) if progress >= ramp_end else 1.0
    if progress <= ramp_start:
        return 1.0
    if progress >= ramp_end:
        return float(peak)
    t = (progress - ramp_start) / (ramp_end - ramp_start)
    s = t * t * (3.0 - 2.0 * t)  # smoothstep
    return 1.0 + (float(peak) - 1.0) * s


def _boosted_entries(resolved_entries, ref_count, gain):
    """Return a copy of ``resolved_entries`` with the LAST ``ref_count`` entries'
    ``strength`` multiplied by ``gain`` (references are appended after keyframes).
    Returns the original list unchanged when there's nothing to do."""
    if not resolved_entries or gain == 1.0 or ref_count <= 0:
        return resolved_entries
    n = len(resolved_entries)
    k = min(int(ref_count), n)
    if k <= 0:
        return resolved_entries
    out = list(resolved_entries)
    for i in range(n - k, n):
        e = dict(out[i])
        e["strength"] = float(e.get("strength", 1.0)) * float(gain)
        out[i] = e
    return out


def fabricate_reference_entries(num_guide_tokens, num_references, num_keyframes, strength):
    """Build ``resolved_guide_entries`` for the trailing reference tokens when the
    guide node created none (the no-IC-LoRA path).

    Splits ``num_guide_tokens`` evenly across the ``num_keyframes + num_references``
    guide frames and gives the trailing reference share ``strength`` while the
    keyframe share stays 1.0 (neutral). The entries match the shape the core
    builder consumes (``strength`` + ``surviving_count`` + optional pixel_mask /
    latent_shape). When ``num_keyframes`` is 0 the whole guide region is treated
    as reference. Returns ``[]`` when there's nothing to amplify.
    """
    n_guide = int(num_guide_tokens)
    n_ref = int(num_references)
    n_kf = max(0, int(num_keyframes))
    if n_guide <= 0 or n_ref <= 0:
        return []
    total_frames = n_ref + n_kf
    ref_tokens = round(n_guide * n_ref / total_frames) if total_frames > 0 else n_guide
    ref_tokens = max(1, min(int(ref_tokens), n_guide))
    kf_tokens = n_guide - ref_tokens
    entries = []
    if kf_tokens > 0:
        entries.append({"strength": 1.0, "surviving_count": kf_tokens, "pixel_mask": None, "latent_shape": None})
    entries.append({"strength": float(strength), "surviving_count": ref_tokens, "pixel_mask": None, "latent_shape": None})
    return entries


def _make_guide_mask_wrapper(underlying, ref_count, kf_count, peak, diag):
    """Bound-method wrapper around ``_build_guide_self_attention_mask``.

    If the guide node already populated entries (IC-LoRA path) scale the trailing
    ``ref_count`` of them by the per-step gain. If not (no-LoRA Ghost Mask path)
    fabricate entries for the trailing reference tokens at the gain strength. Then
    hand off to the original builder, which turns it into the additive bias.
    """

    def wrapped(_self_module, x, transformer_options, merged_args, *args, **kwargs):
        to = transformer_options or {}
        gain = float(to.get(_REF_GAIN_KEY, 1.0))
        ma = merged_args if isinstance(merged_args, dict) else {}
        entries = ma.get("resolved_guide_entries")
        n_guide = int(ma.get("num_guide_tokens") or 0)

        mode = "passthrough"
        if entries:  # IC-LoRA path — scale real entries
            if gain != 1.0 and ref_count > 0:
                boosted = _boosted_entries(entries, ref_count, gain)
                if boosted is not entries:
                    merged_args = {**ma, "resolved_guide_entries": boosted}
                    mode = "scaled"
        elif n_guide > 0 and ref_count > 0 and gain != 1.0:  # no-LoRA path — fabricate
            fab = fabricate_reference_entries(n_guide, ref_count, kf_count, gain)
            if fab:
                merged_args = {**ma, "resolved_guide_entries": fab}
                mode = "fabricated"

        if not diag.get("first"):
            diag["first"] = True
            log.info(
                "[LTXReferenceBindSchedule] active: num_guide_tokens=%s entries=%s "
                "ref_count=%d num_keyframes=%d peak=%.2f (amplifies the trailing "
                "reference tokens at low sigma).",
                n_guide, (len(entries) if entries else 0), int(ref_count),
                int(kf_count), float(peak),
            )
        if mode in ("scaled", "fabricated") and not diag.get(mode):
            diag[mode] = True
            log.info("[LTXReferenceBindSchedule] %s reference attention (gain x%.2f).", mode, gain)
        return underlying(x, transformer_options, merged_args, *args, **kwargs)

    return wrapped


class LTXReferenceBindSchedule:
    """Ramp the LTX Ghost Mask reference's attention up as sigma falls."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "LTX 2.3 model (the same one feeding the guide pipeline)."}),
                "num_references": ("INT", {
                    "default": 1, "min": 0, "max": 64, "step": 1,
                    "tooltip": "How many trailing guide frames are references (= how many images you fed LTX Director's reference_images). They're appended after the keyframes, so the LAST N guide frames are the references.",
                }),
                "peak_strength": ("FLOAT", {
                    "default": 2.0, "min": 1.0, "max": 8.0, "step": 0.05,
                    "tooltip": "Reference attention multiplier at the end of denoise (low sigma). 1.0 = no boost; 2.0 = double the identity pull during refinement. Try 2-4.",
                }),
            },
            "optional": {
                # NOTE: ramp_start/ramp_end stay first to preserve widget order
                # for workflows saved before num_keyframes existed; num_keyframes
                # is appended last so old [num_references, peak_strength, ramp_start,
                # ramp_end] values keep aligning.
                "ramp_start": ("FLOAT", {
                    "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
                    "tooltip": "Denoise progress (0=start, 1=end) where the reference boost begins. Keep above where big motion finalizes (~0.5-0.6) so motion forms first.",
                }),
                "ramp_end": ("FLOAT", {
                    "default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01,
                    "tooltip": "Denoise progress where the boost reaches peak_strength.",
                }),
                "num_keyframes": ("INT", {
                    "default": 0, "min": 0, "max": 256, "step": 1,
                    "tooltip": "No-LoRA path only: how many timeline image keyframes precede the references, so the node knows the keyframe/reference token split. Set this = your number of timeline image segments (e.g. 4). Leave 0 and the WHOLE guide region is treated as reference (also boosts keyframes). Ignored when an IC-LoRA populated real entries.",
                }),
            },
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "patch"
    CATEGORY = "Koolook/LTX"
    DESCRIPTION = (
        "Ramp the reference's guide-attention UP at low sigma so identity locks "
        "during refinement without disturbing early motion. Works WITHOUT an "
        "IC-LoRA: it fabricates the per-guide attention entries for the trailing "
        "reference tokens itself (set num_keyframes so it knows the split), so the "
        "reference gets the same amplified-attention channel an IC-LoRA would use "
        "— inside the Director, no guide-node changes. Mirror of LTX A/V Bind "
        "Schedule. One pass, no re-noise."
    )

    def patch(self, model, num_references, peak_strength, num_keyframes=0, ramp_start=0.5, ramp_end=0.9):
        m = model.clone()

        if int(num_references) <= 0:
            log.info("[LTXReferenceBindSchedule] num_references=0 — model passed through.")
            return (m,)

        try:
            sigma_max = float(m.get_model_object("model_sampling").sigma_max)
        except Exception as e:  # pragma: no cover - live model
            log.warning("[LTXReferenceBindSchedule] no model_sampling.sigma_max: %s", e)
            sigma_max = 1.0

        key = "diffusion_model._build_guide_self_attention_mask"
        try:
            underlying = m.get_model_object(key)
            dm = m.get_model_object("diffusion_model")
        except Exception as e:  # pragma: no cover - live model
            log.warning(
                "[LTXReferenceBindSchedule] no %s on this model (not an LTX guide "
                "model?): %s — passed through unchanged.", key, e,
            )
            return (m,)

        diag = {}
        wrapper = _make_guide_mask_wrapper(
            underlying, int(num_references), int(num_keyframes), float(peak_strength), diag,
        )
        try:
            m.add_object_patch(key, types.MethodType(wrapper, dm))
        except Exception as e:  # pragma: no cover - live model
            log.warning("[LTXReferenceBindSchedule] guide-mask patch failed: %s", e)
            return (m,)

        prev = m.model_options.get("model_function_wrapper")

        def unet_wrapper(model_function, params):
            try:
                sigma = float(params["timestep"].max())
            except Exception:
                sigma = sigma_max
            gain = ref_gain(sigma, sigma_max, ramp_start, ramp_end, peak_strength)
            c = params["c"]
            to = dict(c.get("transformer_options", {}))
            to[_REF_GAIN_KEY] = gain
            c = {**c, "transformer_options": to}
            params = {**params, "c": c}
            if prev is not None:
                return prev(model_function, params)
            return model_function(params["input"], params["timestep"], **params["c"])

        m.set_model_unet_function_wrapper(unet_wrapper)
        log.info(
            "[LTXReferenceBindSchedule] scheduling %d reference(s): peak x%.2f, "
            "ramp %.2f..%.2f of denoise.", int(num_references), float(peak_strength),
            float(ramp_start), float(ramp_end),
        )
        return (m,)


NODE_CLASS_MAPPINGS = {"LTXReferenceBindSchedule": LTXReferenceBindSchedule}
NODE_DISPLAY_NAME_MAPPINGS = {"LTXReferenceBindSchedule": "LTX Reference Bind Schedule (Koolook)"}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]