"""Conditioning-carried guides for LTXV / LTXAV (proof of concept). Unlike core's LTXVAddGuide, the guide latent is NOT appended to the working latent. It rides on the conditioning and is injected as extra tokens *inside* ``LTXBaseModel._forward`` (with rope + ``(1 - strength) * sigma`` per-token timestep), then stripped back off before the output is unpatchified. So there is nothing to crop afterward and the working latent is never touched. Reaching mid-``_forward`` (between ``_process_input`` and ``_prepare_timestep``) is not possible through any official wrapper hook, so this installs a **gated, idempotent monkeypatch** of ``LTXBaseModel._forward`` plus the ``LTXV`` / ``LTXAV`` ``extra_conds``. The patches are a complete no-op unless the ``guide_cond_latents`` conditioning key is present, so normal LTXV/LTXAV sampling is byte-identical to stock core. This deliberately departs from the rest of the pack's "no monkeypatch core" approach because the official hooks cannot express this injection; see the node docstring for the rationale. LTXV uses single tensors throughout. LTXAV carries video/audio as ``[video, audio]`` lists; guides are video-only, so injection targets the video slot (index 0) and leaves audio untouched. """ from __future__ import annotations import torch import node_helpers import comfy.conds import comfy.model_base import comfy.ldm.lightricks.model as ltx_model from comfy.ldm.lightricks.symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords GUIDE_LATENTS_KEY = "guide_cond_latents" GUIDE_COORDS_KEY = "guide_cond_keyframe_idxs" GUIDE_STRENGTH_KEY = "guide_cond_strength" # Standalone patchifier for node-side coord computation (matches LTXVModel.patchifier config). LTXCONDGUIDE_ENABLED = False # Set to False by default. You must enable this manually for now. LTXCONDGUIDE_PATCHIFIER = SymmetricPatchifier(1, start_end=True) LTXCONDGUIDE_PATCHED = False # -------------------------------------------------------------------------------------------- # Node-side helpers (model-free): coordinates + conditioning accumulation # -------------------------------------------------------------------------------------------- def compute_guide_coords(guide_latent, frame_idx, scale_factors, causal_fix, latent_downscale_factor=1): """Pixel-space keyframe coordinates for a guide latent placed at pixel-frame ``frame_idx``. Mirrors LTXVAddGuide.add_keyframe_index. For a plain guide (``latent_downscale_factor == 1``): patchify for the base grid, then offset the temporal channel by frame_idx. For IC-LoRA references (``latent_downscale_factor > 1``) the low-res guide tokens are spread across the full-res grid (dilation) and each token's spatial RoPE *end* is widened to represent the larger patch it covers, exactly as core does. The hole positions are then dropped so only the real tokens remain, in the same order as a plain patchify of the low-res ``guide_latent`` (so they align with the model-side injected tokens). The conditioning path injects no dense grid, so unlike core there is nothing to filter later. Shape: (B, [t,h,w], token, [start,end]). """ if latent_downscale_factor > 1: from comfy_extras.nodes_lt import LTXVAddGuide dilated, dilated_mask = LTXVAddGuide.dilate_latent(guide_latent, latent_downscale_factor) _, latent_coords = LTXCONDGUIDE_PATCHIFIER.patchify(dilated) pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=causal_fix) pixel_coords[:, 0] += frame_idx spatial_end_offset = (latent_downscale_factor - 1) * torch.tensor( scale_factors[1:], device=pixel_coords.device, ).view(1, -1, 1, 1) pixel_coords[:, 1:, :, 1:] += spatial_end_offset.to(pixel_coords.dtype) # Holes are marked < 0 in the dilated mask; keep only the real tokens. mask_tokens, _ = LTXCONDGUIDE_PATCHIFIER.patchify(dilated_mask) keep = mask_tokens[0, :, 0] > 0 return pixel_coords[:, :, keep, :] _, latent_coords = LTXCONDGUIDE_PATCHIFIER.patchify(guide_latent) pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=causal_fix) pixel_coords[:, 0] += frame_idx return pixel_coords def _existing_coords(cond): for t in cond: v = t[1].get(GUIDE_COORDS_KEY) if v is not None: return v return None def append_guide(cond, guide_latent, pixel_coords, strength): """Store one guide on the conditioning, accumulating across chained nodes. Coords are a tensor (concatenated along the token dim; node_helpers append=True would add tensors elementwise, which is wrong). Latents/strengths are lists (append=True concatenates), kept aligned with the coord token order. """ existing = _existing_coords(cond) if existing is not None: pixel_coords = torch.cat([existing, pixel_coords], dim=2) cond = node_helpers.conditioning_set_values(cond, {GUIDE_COORDS_KEY: pixel_coords}) cond = node_helpers.conditioning_set_values( cond, {GUIDE_LATENTS_KEY: [guide_latent], GUIDE_STRENGTH_KEY: [strength]}, append=True ) return cond # -------------------------------------------------------------------------------------------- # Forward-injection helpers (model-side), ported from the core branch; list-aware for LTXAV # -------------------------------------------------------------------------------------------- def _inject(self, x, pixel_coords, timestep, batch_size, merged_args): guide_latents = merged_args[GUIDE_LATENTS_KEY] guide_coords = merged_args[GUIDE_COORDS_KEY] guide_strength = merged_args.get(GUIDE_STRENGTH_KEY, None) is_av = isinstance(x, list) vx = x[0] if is_av else x v_coords = pixel_coords[0] if is_av else pixel_coords main_len = vx.shape[1] # Full-noise reference sigma for the guide timestep: amax over any base per-token timestep, so a # strength-1 guide still resolves to a clean (~0) timestep even under partial denoise. sigma = timestep.reshape(batch_size, -1).amax(dim=1, keepdim=True) # [B, 1] # Preserve a per-token base timestep (denoise_mask / partial-denoise / inpainting) on the main # tokens rather than flattening to uniform sigma. Falls back to uniform when the incoming timestep # is a single per-batch value (the common empty-latent, denoise=1.0 case). base_ts = timestep.reshape(batch_size, -1) main_ts = base_ts if base_ts.shape[1] == main_len else sigma.expand(batch_size, main_len) timestep_parts = [main_ts] guide_tokens = [] for i, lat in enumerate(guide_latents): gx, _ = self.patchifier.patchify(lat.to(device=vx.device, dtype=vx.dtype)) gx = self.patchify_proj(gx) # same input projection (in_channels -> inner_dim) base tokens get guide_tokens.append(gx) s = guide_strength[i] if guide_strength is not None else 1.0 timestep_parts.append(sigma.expand(batch_size, gx.shape[1]) * max(0.0, 1.0 - float(s))) # Enable the model's per-guide self-attention mask (spatial pixel_mask + strength attenuation). # This path never drops guide tokens (no denoise-mask grid filter), so surviving_count == # pre_filter_count for every entry; the rest is reused from core's _build_guide_self_attention_mask. merged_args["num_guide_tokens"] = sum(g.shape[1] for g in guide_tokens) entries = merged_args.get("guide_attention_entries", None) if entries: merged_args["resolved_guide_entries"] = [ {**e, "surviving_count": e["pre_filter_count"]} for e in entries ] vx = torch.cat([vx] + guide_tokens, dim=1) v_coords = torch.cat([v_coords, guide_coords.to(v_coords.device)], dim=2) timestep = torch.cat(timestep_parts, dim=1) if is_av: x = [vx, x[1]] pixel_coords = [v_coords, pixel_coords[1]] else: x, pixel_coords = vx, v_coords return x, pixel_coords, timestep, main_len def _strip(x, embedded_timestep, main_len): if isinstance(x, list): x[0] = x[0][:, :main_len] v_emb = embedded_timestep[0] if not torch.is_tensor(v_emb): # CompressedTimestep (per-frame) -> expand to per-token v_emb = v_emb.expand() embedded_timestep[0] = v_emb[:, :main_len] else: x = x[:, :main_len] embedded_timestep = embedded_timestep[:, :main_len] return x, embedded_timestep # -------------------------------------------------------------------------------------------- # Gated monkeypatch installer # -------------------------------------------------------------------------------------------- def _make_patched_forward(orig_forward): def _forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=None, denoise_mask=None, **kwargs): # Gate: no guides -> delegate to the original forward, byte-identical to stock core. if kwargs.get(GUIDE_LATENTS_KEY) is None and transformer_options.get(GUIDE_LATENTS_KEY) is None: return orig_forward(self, x, timestep, context, attention_mask, frame_rate, transformer_options, keyframe_idxs, denoise_mask=denoise_mask, **kwargs) # Guide path: this mirrors core LTXBaseModel._forward with the two guarded inject/strip calls. if isinstance(x, list): input_dtype = x[0].dtype batch_size = x[0].shape[0] else: input_dtype = x.dtype batch_size = x.shape[0] merged_args = {**transformer_options, **kwargs} x, pixel_coords, additional_args = self._process_input(x, keyframe_idxs, denoise_mask, **merged_args) merged_args.update(additional_args) x, pixel_coords, timestep, main_len = _inject(self, x, pixel_coords, timestep, batch_size, merged_args) timestep, embedded_timestep, prompt_timestep = self._prepare_timestep(timestep, batch_size, input_dtype, **merged_args) merged_args["prompt_timestep"] = prompt_timestep context, attention_mask = self._prepare_context(context, batch_size, x, attention_mask) attention_mask = self._prepare_attention_mask(attention_mask, input_dtype) pe = self._prepare_positional_embeddings(pixel_coords, frame_rate, input_dtype) self_attention_mask = self._build_guide_self_attention_mask(x, transformer_options, merged_args) x = self._process_transformer_blocks( x, context, attention_mask, timestep, pe, transformer_options=transformer_options, self_attention_mask=self_attention_mask, **merged_args, ) x, embedded_timestep = _strip(x, embedded_timestep, main_len) # keyframe_idxs is None on this path (guides never entered the latent), so _process_output # skips its grid-mask scatter branch. x = self._process_output(x, embedded_timestep, None, **merged_args) return x _forward._ltex_guide_cond = True return _forward def _make_patched_extra_conds(orig_extra_conds): def extra_conds(self, **kwargs): out = orig_extra_conds(self, **kwargs) guide_latents = kwargs.get(GUIDE_LATENTS_KEY, None) if guide_latents is not None: # process_latent_in here because the guides bypass the main latent (the latent-carried # path gets it for free via the sampler). Safe for LTXAV: its latent format is LTXAV(LTXV), # so a plain video guide gets the same video normalization as the video stream. out[GUIDE_LATENTS_KEY] = comfy.conds.CONDList([self.process_latent_in(l) for l in guide_latents]) guide_coords = kwargs.get(GUIDE_COORDS_KEY, None) if guide_coords is not None: out[GUIDE_COORDS_KEY] = comfy.conds.CONDRegular(guide_coords) guide_strength = kwargs.get(GUIDE_STRENGTH_KEY, None) if guide_strength is not None: out[GUIDE_STRENGTH_KEY] = comfy.conds.CONDConstant(guide_strength) return out extra_conds._ltex_guide_cond = True return extra_conds def install(): """Install the gated patches once. Idempotent across module reloads (guarded by a marker attribute on the patched callables).""" global LTXCONDGUIDE_PATCHED if LTXCONDGUIDE_PATCHED or not LTXCONDGUIDE_ENABLED: return if not getattr(ltx_model.LTXBaseModel._forward, "_ltex_guide_cond", False): ltx_model.LTXBaseModel._forward = _make_patched_forward(ltx_model.LTXBaseModel._forward) for model_cls in (comfy.model_base.LTXV, getattr(comfy.model_base, "LTXAV", None)): if model_cls is None: continue if not getattr(model_cls.extra_conds, "_ltex_guide_cond", False): model_cls.extra_conds = _make_patched_extra_conds(model_cls.extra_conds) LTXCONDGUIDE_PATCHED = True