From a4388d617f6b81ff1770357e28add735c3fb6632 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:52:12 +0300 Subject: [PATCH 01/15] Best-FaceID identity transfer: full native overlap+ArcFace port Replaces the source-phase-RoPE-only approximation with a faithful native port of ComfyUI-BFSNodes' LTX Identity Transfer (read upstream source directly): the reference face is now injected as separate, non-rendered overlap tokens (clean timestep, RoPE source-phase tag, trimmed before output) instead of blending into Continuity's identity-pin keyframe, plus an optional ArcFace embedding -> IdentityProjector -> context-token channel (identity_projector, source_id, phase_scale, id_strength, arcface_mode, debug_log). Stays on Chain Sampler per the Studio/Sampler boundary law's guide-ops carve-out. New identity_transfer.py module; insightface added as a soft pip dependency. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 11 + identity_transfer.py | 145 +++++++++ movie_editor/frontend/engine_settings.js | 61 +++- requirements.txt | 2 + samplers.py | 379 ++++++++++++++--------- tests/test_identity_transfer.py | 354 +++++++++++++-------- 6 files changed, 672 insertions(+), 280 deletions(-) create mode 100644 identity_transfer.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e4f3a5b..de8a3de5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ ## [Unreleased] +### Changed +- **Best-FaceID identity transfer** (experimental, Engine ▸ Experimental) is now a full native + port of ComfyUI-BFSNodes' LTX Identity Transfer overlap+source_phase+ArcFace conditioning, + replacing 3.2.1's source-phase-only approximation. The reference face is now injected as + separate, non-rendered overlap tokens (never blended into a real frame) instead of tagging + Continuity's Identity pin keyframe, plus an optional ArcFace projector channel (`identity_projector`, + `id_strength`, `arcface_mode`) for a second identity signal. Still keyed off Continuity's + existing Identity pin image — no new picker. Requires the `insightface` package (added to + requirements.txt; downloads the buffalo_l model on first use) only when an ArcFace projector + is selected — the overlap tokens alone work without it. + ## [3.2.1] - 2026-07-06 ### Added diff --git a/identity_transfer.py b/identity_transfer.py new file mode 100644 index 00000000..172c26a3 --- /dev/null +++ b/identity_transfer.py @@ -0,0 +1,145 @@ +"""Best-FaceID identity transfer — native port of ComfyUI-BFSNodes' LTX Identity Transfer +(alisson-anjos/ComfyUI-BFSNodes, LTXIdentityOverlapConditioning): the reference face is +injected as SEPARATE tokens that share the target's frame-0 RoPE grid (overlap) and are +tagged with a per-source RoPE phase (source_phase), exactly as ltx-trainer's flexible +strategy did at train/validation time. The ref tokens are clean (timestep 0), attend to the +target in self-attention, and are sliced off the output (never rendered) — unlike FunPack's +guide-attention path, they never blend into a real frame. Plus optional ArcFace IdentityProjector +tokens appended to the text context (as trained). + +Requires `insightface` (+ buffalo_l, auto-downloaded on first use, CPU) for the ArcFace half; +soft dependency — the overlap tokens alone (identity_projector="None") work without it. +""" +from __future__ import annotations + +import logging +import os + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors.torch import load_file + +log = logging.getLogger("FunPackIdentityTransfer") + +_USE_GPU = os.environ.get("FUNPACK_ARCFACE_GPU", "0") == "1" +_FACE_APP = None + + +def _get_face_app(): + global _FACE_APP + if _FACE_APP is None: + from insightface.app import FaceAnalysis + providers = (["CUDAExecutionProvider", "CPUExecutionProvider"] if _USE_GPU else ["CPUExecutionProvider"]) + app = FaceAnalysis(name="buffalo_l", providers=providers) + app.prepare(ctx_id=0 if _USE_GPU else -1, det_size=(640, 640)) + _FACE_APP = app + return _FACE_APP + + +def arcface_embed(image_bhwc, mode="auto_adjust"): + """Return the ArcFace embedding [512], or None if disabled / no face found. + mode: 'as_is' (detect on the image only), 'auto_adjust' (retry with border-pad + zoom-out + upscale when detection fails), 'disable' (skip ArcFace entirely).""" + if mode == "disable": + return None + import cv2 + app = _get_face_app() + img = np.ascontiguousarray( + (np.clip(image_bhwc[0].detach().cpu().numpy(), 0.0, 1.0) * 255.0).astype(np.uint8)[:, :, ::-1] + ) + attempts = [img] + if mode == "auto_adjust": + h, w = img.shape[:2] + pad = int(0.4 * max(h, w)) + attempts.append(cv2.copyMakeBorder(img, pad, pad, pad, pad, cv2.BORDER_REPLICATE)) + attempts.append(cv2.resize(img, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)) + attempts.append(cv2.resize(img, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)) + for a in attempts: + faces = app.get(a) + if faces: + f = max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) + return torch.from_numpy(f.normed_embedding.astype(np.float32)) + return None + + +class IdentityProjector(nn.Module): + """ArcFace embedding [in_dim] -> num_tokens context tokens [num_tokens, context_dim].""" + + def __init__(self, in_dim=512, context_dim=4096, num_tokens=4, proj=None): + super().__init__() + self.in_dim = in_dim + self.context_dim = context_dim + self.num_tokens = num_tokens + self.proj = proj if proj is not None else nn.Sequential( + nn.Linear(in_dim, 1024), nn.GELU(), nn.Linear(1024, num_tokens * context_dim)) + self.norm = nn.LayerNorm(context_dim) + + def forward(self, e): + return self.norm(self.proj(e).reshape(-1, self.num_tokens, self.context_dim)) + + +def load_identity_projector(path, device): + """Load an IdentityProjector of ANY depth by rebuilding proj.N Linears from the + state_dict (old = 2 linears proj.0/proj.2; enhanced = 3 linears proj.0/proj.2/proj.4).""" + sd = load_file(path) + context_dim = sd["norm.weight"].shape[0] + idxs = sorted({int(k.split(".")[1]) for k in sd if k.startswith("proj.") and k.endswith(".weight")}) + layers = [] + for j, i in enumerate(idxs): + w = sd[f"proj.{i}.weight"] + layers.append(nn.Linear(w.shape[1], w.shape[0])) + if j < len(idxs) - 1: + layers.append(nn.GELU()) + proj = nn.Sequential(*layers) + in_dim = sd["proj.0.weight"].shape[1] + num_tokens = sd[f"proj.{idxs[-1]}.weight"].shape[0] // context_dim + p = IdentityProjector(in_dim=in_dim, context_dim=context_dim, num_tokens=num_tokens, proj=proj) + p.load_state_dict(sd) + return p.to(device=device, dtype=torch.float32).eval() + + +def append_context_tokens(conditioning, tokens): + """Append `tokens` [1 or B, N, D'] onto every entry's text context, padding/truncating + the last dim to match and extending any attention_mask with all-ones for the new tokens.""" + out = [] + for ce, cd in conditioning: + t = tokens.to(device=ce.device, dtype=ce.dtype) + if t.shape[0] != ce.shape[0]: + t = t.expand(ce.shape[0], -1, -1) + if t.shape[-1] < ce.shape[-1]: + t = F.pad(t, (0, ce.shape[-1] - t.shape[-1])) + elif t.shape[-1] > ce.shape[-1]: + t = t[..., :ce.shape[-1]] + nd = cd.copy() + am = nd.get("attention_mask") + if am is not None: + nd["attention_mask"] = torch.cat( + [am, torch.ones((*am.shape[:-1], t.shape[1]), device=am.device, dtype=am.dtype)], dim=-1) + out.append([torch.cat([ce, t], dim=1), nd]) + return out + + +def rotate_overlap_freqs(pe, ref_len, seg_value, theta=10000.0): + """Rotate the LAST ref_len tokens' RoPE freqs by phase = seg_value * theta^(-d/L) + (port of ltx_core rope.apply_segment_phase). pe = (cos, sin, [split_flag]).""" + if ref_len <= 0 or seg_value == 0.0: + return pe + cos, sin = pe[0], pe[1] + rest = tuple(pe[2:]) + L = cos.shape[-1] + d = torch.arange(L, device=cos.device, dtype=torch.float32) + rate = theta ** (-d / float(L)) + phase = seg_value * rate + pc = phase.cos().to(cos.dtype) + ps = phase.sin().to(sin.dtype) + idx = [slice(None)] * cos.dim() + idx[-2] = slice(cos.shape[-2] - ref_len, cos.shape[-2]) + idx = tuple(idx) + c0, s0 = cos[idx], sin[idx] + cos = cos.clone() + sin = sin.clone() + cos[idx] = c0 * pc - s0 * ps + sin[idx] = s0 * pc + c0 * ps + return (cos, sin, *rest) diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 6eea3300..0dbd4870 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -5,6 +5,7 @@ (function () { const { el, clear } = window.dom; const S = window.Store; + const API = window.MovieEditorAPI; let _mounted = null; // { scroller (pane, set per render), content (shell root) } let unsub = null; @@ -180,9 +181,15 @@ { name: "dynashift_strength", label: "DynaShift strength", kind: "float", default: 0.3, min: 0.05, max: 1.0, step: 0.05, dependsOn: "dynashift" }, { name: "dynashift_threshold", label: "DynaShift match threshold", kind: "float", default: 0.6, min: 0.3, max: 0.95, step: 0.05, dependsOn: "dynashift" }, { name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false, - hint: "Tags Continuity's Identity pin guide (Engine → Continuity) with the source-phase RoPE rotation Best-FaceID-style identity LoRAs were trained on. Load the LoRA itself the normal way — Models → add a LoRA loader onto the model path. No effect without an Identity pin image set." }, - { name: "identity_transfer_phase", label: "Source-phase id", kind: "float", default: 2.0, min: 0.0, max: 8.0, step: 0.5, dependsOn: "identity_transfer_enabled", - hint: "Matches the LoRA's training convention (ltx-trainer used 2)." }, + hint: "Full native port of the overlap+source_phase+ArcFace conditioning Best-FaceID-style identity LoRAs were trained on. Replaces Continuity's Identity pin guide (Engine → Continuity) with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself the normal way — Models → add a LoRA loader onto the model path. No effect without an Identity pin image set." }, + { name: "source_id", label: "Source-phase id", kind: "float", default: 2.0, min: 0.0, max: 8.0, step: 1.0, dependsOn: "identity_transfer_enabled", + hint: "Matches the LoRA's training convention (ltx-trainer used 2). 0 disables the RoPE rotation." }, + { name: "phase_scale", label: "Phase scale", kind: "float", default: 1.0, min: 0.0, max: 4.0, step: 0.1, dependsOn: "identity_transfer_enabled" }, + { name: "id_strength", label: "ArcFace token strength", kind: "float", default: 1.0, min: 0.0, max: 50.0, step: 0.5, dependsOn: "identity_transfer_enabled", + hint: "Only applies when an ArcFace projector is selected below — weak channel, push high (5-20) to test." }, + { name: "arcface_mode", label: "ArcFace detection mode", kind: "combo", choices: ["auto_adjust", "as_is", "disable"], default: "auto_adjust", dependsOn: "identity_transfer_enabled" }, + { name: "debug_log", label: "Debug log", kind: "bool", default: false, dependsOn: "identity_transfer_enabled", + hint: "Print per-scene identity-transfer shape/status logs to the ComfyUI console." }, ]; const SAMPLER_KNOB_MAP = Object.fromEntries(SAMPLER_KNOBS.map((k) => [k.name, k])); @@ -287,7 +294,7 @@ chain_timing: ["frame_overlap", "transition_duration", "use_same_seed"], chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "output_guidance", "output_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold"], chain_decode: ["decode_noise_scale", "decode_timestep", "decode_tile_size"], - chain_experimental: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "identity_transfer_enabled", "identity_transfer_phase"], + chain_experimental: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], }; function countChainView(p, id) { @@ -563,6 +570,49 @@ renderKnobList(g, st, CHAIN_VIEW_KNOBS[id]); } + // ── Best-FaceID ArcFace projector: the one identity-transfer field that needs a + // live (server-fetched) choice list, so it can't live in the static SAMPLER_KNOBS combo + // shape. Reuses the same loras folder listing LoraLoaderModelOnly exposes. + let _loraChoices = null; + + async function ensureLoraChoices() { + if (_loraChoices) return _loraChoices; + try { + const spec = await API.nodeSpec("LoraLoaderModelOnly"); + const w = (spec?.inputs || []).find((i) => i.name === "lora_name"); + _loraChoices = (w && w.choices) || []; + } catch (_) { + _loraChoices = []; + } + render(); + return _loraChoices; + } + + function renderChainExperimental(pane, st) { + renderChainKnobsView(pane, st, "chain_experimental", "Experimental", + "Research techniques — off by default; expect quality/overhead trade-offs."); + const si = st.project.sampler_inputs || {}; + if (!si.identity_transfer_enabled) return; + const g = group(pane, "Best-FaceID: ArcFace projector"); + if (!_loraChoices) { + g.append(hintEl("Loading projector list…")); + ensureLoraChoices(); + return; + } + const sel = el("select"); sel.dataset.k = "si-identity_projector"; + const noneOpt = el("option", null, "None"); noneOpt.value = "None"; + if (!si.identity_projector || si.identity_projector === "None") noneOpt.selected = true; + sel.append(noneOpt); + _loraChoices.forEach((c) => { + const o = el("option", null, c); o.value = c; + if (c === si.identity_projector) o.selected = true; + sel.append(o); + }); + sel.onchange = () => S.setSamplerInputNow("identity_projector", sel.value); + g.append(field("ArcFace projector", sel, + "Optional secondary identity channel — the overlap reference tokens above carry the bulk of identity even with this set to None.")); + } + function renderPane(pane, st) { const p = st.project; if (!p) { pane.append(el("div", "pj-meta", "No project open.")); return; } @@ -582,8 +632,7 @@ case "chain_timing": return renderChainTiming(pane, st); case "chain_guidance": return renderChainKnobsView(pane, st, "chain_guidance", "Guidance"); case "chain_decode": return renderChainKnobsView(pane, st, "chain_decode", "Decode"); - case "chain_experimental": return renderChainKnobsView(pane, st, "chain_experimental", "Experimental", - "Research techniques — off by default; expect quality/overhead trade-offs."); + case "chain_experimental": return renderChainExperimental(pane, st); default: return renderOverview(pane, st); } } diff --git a/requirements.txt b/requirements.txt index 7770ec26..4b472b8b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ huggingface_hub safetensors numpy Pillow +insightface +opencv-python-headless diff --git a/samplers.py b/samplers.py index 6b79d5db..9bfc80a9 100644 --- a/samplers.py +++ b/samplers.py @@ -2398,11 +2398,31 @@ class FunPackLTXAVSceneChainSampler: }), "identity_transfer_enabled": ("BOOLEAN", { "default": False, - "tooltip": "EXPERIMENTAL Best-FaceID compatibility: tags Continuity's 'Identity pin' guide (Engine settings) with an extra source-phase RoPE rotation, matching the exact overlap+source_phase conditioning Best-FaceID-style identity LoRAs were trained on. No-op without an identity pin image set. Load the LoRA itself the normal way (Models -> add a LoRA loader onto the model path).", + "tooltip": "EXPERIMENTAL Best-FaceID compatibility: full native port of the overlap+source_phase+ArcFace conditioning Best-FaceID-style identity LoRAs were trained on. Replaces Continuity's 'Identity pin' guide (Engine settings) with separate, non-rendered reference tokens (never blended into frame 0) plus optional ArcFace projector tokens on the text context. No-op without an identity pin image set. Load the LoRA itself the normal way (Models -> add a LoRA loader onto the model path).", }), - "identity_transfer_phase": ("FLOAT", { - "default": 2.0, "min": 0.0, "max": 8.0, "step": 0.5, - "tooltip": "Source-phase segment id for the identity-pin guide tokens (ltx-trainer's overlap+source_phase convention used 2). 0 disables the rotation while leaving the guide itself active.", + "identity_projector": (cls._identity_projector_choices(), { + "default": "None", + "tooltip": "ArcFace projector .safetensors (from models/loras). 'None' = overlap only (the projector is a weak secondary channel; the overlap latent carries the bulk of identity).", + }), + "source_id": ("FLOAT", { + "default": 2.0, "min": 0.0, "max": 8.0, "step": 1.0, + "tooltip": "source_phase segment id for the overlap reference tokens (ltx-trainer's overlap+source_phase convention used 2). 0 disables the RoPE rotation while leaving the overlap tokens active.", + }), + "phase_scale": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 4.0, "step": 0.1, + "tooltip": "Multiplier on source_id before the RoPE rotation.", + }), + "id_strength": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 50.0, "step": 0.5, + "tooltip": "Multiplies the ArcFace projector tokens (only when identity_projector is set). Weak channel; push high (5-20) to test, very high may add artifacts.", + }), + "arcface_mode": (["auto_adjust", "as_is", "disable"], { + "default": "auto_adjust", + "tooltip": "auto_adjust: retry face detection with zoom-out/upscale, skip projector tokens if none found. as_is: detect on the image only. disable: skip ArcFace, use only the overlap latent.", + }), + "debug_log": ("BOOLEAN", { + "default": False, + "tooltip": "Print per-scene identity-transfer shape/status logs to the console.", }), }, "hidden": { @@ -2420,6 +2440,14 @@ class FunPackLTXAVSceneChainSampler: "Use with FunPack Studio split-by-transitions output." ) + @classmethod + def _identity_projector_choices(cls): + try: + import folder_paths + return ["None"] + folder_paths.get_filename_list("loras") + except Exception: + return ["None"] + def _is_nested(self, samples): return bool(getattr(samples, "is_nested", False)) @@ -3576,34 +3604,8 @@ class FunPackLTXAVSceneChainSampler: return 0, True return 1 + (at - 1) * time_scale, False - def _tag_last_guide_entry(self, cond, key, value): - """Best-effort: stamp `key=value` onto the guide_attention_entry _append_guide_attention_entry - just added (the last one), without touching any earlier entry. Used to mark which guide - (e.g. the Best-FaceID identity pin) is eligible for the source-phase RoPE patch — see - _install_identity_phase. No-op (returns cond unchanged) if node_helpers isn't importable - or no entries are present; never blocks generation.""" - try: - import node_helpers - except ImportError: - return cond - try: - for t in cond: - entries = t[1].get("guide_attention_entries") - if entries: - tagged = [*entries[:-1], {**entries[-1], key: value}] - return node_helpers.conditioning_set_values(cond, {"guide_attention_entries": tagged}) - except Exception: - pass - return cond - - def _append_guide_latent(self, chunk, guide_frame, apply_at, strength, positive, negative, vae, - phase_tag=0.0): - """Append one guide latent frame with LTX guide attention at apply_at. - - `phase_tag`: when non-zero, marks this guide's entry with `funpack_identity_phase` so - _install_identity_phase can find its exact token range and apply the source-phase RoPE - rotation the Best-FaceID LoRA was trained against — see [[project_identity_transfer_native]]. - """ + def _append_guide_latent(self, chunk, guide_frame, apply_at, strength, positive, negative, vae): + """Append one guide latent frame with LTX guide attention at apply_at.""" try: from comfy_extras.nodes_lt import LTXVAddGuide, _append_guide_attention_entry except ImportError: @@ -3629,9 +3631,6 @@ class FunPackLTXAVSceneChainSampler: positive, negative = _append_guide_attention_entry( positive, negative, pre_filter_count, guide_latent_shape, strength=float(strength) ) - if phase_tag: - positive = self._tag_last_guide_entry(positive, "funpack_identity_phase", float(phase_tag)) - negative = self._tag_last_guide_entry(negative, "funpack_identity_phase", float(phase_tag)) result = self._clone_latent(chunk) tensors = self._latent_tensors(result) @@ -3845,7 +3844,7 @@ class FunPackLTXAVSceneChainSampler: return None def _append_media_guide_at(self, chunk, filename, frame_idx, apply_at, strength, - positive, negative, vae, phase_tag=0.0): + positive, negative, vae): chunk_tensors = self._latent_tensors(chunk) if not chunk_tensors: return chunk, positive, negative, 0, 0 @@ -3856,15 +3855,15 @@ class FunPackLTXAVSceneChainSampler: at = self._resolve_frame_index(total, frame_idx) guide_slice = self._time_slice(guide_frame, at, at + 1) return self._append_guide_latent( - chunk, guide_slice, apply_at, strength, positive, negative, vae, phase_tag=phase_tag, + chunk, guide_slice, apply_at, strength, positive, negative, vae, ) + (0,) def _apply_configured_guides(self, chunk, scene_index, guide_list, latent_template, scene_outputs, scene_media_by_ref, positive, negative, vae, - identity_transfer_enabled=False, identity_transfer_phase=2.0): + identity_transfer_enabled=False): head_crop = 0 tail_crop = 0 - identity_phase_applied = False + identity_ref_filename = None for g in guide_list or []: if not g or not g.get("enabled", True): continue @@ -3894,20 +3893,21 @@ class FunPackLTXAVSceneChainSampler: fn = (scene_media_by_ref or {}).get(ref) if ref else None if not fn: continue - # The project's Best-FaceID reference face (Continuity's "Identity pin") is the - # ONLY guide eligible for the source-phase RoPE tag — never prior-scene/mid-scene - # guides, which serve motion continuity rather than a trained-LoRA identity lock. - phase_tag = identity_transfer_phase if (identity_transfer_enabled and g.get("identity_pin")) else 0.0 + if identity_transfer_enabled and g.get("identity_pin"): + # Best-FaceID full port takes this entry over entirely: the LoRA was + # trained on separate overlap tokens + ArcFace projector tokens, not a + # blended frame-0 keyframe, so skip the plain guide-attention append and + # hand the reference filename back for _install_identity_overlap. + identity_ref_filename = fn + continue chunk, positive, negative, head, tail = self._append_media_guide_at( - chunk, fn, frame_idx, apply_at, strength, positive, negative, vae, phase_tag=phase_tag, + chunk, fn, frame_idx, apply_at, strength, positive, negative, vae, ) - if phase_tag: - identity_phase_applied = True else: continue head_crop += head tail_crop += tail - return chunk, positive, negative, head_crop, tail_crop, identity_phase_applied + return chunk, positive, negative, head_crop, tail_crop, identity_ref_filename def _append_mid_scene_guide(self, chunk, previous_output, positive, negative, vae, strength): """Append the middle frame of the previous scene as a guide for the current chunk @@ -4111,127 +4111,220 @@ class FunPackLTXAVSceneChainSampler: pass # --------------------------------------------------------------------------- - # identity_transfer: source-phase RoPE tag for the Best-FaceID-style identity pin. + # identity_transfer: native port of ComfyUI-BFSNodes' LTX Identity Transfer + # (LTXIdentityOverlapConditioning) — see identity_transfer.py for the ArcFace/projector/ + # RoPE-rotation pieces. The reference face is injected as SEPARATE tokens appended after + # the target's video tokens (an "overlap" placement, sharing the frame-0 pixel-coord grid), + # given a clean (timestep 0) denoise state, tagged with a source-phase RoPE rotation, and + # trimmed back off before unpatchify — never blended into a real rendered frame, unlike + # FunPack's guide-attention path. Optional ArcFace IdentityProjector tokens are appended to + # the text context on top. # - # FunPack's guide-attention path (LTXVAddGuide.add_keyframe_index + causal_fix, see - # _append_guide_latent) already places a reference image on the frame-0 grid with a - # tunable attention strength — structurally the same "overlap" placement an identity - # LoRA like Best-FaceID expects. The one piece that path doesn't do is the LoRA's - # trained-time "source_phase" tag: an extra per-token RoPE rotation that marks those - # tokens as coming from a reference image, not the real timeline. This patches that in, - # scoped ONLY to the specific guide entry marked funpack_identity_phase (see - # _tag_last_guide_entry) — never to co-active guides (mid_scene_guide, carry, prior- - # scene continuity) that happen to share the same keyframe_idxs tensor. - # - # Two bound methods are wrapped on model.model.diffusion_model: - # _prepare_timestep (receives resolved_guide_entries / num_guide_tokens via - # kwargs already computed by the model's own guide-token - # accounting) -> locates the tagged entry's exact token - # range and stashes it on self for the very next call. - # _prepare_positional_embeddings (runs right after, same forward pass, no kwargs) -> - # rotates just that token range's RoPE freqs. + # Four bound methods are wrapped on model.model.diffusion_model, matching LTXBaseModel's own + # _forward pipeline: _process_input (append ref tokens) -> _prepare_timestep (clean timestep + # for ref tokens) -> _prepare_positional_embeddings (source-phase RoPE rotation) -> + # _process_output (trim ref tokens before unpatchify). # Idempotent + tag/strip per scene (same discipline as _install_v2a_scale / # _strip_funpack_scene_wrappers) so an interrupt mid-scene can never leave a stale patch # steering a later, unrelated run — see [[project_hook_leak_bug]]. # --------------------------------------------------------------------------- - _IDENTITY_PHASE_TAG = "_funpack_identity_phase_patched" + _IDENTITY_OVERLAP_TAG = "_funpack_identity_overlap_patched" - def _install_identity_phase(self, model): + def _install_identity_overlap(self, model, ref_latent, seg_value): + try: + from .identity_transfer import rotate_overlap_freqs as _rotate_overlap_freqs + except ImportError: + from identity_transfer import rotate_overlap_freqs as _rotate_overlap_freqs try: ltxv = model.model.diffusion_model except Exception: return None - if getattr(ltxv, self._IDENTITY_PHASE_TAG, False): + if getattr(ltxv, self._IDENTITY_OVERLAP_TAG, False): return None # already installed (idempotent across scenes/runs) + orig_process_input = ltxv._process_input orig_prepare_timestep = ltxv._prepare_timestep orig_prepare_pe = ltxv._prepare_positional_embeddings + orig_process_output = ltxv._process_output + + def process_input(self_ltxv, x, keyframe_idxs, denoise_mask, **kw): + out = orig_process_input(x, keyframe_idxs, denoise_mask, **kw) + self_ltxv._funpack_id_ref_len = 0 + try: + from comfy.ldm.lightricks.model import latent_to_pixel_coords + xx, pix, add = out + is_av = isinstance(xx, (list, tuple)) + vx = xx[0] if is_av else xx + vco = pix[0] if is_av else pix + rt, rlc = self_ltxv.patchifier.patchify(ref_latent.to(dtype=vx.dtype, device=vx.device)) + rpc = latent_to_pixel_coords(latent_coords=rlc, scale_factors=self_ltxv.vae_scale_factors, + causal_fix=self_ltxv.causal_temporal_positioning) + rt = self_ltxv.patchify_proj(rt) + if rt.shape[0] != vx.shape[0]: + rt = rt.expand(vx.shape[0], -1, -1) + if rpc.shape[0] != vco.shape[0]: + rpc = rpc.expand(vco.shape[0], *([-1] * (rpc.dim() - 1))) + ref_len = rt.shape[1] + self_ltxv._funpack_id_target_len = vx.shape[1] + vx = torch.cat([vx, rt], dim=1) + vco = torch.cat([vco, rpc.to(vco)], dim=2) + self_ltxv._funpack_id_ref_len = ref_len + if is_av: + xx = [vx, xx[1]]; pix = [vco, pix[1]] + else: + xx, pix = vx, vco + return xx, pix, add + except Exception: + self_ltxv._funpack_id_ref_len = 0 + return out def prepare_timestep(self_ltxv, timestep, batch_size, hidden_dtype, **kwargs): - self_ltxv._funpack_identity_range = None - self_ltxv._funpack_identity_phase_value = 0.0 - try: - entries = kwargs.get("resolved_guide_entries") - num_guide_tokens = kwargs.get("num_guide_tokens", 0) - if entries and num_guide_tokens: - offset = 0 - for entry in entries: - phase = entry.get("funpack_identity_phase") - surviving = int(entry.get("surviving_count", 0)) - if phase: - self_ltxv._funpack_identity_range = (offset, surviving, int(num_guide_tokens)) - self_ltxv._funpack_identity_phase_value = float(phase) - break - offset += surviving - except Exception: - self_ltxv._funpack_identity_range = None + ref_len = getattr(self_ltxv, "_funpack_id_ref_len", 0) + if ref_len: + target_len = getattr(self_ltxv, "_funpack_id_target_len", None) + if timestep.dim() <= 1 and target_len is not None: + timestep = timestep.view(-1, 1).expand(batch_size, target_len).contiguous() + if timestep.dim() >= 2: + ref_ts = torch.zeros(batch_size, ref_len, *timestep.shape[2:], + device=timestep.device, dtype=timestep.dtype) + timestep = torch.cat([timestep, ref_ts], dim=1) return orig_prepare_timestep(timestep, batch_size, hidden_dtype, **kwargs) def prepare_pe(self_ltxv, pixel_coords, frame_rate, x_dtype): pe = orig_prepare_pe(pixel_coords, frame_rate, x_dtype) - rng = getattr(self_ltxv, "_funpack_identity_range", None) - phase = getattr(self_ltxv, "_funpack_identity_phase_value", 0.0) - if not rng or not phase: + ref_len = getattr(self_ltxv, "_funpack_id_ref_len", 0) + if not ref_len or not seg_value: return pe try: - offset, length, num_guide_tokens = rng # AV models return [(v_pe, av_cross_video), (a_pe, av_cross_audio)]; plain - # video models return v_pe directly. Only the video PE carries guide tokens. + # video models return v_pe directly. Only the video PE carries the ref tokens. if isinstance(pe, list) and len(pe) and isinstance(pe[0], (list, tuple)) and isinstance(pe[0][0], (list, tuple)): v_pe, cross_v = pe[0][0], pe[0][1] - v_pe = self._rotate_identity_phase(v_pe, offset, length, num_guide_tokens, phase) + v_pe = _rotate_overlap_freqs(v_pe, ref_len, seg_value) return [(v_pe, cross_v), pe[1]] - return self._rotate_identity_phase(pe, offset, length, num_guide_tokens, phase) + return _rotate_overlap_freqs(pe, ref_len, seg_value) except Exception: return pe + def process_output(self_ltxv, x, embedded_timestep, keyframe_idxs, **kw): + ref_len = getattr(self_ltxv, "_funpack_id_ref_len", 0) + if ref_len: + try: + from comfy.ldm.lightricks.av_model import CompressedTimestep + import copy + if isinstance(x, (list, tuple)): + x = [x[0][:, :x[0].shape[1] - ref_len], *x[1:]] + et_list = list(embedded_timestep) if isinstance(embedded_timestep, (list, tuple)) else [embedded_timestep] + v_et = et_list[0] + if isinstance(v_et, CompressedTimestep): + ppf = max(1, getattr(v_et, "patches_per_frame", 1) or 1) + n_ref_frames = max(1, ref_len // ppf) + v_et2 = copy.copy(v_et) + v_et2.data = v_et.data[:, : v_et.num_frames - n_ref_frames].contiguous() + v_et2.num_frames = v_et.num_frames - n_ref_frames + et_list[0] = v_et2 + elif hasattr(v_et, "shape") and v_et.dim() >= 2 and v_et.shape[1] > 1: + et_list[0] = v_et[:, : v_et.shape[1] - ref_len] + embedded_timestep = et_list + else: + x = x[:, :x.shape[1] - ref_len] + if hasattr(embedded_timestep, "shape") and embedded_timestep.dim() >= 2 and embedded_timestep.shape[1] > 1: + embedded_timestep = embedded_timestep[:, : embedded_timestep.shape[1] - ref_len] + except Exception: + pass + return orig_process_output(x, embedded_timestep, keyframe_idxs, **kw) + + ltxv._process_input = types.MethodType(process_input, ltxv) ltxv._prepare_timestep = types.MethodType(prepare_timestep, ltxv) ltxv._prepare_positional_embeddings = types.MethodType(prepare_pe, ltxv) - setattr(ltxv, self._IDENTITY_PHASE_TAG, True) - return (ltxv, orig_prepare_timestep, orig_prepare_pe) + ltxv._process_output = types.MethodType(process_output, ltxv) + setattr(ltxv, self._IDENTITY_OVERLAP_TAG, True) + return (ltxv, orig_process_input, orig_prepare_timestep, orig_prepare_pe, orig_process_output) - def _strip_identity_phase(self, handle): + def _strip_identity_overlap(self, handle): if not handle: return - ltxv, orig_prepare_timestep, orig_prepare_pe = handle + ltxv, orig_process_input, orig_prepare_timestep, orig_prepare_pe, orig_process_output = handle try: + ltxv._process_input = orig_process_input ltxv._prepare_timestep = orig_prepare_timestep ltxv._prepare_positional_embeddings = orig_prepare_pe - setattr(ltxv, self._IDENTITY_PHASE_TAG, False) - for attr in ("_funpack_identity_range", "_funpack_identity_phase_value"): + ltxv._process_output = orig_process_output + setattr(ltxv, self._IDENTITY_OVERLAP_TAG, False) + for attr in ("_funpack_id_ref_len", "_funpack_id_target_len"): if hasattr(ltxv, attr): delattr(ltxv, attr) except Exception: pass - def _rotate_identity_phase(self, pe, offset, length, num_guide_tokens, phase, theta=10000.0): - """Givens-rotate the RoPE freqs of exactly [guide_start+offset, guide_start+offset+length) - (a single guide entry's tokens, identified upstream — never the whole guide-token tail) - by a phase derived from `phase` — the same source_phase tag the LoRA saw at train time. - pe = (cos, sin, split_mode) with cos/sin shaped [..., T, L]. Returns a new pe tuple.""" - if length <= 0 or not phase: - return pe - cos, sin = pe[0], pe[1] - rest = tuple(pe[2:]) - T = cos.shape[-2] - guide_start = T - int(num_guide_tokens) - start = guide_start + int(offset) - end = start + int(length) - if start < 0 or end > T or start >= end: - return pe - L = cos.shape[-1] - d = torch.arange(L, device=cos.device, dtype=torch.float32) - rate = theta ** (-d / float(L)) - rot = (float(phase) * rate) - pc = rot.cos().to(cos.dtype); ps = rot.sin().to(sin.dtype) - idx = [slice(None)] * cos.dim() - idx[-2] = slice(start, end) - idx = tuple(idx) - c0, s0 = cos[idx], sin[idx] - cos = cos.clone(); sin = sin.clone() - cos[idx] = c0 * pc - s0 * ps - sin[idx] = s0 * pc + c0 * ps - return (cos, sin, *rest) + def _resolve_identity_overlap(self, state, filename, vae, chunk, identity_projector, source_id, + phase_scale, id_strength, arcface_mode, debug_log): + """Lazily encode the reference face + (optional) ArcFace projector tokens once per + run() call and cache them in `state` (a dict owned by the caller). Cheap on every + later scene — just returns the cached tuple. Returns + (ref_latent_or_None, seg_value, pos_tokens_or_None, neg_tokens_or_None).""" + if state.get("ready"): + return state["ref_latent"], state["seg_value"], state["pos_tokens"], state["neg_tokens"] + state.update(ready=True, ref_latent=None, seg_value=float(source_id) * float(phase_scale), + pos_tokens=None, neg_tokens=None) + image = self._load_image_tensor(filename) + if image is None: + print(f"[FunPackSceneChain] identity_transfer: couldn't load reference image '{filename}' — skipped.") + return None, 0.0, None, None + chunk_tensors = self._latent_tensors(chunk) + if not chunk_tensors: + return None, 0.0, None, None + try: + import comfy.utils + scale_factors = getattr(vae, "downscale_index_formula", [8, 8, 8]) + _, w_sf, h_sf = scale_factors + _, _, _, lat_h, lat_w = chunk_tensors[0].shape + ref_px = comfy.utils.common_upscale( + image.movedim(-1, 1), lat_w * w_sf, lat_h * h_sf, "bilinear", "center", + ).movedim(1, -1)[:1, :, :, :3] + state["ref_latent"] = vae.encode(ref_px) + if debug_log: + print(f"[FunPackSceneChain] identity_transfer: overlap ref latent " + f"{list(state['ref_latent'].shape)}, seg={state['seg_value']}") + except Exception as e: + print(f"[FunPackSceneChain] identity_transfer: overlap ref-latent encode failed ({e}) — skipped.") + return None, 0.0, None, None + + if identity_projector in (None, "", "None"): + return state["ref_latent"], state["seg_value"], None, None + try: + try: + from .identity_transfer import arcface_embed, load_identity_projector + except ImportError: + from identity_transfer import arcface_embed, load_identity_projector + emb = arcface_embed(image, mode=arcface_mode) + if emb is None: + print(f"[FunPackSceneChain] identity_transfer: ArcFace mode={arcface_mode} found no " + "face — projector tokens skipped (overlap only).") + return state["ref_latent"], state["seg_value"], None, None + import folder_paths + path = folder_paths.get_full_path("loras", identity_projector) or identity_projector + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + projector = load_identity_projector(path, device) + emb = emb.to(device=device, dtype=torch.float32).unsqueeze(0) + with torch.no_grad(): + state["pos_tokens"] = projector(emb) * float(id_strength) + state["neg_tokens"] = projector(torch.zeros(1, projector.in_dim, device=device)) + if debug_log: + print(f"[FunPackSceneChain] identity_transfer: ArcFace projector '{identity_projector}' " + f"OK, id_strength={id_strength}") + except Exception as e: + print(f"[FunPackSceneChain] identity_transfer: ArcFace/projector failed ({e}) — overlap tokens only " + "(insightface installed? `pip install insightface`).") + return state["ref_latent"], state["seg_value"], None, None + return state["ref_latent"], state["seg_value"], state["pos_tokens"], state["neg_tokens"] + + def _append_identity_context_tokens(self, positive, negative, pos_tokens, neg_tokens): + try: + from .identity_transfer import append_context_tokens + except ImportError: + from identity_transfer import append_context_tokens + return append_context_tokens(positive, pos_tokens), append_context_tokens(negative, neg_tokens) def _bounded_attention_region_mask(self, t, h, w, device): """[T*H*W] region id per video token: 0 = left half (by width), 1 = right half. @@ -4359,7 +4452,8 @@ class FunPackLTXAVSceneChainSampler: output_guidance=False, output_guidance_strength=0.02, dynashift=False, dynashift_strength=0.3, dynashift_threshold=0.6, alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975, - identity_transfer_enabled=False, identity_transfer_phase=2.0, + identity_transfer_enabled=False, identity_projector="None", source_id=2.0, + phase_scale=1.0, id_strength=1.0, arcface_mode="auto_adjust", debug_log=False, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -4537,6 +4631,7 @@ class FunPackLTXAVSceneChainSampler: scene_media_by_ref = self._parse_scene_anchors(funpack_scene_media_refs) scene_runs: list = [] joyai_bank = _JoyAIMemoryBank(joyai_memory_size, joyai_fix_frames) if joyai_memory else None + _identity_overlap_state: dict = {} for scene_index, scene_cond in enumerate(scene_conditionings): scene_positive = [scene_cond] @@ -4551,7 +4646,7 @@ class FunPackLTXAVSceneChainSampler: soft_carried = 0 guide_tail = 0 audio_tail = 0 - identity_phase_applied = False + identity_ref_filename = None run_mechanisms: list = [] anchor_meta = (scene_anchors or {}).get(str(scene_index)) if output is None: @@ -4561,11 +4656,10 @@ class FunPackLTXAVSceneChainSampler: custom_guides = per_scene_guides[scene_index] if custom_guides: run_mechanisms.append("custom_guide_stack") - chunk, scene_positive, scene_negative, carried, guide_tail, identity_phase_applied = self._apply_configured_guides( + chunk, scene_positive, scene_negative, carried, guide_tail, identity_ref_filename = self._apply_configured_guides( chunk, scene_index, custom_guides, latent_template, scene_outputs, scene_media_by_ref, scene_positive, scene_negative, vae, identity_transfer_enabled=identity_transfer_enabled, - identity_transfer_phase=identity_transfer_phase, ) carried_guide_frames = max(carried_guide_frames, carried) elif anchor_meta: @@ -4605,11 +4699,10 @@ class FunPackLTXAVSceneChainSampler: custom_guides = per_scene_guides[scene_index] if custom_guides: run_mechanisms.append("custom_guide_stack") - chunk, scene_positive, scene_negative, carried, guide_tail, identity_phase_applied = self._apply_configured_guides( + chunk, scene_positive, scene_negative, carried, guide_tail, identity_ref_filename = self._apply_configured_guides( chunk, scene_index, custom_guides, latent_template, scene_outputs, scene_media_by_ref, scene_positive, scene_negative, vae, identity_transfer_enabled=identity_transfer_enabled, - identity_transfer_phase=identity_transfer_phase, ) carried_guide_frames = max(carried_guide_frames, carried) elif carry_i2v_guides: @@ -4644,7 +4737,7 @@ class FunPackLTXAVSceneChainSampler: # mode as the block-hook leak; see _strip_funpack_scene_wrappers). _scene_base_wrapper = model.model_options.get("model_function_wrapper") _v2a_handles = [] - _identity_phase_handle = None + _identity_overlap_handle = None try: if embed_guidance and _value_fn is not None and _value_fn.is_ready(): run_mechanisms.append("embed_guidance_vf_ascend") @@ -4696,10 +4789,20 @@ class FunPackLTXAVSceneChainSampler: _v2a_handles = self._install_v2a_scale(model, v2a_grad_scale) if _v2a_handles: run_mechanisms.append(f"v2a_grad_scale({v2a_grad_scale})") - if identity_transfer_enabled and identity_phase_applied: - _identity_phase_handle = self._install_identity_phase(model) - if _identity_phase_handle: - run_mechanisms.append(f"identity_transfer_phase({identity_transfer_phase})") + if identity_transfer_enabled and identity_ref_filename: + _id_ref_latent, _id_seg_value, _id_pos_tokens, _id_neg_tokens = self._resolve_identity_overlap( + _identity_overlap_state, identity_ref_filename, vae, chunk, + identity_projector, source_id, phase_scale, id_strength, arcface_mode, debug_log, + ) + if _id_ref_latent is not None: + _identity_overlap_handle = self._install_identity_overlap(model, _id_ref_latent, _id_seg_value) + if _identity_overlap_handle: + run_mechanisms.append(f"identity_transfer_overlap(seg={_id_seg_value})") + if _id_pos_tokens is not None: + scene_positive, scene_negative = self._append_identity_context_tokens( + scene_positive, scene_negative, _id_pos_tokens, _id_neg_tokens, + ) + run_mechanisms.append(f"identity_transfer_arcface(id_strength={id_strength})") _t_sample0 = _time.perf_counter() sampled = self._sample_chunk( model, sampler, sigmas, scene_seed, cfg, scene_positive, scene_negative, chunk, @@ -4713,7 +4816,7 @@ class FunPackLTXAVSceneChainSampler: _phase_sampling += _scene_sample_s finally: self._remove_v2a_scale(_v2a_handles) - self._strip_identity_phase(_identity_phase_handle) + self._strip_identity_overlap(_identity_overlap_handle) if model.model_options.get("model_function_wrapper") is not _scene_base_wrapper: if _scene_base_wrapper is not None: model.model_options["model_function_wrapper"] = _scene_base_wrapper diff --git a/tests/test_identity_transfer.py b/tests/test_identity_transfer.py index 6293a76f..3af38d48 100644 --- a/tests/test_identity_transfer.py +++ b/tests/test_identity_transfer.py @@ -1,15 +1,16 @@ -"""Unit tests for the native Best-FaceID identity_transfer port. +"""Unit tests for the native Best-FaceID identity_transfer port (overlap tokens + ArcFace +projector), a port of ComfyUI-BFSNodes' LTX Identity Transfer. -Covers the three pieces added to FunPackLTXAVSceneChainSampler: - - _tag_last_guide_entry: best-effort marks the most-recently-appended - guide_attention_entry, never touching earlier entries. - - _apply_configured_guides: routes the phase tag ONLY to the guide entry the - editor marked identity_pin (never prior-scene/mid-scene/template guides that - happen to share the same call). - - _install_identity_phase / _strip_identity_phase / _rotate_identity_phase: the - tag+strip RoPE patch that reproduces the LoRA's trained-time source_phase - tag, scoped to the exact token range of the tagged guide (not the whole - guide-token tail, which may include co-active unrelated guides). +Covers: + - identity_transfer.rotate_overlap_freqs / append_context_tokens: pure-tensor pieces, + no torch model needed. + - FunPackLTXAVSceneChainSampler._apply_configured_guides: routes the identity_pin entry + to the new mechanism (returns its filename, skips the keyframe-blend append) only when + identity_transfer_enabled; unaffected guides append normally either way. + - _install_identity_overlap / _strip_identity_overlap: tag+strip idempotency, ref-token + append (_process_input), clean timestep for ref tokens (_prepare_timestep), source-phase + RoPE rotation (_prepare_positional_embeddings), and ref-token trim (_process_output). + - _resolve_identity_overlap: lazy, memoized once per run() call. """ import sys import types @@ -23,217 +24,298 @@ for _name in ( "comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling", "comfy.model_sampling", "comfy.nested_tensor", "comfy.sample", "comfy.samplers", "comfy.utils", + "comfy.ldm", "comfy.ldm.lightricks", "comfy.ldm.lightricks.model", "comfy.ldm.lightricks.av_model", ): sys.modules.setdefault(_name, types.ModuleType(_name)) sys.modules["comfy.nested_tensor"].NestedTensor = object +sys.modules["comfy.ldm.lightricks.model"].latent_to_pixel_coords = ( + lambda latent_coords, scale_factors, causal_fix: latent_coords) + + +class _CompressedTimestep: + def __init__(self, data, num_frames, patches_per_frame=1): + self.data = data + self.num_frames = num_frames + self.patches_per_frame = patches_per_frame + + +sys.modules["comfy.ldm.lightricks.av_model"].CompressedTimestep = _CompressedTimestep import samplers # noqa: E402 +import identity_transfer as idt # noqa: E402 def _sampler(): return samplers.FunPackLTXAVSceneChainSampler() -def _install_node_helpers(): - """node_helpers.conditioning_set_values, the same real implementation shape - _append_guide_attention_entry itself relies on (copy metadata dict, overwrite - the given keys) — reimplemented here so the test runs without a full ComfyUI env.""" - mod = types.ModuleType("node_helpers") +# ── identity_transfer.rotate_overlap_freqs ───────────────────────────────────── - def conditioning_set_values(conditioning, values=None, append=False): - values = values or {} - out = [] - for t in conditioning: - n = [t[0], t[1].copy()] - for k, v in values.items(): - n[1][k] = v - out.append(n) - return out - - mod.conditioning_set_values = conditioning_set_values - sys.modules["node_helpers"] = mod - return mod +def test_rotate_overlap_freqs_noop_for_zero_ref_len_or_seg(): + cos = torch.arange(20, dtype=torch.float32).reshape(5, 4) + sin = -cos.clone() + pe = (cos, sin, False) + assert idt.rotate_overlap_freqs(pe, 0, 2.0) is pe + assert idt.rotate_overlap_freqs(pe, 2, 0.0) is pe -def _cond_with_entries(entries): - return [[torch.zeros(1, 1, 4), {"guide_attention_entries": list(entries)}]] +def test_rotate_overlap_freqs_rotates_only_last_ref_len_tokens(): + T, L = 10, 4 + cos = torch.arange(T * L, dtype=torch.float32).reshape(T, L) + sin = torch.arange(T * L, dtype=torch.float32).reshape(T, L) * -1.0 + pe = (cos, sin, False) + ref_len = 3 + out_cos, out_sin, split = idt.rotate_overlap_freqs(pe, ref_len, 2.0) + untouched = list(range(0, T - ref_len)) + for i in untouched: + assert torch.equal(out_cos[i], cos[i]) + assert torch.equal(out_sin[i], sin[i]) + rotated = list(range(T - ref_len, T)) + changed = any(not torch.equal(out_cos[i], cos[i]) for i in rotated) + assert changed + assert split is False -# ── _tag_last_guide_entry ────────────────────────────────────────────────────── +# ── identity_transfer.append_context_tokens ──────────────────────────────────── -def test_tag_last_guide_entry_marks_only_newest_entry(): - _install_node_helpers() - s = _sampler() - cond = _cond_with_entries([{"pre_filter_count": 4}, {"pre_filter_count": 8}]) - tagged = s._tag_last_guide_entry(cond, "funpack_identity_phase", 2.0) - entries = tagged[0][1]["guide_attention_entries"] - assert entries[0].get("funpack_identity_phase") is None - assert entries[1]["funpack_identity_phase"] == 2.0 - # Original untouched (conditioning_set_values copies, doesn't mutate in place). - assert "funpack_identity_phase" not in cond[0][1]["guide_attention_entries"][1] +def test_append_context_tokens_pads_and_extends_mask(): + ce = torch.zeros(2, 5, 8) # batch=2, tokens=5, dim=8 + cond = [[ce, {"attention_mask": torch.ones(2, 5)}]] + tokens = torch.ones(1, 3, 4) # batch=1 (broadcast), 3 new tokens, smaller dim + out = idt.append_context_tokens(cond, tokens) + new_ce, new_extra = out[0] + assert new_ce.shape == (2, 8, 8) # 5+3 tokens, dim padded back to 8 + assert torch.equal(new_ce[:, 5:, :4], torch.ones(2, 3, 4)) + assert torch.equal(new_ce[:, 5:, 4:], torch.zeros(2, 3, 4)) + assert new_extra["attention_mask"].shape == (2, 8) + assert torch.equal(new_extra["attention_mask"][:, 5:], torch.ones(2, 3)) + # Original conditioning list untouched. + assert ce.shape == (2, 5, 8) -def test_tag_last_guide_entry_noop_without_entries(): - _install_node_helpers() - s = _sampler() - cond = [[torch.zeros(1, 1, 4), {}]] - out = s._tag_last_guide_entry(cond, "funpack_identity_phase", 2.0) - assert out is cond # unchanged object, nothing to tag +def test_append_context_tokens_truncates_wider_tokens(): + ce = torch.zeros(1, 2, 4) + cond = [[ce, {}]] + tokens = torch.ones(1, 1, 8) # wider than context dim + out = idt.append_context_tokens(cond, tokens) + new_ce, _ = out[0] + assert new_ce.shape == (1, 3, 4) -def test_tag_last_guide_entry_noop_without_node_helpers(): - sys.modules.pop("node_helpers", None) - s = _sampler() - cond = _cond_with_entries([{"pre_filter_count": 4}]) - out = s._tag_last_guide_entry(cond, "funpack_identity_phase", 2.0) - assert out is cond +# ── _apply_configured_guides identity-pin routing ────────────────────────────── - -# ── _apply_configured_guides phase-tag routing ───────────────────────────────── - -def test_phase_tag_applies_only_to_identity_pin_entry(monkeypatch): +def test_identity_pin_routes_to_overlap_when_enabled(monkeypatch): s = _sampler() calls = [] - - def fake_append_media(chunk, filename, frame_idx, apply_at, strength, positive, negative, vae, - phase_tag=0.0): - calls.append((filename, phase_tag)) - return chunk, positive, negative, 0, 0 - - monkeypatch.setattr(s, "_append_media_guide_at", fake_append_media) + monkeypatch.setattr(s, "_append_media_guide_at", + lambda chunk, filename, *a, **kw: (calls.append(filename), chunk, [], [], 0, 0)[1:]) guide_list = [ {"enabled": True, "source": "image", "media_ref": "pin", "identity_pin": True, "strength": 0.35}, {"enabled": True, "source": "image", "media_ref": "prior", "strength": 0.35}, # NOT the pin ] scene_media_by_ref = {"pin": "pin.png", "prior": "prior.png"} - chunk, pos, neg, head, tail, phase_applied = s._apply_configured_guides( + chunk, pos, neg, head, tail, identity_ref_filename = s._apply_configured_guides( {}, 1, guide_list, {}, [], scene_media_by_ref, [], [], object(), - identity_transfer_enabled=True, identity_transfer_phase=3.0, + identity_transfer_enabled=True, ) - assert calls == [("pin.png", 3.0), ("prior.png", 0.0)] - assert phase_applied is True + assert calls == ["prior.png"] # pin entry never went through the plain append + assert identity_ref_filename == "pin.png" -def test_phase_tag_off_when_identity_transfer_disabled(monkeypatch): +def test_identity_pin_appends_normally_when_disabled(monkeypatch): s = _sampler() calls = [] monkeypatch.setattr(s, "_append_media_guide_at", - lambda chunk, filename, *a, phase_tag=0.0, **kw: (calls.append(phase_tag), chunk, [], [], 0, 0)[1:]) + lambda chunk, filename, *a, **kw: (calls.append(filename), chunk, [], [], 0, 0)[1:]) guide_list = [{"enabled": True, "source": "image", "media_ref": "pin", "identity_pin": True, "strength": 0.35}] - _, _, _, _, _, phase_applied = s._apply_configured_guides( + _, _, _, _, _, identity_ref_filename = s._apply_configured_guides( {}, 1, guide_list, {}, [], {"pin": "pin.png"}, [], [], object(), - identity_transfer_enabled=False, identity_transfer_phase=2.0, + identity_transfer_enabled=False, ) - assert calls == [0.0] - assert phase_applied is False + assert calls == ["pin.png"] + assert identity_ref_filename is None -# ── _install_identity_phase / _strip_identity_phase / _rotate_identity_phase ── +# ── _install_identity_overlap / _strip_identity_overlap ──────────────────────── + +class _FakePatchifier: + def patchify(self, latent): + # latent: [1, C, 1, H, W] -> pretend 2 ref tokens: rt=[B,N,C], rlc=[B,3,N] (coords). + b = latent.shape[0] + return torch.zeros(b, 2, 8), torch.zeros(b, 3, 2) + class _FakeLTXV: + def __init__(self): + self.patchifier = _FakePatchifier() + self.vae_scale_factors = [8, 8, 8] + self.causal_temporal_positioning = False + + def _process_input(self, x, keyframe_idxs, denoise_mask, **kw): + # pixel_coords mirrors x's structure: [B, 3, T] (or a [video, audio] pair for AV). + if isinstance(x, (list, tuple)): + pix = [torch.zeros(t.shape[0], 3, t.shape[1]) for t in x] + else: + pix = torch.zeros(x.shape[0], 3, x.shape[1]) + return x, pix, {} + + def patchify_proj(self, rt): + return rt # identity, dim already matches x for the test + def _prepare_timestep(self, timestep, batch_size, hidden_dtype, **kwargs): - return ("timestep", batch_size, hidden_dtype) + return ("timestep", timestep, batch_size, hidden_dtype) def _prepare_positional_embeddings(self, pixel_coords, frame_rate, x_dtype): - # cos/sin shaped [T=10, L=4]; identity token slice under test is [6:9). T, L = 10, 4 cos = torch.arange(T * L, dtype=torch.float32).reshape(T, L) sin = torch.arange(T * L, dtype=torch.float32).reshape(T, L) * -1.0 return (cos, sin, False) + def _process_output(self, x, embedded_timestep, keyframe_idxs, **kw): + return ("output", x, embedded_timestep) + def _fake_model(ltxv): inner = types.SimpleNamespace(diffusion_model=ltxv) return types.SimpleNamespace(model=inner) +def _ref_latent(): + return torch.zeros(1, 8, 1, 4, 4) + + def test_install_is_idempotent_and_strip_restores_originals(): s = _sampler() ltxv = _FakeLTXV() + orig_pi = ltxv._process_input orig_pt = ltxv._prepare_timestep orig_pe = ltxv._prepare_positional_embeddings + orig_po = ltxv._process_output - handle1 = s._install_identity_phase(_fake_model(ltxv)) + handle1 = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) assert handle1 is not None - assert ltxv._prepare_timestep is not orig_pt - handle2 = s._install_identity_phase(_fake_model(ltxv)) + assert ltxv._process_input is not orig_pi + handle2 = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) assert handle2 is None # already installed — idempotent, no double-wrap - s._strip_identity_phase(handle1) - # Bound-method objects aren't identity-stable across separate attribute accesses in - # CPython (a fresh wrapper is created each time) — compare by equality (same __func__ + - # __self__) instead of `is`. + s._strip_identity_overlap(handle1) + assert ltxv._process_input == orig_pi assert ltxv._prepare_timestep == orig_pt assert ltxv._prepare_positional_embeddings == orig_pe - assert not hasattr(ltxv, "_funpack_identity_range") + assert ltxv._process_output == orig_po + assert not hasattr(ltxv, "_funpack_id_ref_len") -def test_prepare_timestep_stashes_range_for_tagged_entry_only(): +def test_process_input_appends_ref_tokens_and_stashes_len(): s = _sampler() ltxv = _FakeLTXV() - handle = s._install_identity_phase(_fake_model(ltxv)) - entries = [ - {"pre_filter_count": 16, "surviving_count": 2}, # unrelated guide, first - {"pre_filter_count": 16, "surviving_count": 3, "funpack_identity_phase": 2.5}, # the tagged one - {"pre_filter_count": 16, "surviving_count": 4}, # unrelated guide, after - ] - ltxv._prepare_timestep(torch.zeros(1), 1, torch.float32, - resolved_guide_entries=entries, num_guide_tokens=9) - assert ltxv._funpack_identity_range == (2, 3, 9) # offset past the first entry's 2 tokens - assert ltxv._funpack_identity_phase_value == 2.5 - s._strip_identity_phase(handle) + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) + x = torch.zeros(1, 5, 8) # 5 target tokens + xx, pix, add = ltxv._process_input(x, None, None) + assert xx.shape[1] == 7 # 5 target + 2 ref tokens + assert ltxv._funpack_id_ref_len == 2 + assert ltxv._funpack_id_target_len == 5 + s._strip_identity_overlap(handle) -def test_prepare_timestep_clears_range_without_tagged_entry(): +def test_prepare_timestep_gives_ref_tokens_clean_timestep(): s = _sampler() ltxv = _FakeLTXV() - handle = s._install_identity_phase(_fake_model(ltxv)) - ltxv._funpack_identity_range = (0, 1, 1) # stale from a previous scene - entries = [{"pre_filter_count": 16, "surviving_count": 5}] - ltxv._prepare_timestep(torch.zeros(1), 1, torch.float32, - resolved_guide_entries=entries, num_guide_tokens=5) - assert ltxv._funpack_identity_range is None - s._strip_identity_phase(handle) + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) + ltxv._process_input(torch.zeros(1, 5, 8), None, None) # sets ref_len=2, target_len=5 + _, timestep, batch_size, _ = ltxv._prepare_timestep(torch.full((1,), 0.7), 1, torch.float32) + assert timestep.shape == (1, 7) + assert torch.all(timestep[:, :5] == 0.7) + assert torch.all(timestep[:, 5:] == 0.0) # ref tokens: clean timestep + s._strip_identity_overlap(handle) -def test_prepare_pe_rotates_only_the_tagged_range(): +def test_prepare_pe_rotates_only_the_ref_token_tail(): s = _sampler() ltxv = _FakeLTXV() - handle = s._install_identity_phase(_fake_model(ltxv)) + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) baseline_cos, baseline_sin, _ = ltxv._prepare_positional_embeddings(None, 25, torch.float32) - - entries = [ - {"pre_filter_count": 16, "surviving_count": 2}, - {"pre_filter_count": 16, "surviving_count": 3, "funpack_identity_phase": 2.5}, - {"pre_filter_count": 16, "surviving_count": 4}, - ] - ltxv._prepare_timestep(torch.zeros(1), 1, torch.float32, - resolved_guide_entries=entries, num_guide_tokens=9) + ltxv._process_input(torch.zeros(1, 5, 8), None, None) # ref_len=2 cos, sin, split_mode = ltxv._prepare_positional_embeddings(None, 25, torch.float32) - - # T=10, num_guide_tokens=9 -> guide_start=1; offset=2, length=3 -> rotated slice [3:6). - untouched = list(range(0, 3)) + list(range(6, 10)) + untouched = list(range(0, 8)) for i in untouched: - assert torch.equal(cos[i], baseline_cos[i]), f"row {i} should be untouched" - assert torch.equal(sin[i], baseline_sin[i]), f"row {i} should be untouched" - rotated_rows = [3, 4, 5] - changed = any(not torch.equal(cos[i], baseline_cos[i]) for i in rotated_rows) - assert changed, "rotated slice should differ from baseline" - assert split_mode is False # passthrough element preserved - s._strip_identity_phase(handle) + assert torch.equal(cos[i], baseline_cos[i]) + assert torch.equal(sin[i], baseline_sin[i]) + rotated = [8, 9] + changed = any(not torch.equal(cos[i], baseline_cos[i]) for i in rotated) + assert changed + assert split_mode is False + s._strip_identity_overlap(handle) -def test_rotate_identity_phase_is_noop_for_zero_phase_or_length(): +def test_prepare_pe_noop_when_ref_len_zero(): s = _sampler() - pe = (torch.ones(5, 4), torch.zeros(5, 4), False) - assert s._rotate_identity_phase(pe, 0, 0, 5, 2.0) is pe - assert s._rotate_identity_phase(pe, 0, 2, 5, 0.0) is pe + ltxv = _FakeLTXV() + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) + baseline = ltxv._prepare_positional_embeddings(None, 25, torch.float32) + # No _process_input call this "forward pass" -> _funpack_id_ref_len defaults to 0. + out = ltxv._prepare_positional_embeddings(None, 25, torch.float32) + assert torch.equal(out[0], baseline[0]) + assert torch.equal(out[1], baseline[1]) + s._strip_identity_overlap(handle) -def test_rotate_identity_phase_rejects_out_of_bounds_range(): +def test_process_output_trims_ref_tokens_from_plain_tensor(): s = _sampler() - cos = torch.arange(20, dtype=torch.float32).reshape(5, 4) - sin = -cos.clone() - pe = (cos, sin, False) - out = s._rotate_identity_phase(pe, offset=10, length=3, num_guide_tokens=5, phase=2.0) - assert out is pe # start/end fall outside T -> safe no-op + ltxv = _FakeLTXV() + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) + ltxv._process_input(torch.zeros(1, 5, 8), None, None) # ref_len=2 + x = torch.arange(7 * 3, dtype=torch.float32).reshape(1, 7, 3) + embedded_timestep = torch.zeros(1, 7, 1) + tag, out_x, out_et = ltxv._process_output(x, embedded_timestep, None) + assert out_x.shape[1] == 5 + assert torch.equal(out_x, x[:, :5]) + assert out_et.shape[1] == 5 + s._strip_identity_overlap(handle) + + +def test_process_output_trims_compressed_timestep_frames(): + s = _sampler() + ltxv = _FakeLTXV() + handle = s._install_identity_overlap(_fake_model(ltxv), _ref_latent(), 2.0) + ltxv._process_input([torch.zeros(1, 5, 8), torch.zeros(1, 3, 8)], None, None) + x = [torch.arange(7 * 3, dtype=torch.float32).reshape(1, 7, 3), torch.zeros(1, 3, 3)] + ct = _CompressedTimestep(data=torch.zeros(1, 7), num_frames=7, patches_per_frame=1) + tag, out_x, out_et = ltxv._process_output(x, [ct, "audio_et"], None) + assert out_x[0].shape[1] == 5 + assert out_et[0].num_frames == 5 + assert out_et[0].data.shape[1] == 5 + assert out_et[1] == "audio_et" + s._strip_identity_overlap(handle) + + +# ── _resolve_identity_overlap memoization ────────────────────────────────────── + +def test_resolve_identity_overlap_computes_once_per_state(monkeypatch): + s = _sampler() + monkeypatch.setattr(s, "_load_image_tensor", lambda filename: torch.zeros(1, 32, 32, 3)) + monkeypatch.setattr(s, "_latent_tensors", lambda chunk: [torch.zeros(1, 8, 1, 4, 4)]) + fake_comfy_utils = types.ModuleType("comfy.utils") + calls = [] + + def fake_upscale(img, w, h, mode, crop): + calls.append((w, h)) + return img + + fake_comfy_utils.common_upscale = fake_upscale + sys.modules["comfy.utils"] = fake_comfy_utils + sys.modules["comfy"].utils = fake_comfy_utils + + class _FakeVAE: + downscale_index_formula = [8, 8, 8] + + def encode(self, px): + return torch.zeros(1, 8, 1, 4, 4) + + state = {} + r1 = s._resolve_identity_overlap(state, "pin.png", _FakeVAE(), {}, "None", 2.0, 1.0, 1.0, "auto_adjust", False) + assert len(calls) == 1 + r2 = s._resolve_identity_overlap(state, "pin.png", _FakeVAE(), {}, "None", 2.0, 1.0, 1.0, "auto_adjust", False) + assert len(calls) == 1 # second call served from cache, no re-encode + assert r1 == r2 + assert r1[2] is None and r1[3] is None # no projector selected -> no ArcFace tokens From 864d222ee589d543754e047a894a295c0fc820a2 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Mon, 6 Jul 2026 22:50:09 +0300 Subject: [PATCH 02/15] Preview stability: route timeline filmstrip/waveform through segments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Today's preview-stability fix (4bc56e8) made the player pool fetch every chain-shared clip through a trimmed, faststart preview segment instead of deep-seeking the raw chain output (moov-at-end, stalls/blacks in the browser). Two other consumers of the raw file were missed: - appendFilmstrip built a