v1.3.0 New Nodes: LTX Director and LTX Director Guide

patches.py

@@ -0,0 +1,167 @@
+import types
+import torch
+import comfy.ldm.modules.attention
+
+
+def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs):
+    # Bypass wrap_attn (sage/etc may ignore masks) by calling attention_pytorch directly.
+    return comfy.ldm.modules.attention.attention_pytorch(
+        q, k, v, heads, mask=mask,
+        _inside_attn_wrapper=True,
+        transformer_options=transformer_options,
+        **kwargs,
+    )
+
+
+def _wan_t2v_forward(self, mask_fn, x, context, transformer_options={}, **kwargs):
+    q = self.norm_q(self.q(x))
+    k = self.norm_k(self.k(context))
+    v = self.v(context)
+
+    mask = mask_fn(q, k, transformer_options)
+    if mask is not None:
+        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
+                              transformer_options=transformer_options)
+    else:
+        x = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
+        )
+    return self.o(x)
+
+
+def _wan_i2v_forward(self, mask_fn, x, context, context_img_len, transformer_options={}, **kwargs):
+    context_img = context[:, :context_img_len]
+    context_text = context[:, context_img_len:]
+
+    q = self.norm_q(self.q(x))
+
+    k_img = self.norm_k_img(self.k_img(context_img))
+    v_img = self.v_img(context_img)
+    img_x = comfy.ldm.modules.attention.optimized_attention(
+        q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options,
+    )
+
+    k = self.norm_k(self.k(context_text))
+    v = self.v(context_text)
+
+    mask = mask_fn(q, k, transformer_options)
+    if mask is not None:
+        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
+                              transformer_options=transformer_options)
+    else:
+        x = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
+        )
+
+    return self.o(x + img_x)
+
+
+def _ltx_forward(self, mask_fn, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
+    from comfy.ldm.lightricks.model import apply_rotary_emb
+
+    is_self_attn = context is None
+    context = x if is_self_attn else context
+
+    q = self.q_norm(self.to_q(x))
+    k = self.k_norm(self.to_k(context))
+    v = self.to_v(context)
+
+    if pe is not None:
+        q = apply_rotary_emb(q, pe)
+        k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
+
+    if not is_self_attn:
+        temporal_mask = mask_fn(q, k, transformer_options)
+        if temporal_mask is not None:
+            mask = temporal_mask if mask is None else mask + temporal_mask
+
+    if mask is None:
+        out = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, self.heads, attn_precision=self.attn_precision,
+            transformer_options=transformer_options,
+        )
+    else:
+        out = _masked_attention(q, k, v, self.heads, mask=mask,
+                                attn_precision=self.attn_precision,
+                                transformer_options=transformer_options)
+
+    if self.to_gate_logits is not None:
+        gate_logits = self.to_gate_logits(x)
+        b, t, _ = out.shape
+        out = out.view(b, t, self.heads, self.dim_head)
+        out = out * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
+        out = out.view(b, t, self.heads * self.dim_head)
+
+    return self.to_out(out)
+
+
+class _CrossAttnPatch:
+    """Descriptor that binds (impl, mask_fn) as a method onto a cross-attn module."""
+
+    def __init__(self, impl, mask_fn):
+        self.impl = impl
+        self.mask_fn = mask_fn
+
+    def __get__(self, obj, objtype=None):
+        impl, mask_fn = self.impl, self.mask_fn
+
+        def wrapped(self_module, *args, **kwargs):
+            return impl(self_module, mask_fn, *args, **kwargs)
+
+        return types.MethodType(wrapped, obj)
+
+
+def detect_model_type(model):
+    """Return (arch, patch_size, temporal_stride) for latent geometry.
+
+    temporal_stride is the VAE's pixel→latent temporal compression factor,
+    used to convert user-facing pixel frame counts to latent frames.
+    """
+    diff_model = model.model.diffusion_model
+
+    if hasattr(diff_model, "patch_size") and not hasattr(diff_model, "patchifier"):
+        return "wan", tuple(diff_model.patch_size), 4
+
+    if hasattr(diff_model, "patchifier"):
+        return "ltx", (1, 1, 1), int(diff_model.vae_scale_factors[0])
+
+    raise ValueError(
+        f"Unsupported model type: {type(diff_model).__name__}. "
+        f"Currently supports Wan and LTX models."
+    )
+
+
+def _check_unpatched(model_clone, key):
+    if key in getattr(model_clone, "object_patches", {}):
+        raise RuntimeError(
+            f"PromptRelay: cross-attention forward at '{key}' is already patched by "
+            "another node (e.g. KJNodes NAG). Stacking is not supported — remove the "
+            "conflicting node."
+        )
+
+
+def apply_patches(model_clone, arch, mask_fn):
+    diffusion_model = model_clone.get_model_object("diffusion_model")
+
+    if arch == "wan":
+        from comfy.ldm.wan.model import WanI2VCrossAttention
+        for idx, block in enumerate(diffusion_model.blocks):
+            key = f"diffusion_model.blocks.{idx}.cross_attn.forward"
+            _check_unpatched(model_clone, key)
+            cross_attn = block.cross_attn
+            impl = _wan_i2v_forward if isinstance(cross_attn, WanI2VCrossAttention) else _wan_t2v_forward
+            model_clone.add_object_patch(key, _CrossAttnPatch(impl, mask_fn).__get__(cross_attn, cross_attn.__class__))
+        return
+
+    if arch == "ltx":
+        for idx, block in enumerate(diffusion_model.transformer_blocks):
+            for attr in ("attn2", "audio_attn2"):
+                module = getattr(block, attr, None)
+                if module is None:
+                    continue
+                key = f"diffusion_model.transformer_blocks.{idx}.{attr}.forward"
+                _check_unpatched(model_clone, key)
+                model_clone.add_object_patch(key, _CrossAttnPatch(_ltx_forward, mask_fn).__get__(module, module.__class__))
+        return
+
+    raise ValueError(f"Unknown model arch: {arch}")