187 lines
8.4 KiB
Python
187 lines
8.4 KiB
Python
"""Per-window guide injection for LTXV context windows.
|
|
|
|
A per-window guide is a clean guide latent injected into each context window at a fixed
|
|
window-relative position, rather than at one absolute position like a normal LTXVAddGuide guide.
|
|
Its per-window strength can vary, and a strength of 0 skips a given window. It rides on the
|
|
conditioning (key ``PERWINDOW_KEY``) and never touches the working latent. The engine appends it
|
|
to a window's token sequence, manufactures the matching ``keyframe_idxs``,
|
|
``guide_attention_entries``, and ``denoise_mask`` for that window, then strips it back off the
|
|
output. The guide is physically a tail token, and its position is purely its RoPE coordinate, so
|
|
``rel_index`` may be negative (a coordinate before the window's first frame).
|
|
"""
|
|
from __future__ import annotations
|
|
import os
|
|
import logging
|
|
import torch
|
|
import comfy.conds
|
|
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
|
|
|
|
PERWINDOW_KEY = "ltex_perwindow_guides"
|
|
|
|
LTEX_DEBUG = bool(os.environ.get("LTEX_DEBUG")) # off by default; set LTEX_DEBUG=1 to trace per-window guide injection
|
|
|
|
|
|
def _dbg(msg):
|
|
if LTEX_DEBUG:
|
|
logging.info("[LTEx perwindow] " + msg)
|
|
|
|
|
|
def get_perwindow_guides(conds) -> list | None:
|
|
"""Return the raw per-window guide dicts carried on the conditioning, if any.
|
|
|
|
The guides ride as a top-level cond-dict key (set via conditioning_set_values); unrecognized
|
|
keys are not wrapped into model_conds by extra_conds, so they live at the top level here.
|
|
"""
|
|
for cond_list in conds:
|
|
if cond_list is None:
|
|
continue
|
|
for cond_dict in cond_list:
|
|
entry = cond_dict.get(PERWINDOW_KEY)
|
|
if entry:
|
|
return entry
|
|
return None
|
|
|
|
|
|
def prepare(model, guides: list) -> list:
|
|
"""Move guide latents into model latent space once, caching alongside the metadata.
|
|
|
|
Returns a list of dicts: {latent (process_latent_in'd, on intermediate device), rel_index,
|
|
strength, pixel_mask, latent_shape [F,H,W]}.
|
|
"""
|
|
prepared = []
|
|
for g in guides:
|
|
latent = model.process_latent_in(g["latent"])
|
|
prepared.append({
|
|
"latent": latent,
|
|
"rel_index": int(g["rel_index"]),
|
|
"strength": g.get("strength", 1.0), # float, or a per-window list of floats
|
|
"pixel_mask": g.get("pixel_mask"),
|
|
"latent_shape": list(latent.shape[2:]), # [F, H, W]
|
|
})
|
|
return prepared
|
|
|
|
|
|
def resolve_strength(guide, canonical_index: int) -> float:
|
|
"""Resolve a guide's strength for a given canonical (standard_static) window index.
|
|
|
|
A scalar applies to every window; a list indexes by canonical window, clamping to the last entry.
|
|
"""
|
|
s = guide["strength"]
|
|
if isinstance(s, (list, tuple)):
|
|
if not s:
|
|
return 1.0
|
|
return float(s[min(int(canonical_index), len(s) - 1)])
|
|
return float(s)
|
|
|
|
|
|
def active_guides(prepared: list, canonical_index: int) -> list:
|
|
"""Guides whose resolved strength for this window is > 0; the rest are skipped entirely."""
|
|
return [p for p in prepared if resolve_strength(p, canonical_index) > 0.0]
|
|
|
|
|
|
def total_frame_count(prepared: list) -> int:
|
|
return sum(p["latent_shape"][0] for p in prepared)
|
|
|
|
|
|
def window_guide_latent(prepared: list, dim: int, device) -> torch.Tensor:
|
|
"""Concatenated guide latents to append to a window's primary-modality slice."""
|
|
frames = [p["latent"].to(device) for p in prepared]
|
|
return torch.cat(frames, dim=dim)
|
|
|
|
|
|
def _guide_pixel_coords(model, guide_latent, rel_index, device):
|
|
"""Pixel (RoPE) coords for one guide latent placed at window-local position rel_index.
|
|
|
|
Mirrors LTXVAddGuide.add_keyframe_index: patchify for the base grid, then offset the temporal
|
|
channel by rel_index in pixel space (one latent step == time scale factor), bypassing the causal
|
|
clamp so negative positions survive. rel_index is measured from the window's first slot, so a
|
|
causal_window_fix anchor (which occupies that slot) sits between rel_index 0 and -1: the guide at
|
|
-1 lands one step before the anchor, not on it.
|
|
"""
|
|
patchifier = model.diffusion_model.patchifier
|
|
scale_factors = model.diffusion_model.vae_scale_factors
|
|
time_scale_factor = scale_factors[0]
|
|
_, latent_coords = patchifier.patchify(guide_latent.to(device))
|
|
# causal_fix only affects the base frame-0 token; the explicit offset below carries position.
|
|
pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=False)
|
|
pixel_coords[:, 0] += rel_index * time_scale_factor
|
|
return pixel_coords
|
|
|
|
|
|
def apply_to_window_cond(new_cond_item: dict, prepared: list, model, window, x_in: torch.Tensor):
|
|
"""Extend (or create) keyframe_idxs / guide_attention_entries / denoise_mask for a window so
|
|
the appended per-window guide tokens are treated as clean, RoPE-positioned guide tokens.
|
|
|
|
Order is [existing guide tokens][per-window guide tokens], matching the latent append order.
|
|
"""
|
|
device = x_in.device # build every per-window cond on the working-latent device
|
|
dim = window.dim
|
|
H, W = x_in.shape[3], x_in.shape[4]
|
|
anchor_idx = getattr(window, "causal_anchor_index", None)
|
|
anchor_shift = 1 if (anchor_idx is not None and anchor_idx >= 0) else 0
|
|
# window-local video length (frames the model will denoise), excluding any appended guides
|
|
video_len = len(window.index_list) + anchor_shift
|
|
|
|
ci = getattr(window, "canonical_index", 0)
|
|
active = active_guides(prepared, ci)
|
|
if not active:
|
|
_dbg(f"window {window.index_list[0]}-{window.index_list[-1]} canonical={ci} | no guide (all strengths 0)")
|
|
return # every guide is strength 0 for this window: inject nothing, leave conds untouched
|
|
|
|
win = f"{window.index_list[0]}-{window.index_list[-1]}"
|
|
coords = []
|
|
entries = []
|
|
strengths = []
|
|
for i, p in enumerate(active):
|
|
s = resolve_strength(p, ci)
|
|
strengths.append(s)
|
|
c = _guide_pixel_coords(model, p["latent"], p["rel_index"], device)
|
|
coords.append(c)
|
|
pixel_mask = p["pixel_mask"]
|
|
entries.append({
|
|
"pre_filter_count": p["latent_shape"][0] * H * W,
|
|
"strength": s,
|
|
"pixel_mask": pixel_mask.to(device) if isinstance(pixel_mask, torch.Tensor) else None,
|
|
"latent_shape": [p["latent_shape"][0], H, W],
|
|
})
|
|
if LTEX_DEBUG:
|
|
_dbg(f"window {win} canonical={ci} | guide#{i} rel_index={p['rel_index']} "
|
|
f"coord_t={c[0, 0, 0, 0].item():.1f} strength={s:.3f}")
|
|
pw_coords = torch.cat(coords, dim=2)
|
|
|
|
# keyframe_idxs: append our coords after any existing (already window-local) guide coords.
|
|
existing_kf = new_cond_item.get("keyframe_idxs")
|
|
if existing_kf is not None and hasattr(existing_kf, "cond") and existing_kf.cond is not None:
|
|
merged = torch.cat([existing_kf.cond.to(device), pw_coords], dim=2)
|
|
new_cond_item["keyframe_idxs"] = existing_kf._copy_with(merged)
|
|
else:
|
|
new_cond_item["keyframe_idxs"] = comfy.conds.CONDRegular(pw_coords)
|
|
|
|
# guide_attention_entries: append after any existing entries.
|
|
existing_entries = new_cond_item.get("guide_attention_entries")
|
|
if existing_entries is not None and hasattr(existing_entries, "cond") and existing_entries.cond:
|
|
merged_entries = [*existing_entries.cond, *entries]
|
|
new_cond_item["guide_attention_entries"] = existing_entries._copy_with(merged_entries)
|
|
else:
|
|
new_cond_item["guide_attention_entries"] = comfy.conds.CONDConstant(entries)
|
|
|
|
# denoise_mask: append a keep-mask (1 - strength) for our tokens; build a video portion if none.
|
|
B = x_in.shape[0]
|
|
pw_mask_parts = [
|
|
torch.full((B, 1, p["latent_shape"][0], H, W), max(0.0, 1.0 - s),
|
|
device=device, dtype=x_in.dtype)
|
|
for p, s in zip(active, strengths)
|
|
]
|
|
pw_mask = torch.cat(pw_mask_parts, dim=dim)
|
|
existing_dm = new_cond_item.get("denoise_mask")
|
|
if existing_dm is not None and hasattr(existing_dm, "cond") and isinstance(existing_dm.cond, torch.Tensor):
|
|
base = existing_dm.cond.to(device)
|
|
if base.shape[3] != H or base.shape[4] != W:
|
|
base = base.expand(-1, -1, -1, H, W)
|
|
merged_mask = torch.cat([base, pw_mask], dim=dim)
|
|
new_cond_item["denoise_mask"] = existing_dm._copy_with(merged_mask)
|
|
else:
|
|
video_mask = torch.ones((B, 1, video_len, H, W), device=device, dtype=x_in.dtype)
|
|
merged_mask = torch.cat([video_mask, pw_mask], dim=dim)
|
|
new_cond_item["denoise_mask"] = comfy.conds.CONDRegular(merged_mask)
|