71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
"""Shared temporal placement helpers for clean LTX reference tokens.
|
|
|
|
LTX VAEs normally compress eight pixel frames into one latent-time step. The
|
|
public node input is expressed in latent steps so ``-1`` consistently means
|
|
"one VAE temporal step before frame zero" without hard-coding a particular
|
|
VAE scale factor.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
DEFAULT_REFERENCE_TEMPORAL_OFFSET_LATENTS = 0
|
|
|
|
|
|
def reference_temporal_offset_input():
|
|
"""Return a fresh ComfyUI input declaration for reference-only placement."""
|
|
return (
|
|
"INT",
|
|
{
|
|
"default": DEFAULT_REFERENCE_TEMPORAL_OFFSET_LATENTS,
|
|
"min": -16,
|
|
"max": 16,
|
|
"step": 1,
|
|
"tooltip": (
|
|
"Temporal RoPE offset for appearance/reference tokens only, in VAE latent-time steps. "
|
|
"0 keeps the standard frame-0 placement used by existing LoRAs. "
|
|
"-1 places the reference one latent step before frame 0 (normally -8 pixel frames), "
|
|
"which can be tested to reduce overlap frame-0 reference leakage. "
|
|
"Guide, mask, and target positions are never changed."
|
|
),
|
|
},
|
|
)
|
|
|
|
|
|
def reference_temporal_pixel_offset(
|
|
offset_latents: int, temporal_scale_factor: int
|
|
) -> int:
|
|
"""Convert a latent-step offset to the raw pixel-frame units used by LTX RoPE."""
|
|
return int(offset_latents) * int(temporal_scale_factor)
|
|
|
|
|
|
def shift_reference_temporal_positions(
|
|
positions: torch.Tensor,
|
|
offset_latents: int,
|
|
temporal_scale_factor: int,
|
|
) -> torch.Tensor:
|
|
"""Shift only the temporal axis of a reference coordinate tensor.
|
|
|
|
The helper accepts both ComfyUI's ``[B, axes, tokens]`` coordinates and
|
|
trainer-style tensors with extra trailing dimensions. A zero offset is an
|
|
exact no-op and returns the original tensor for backward-compatible specs.
|
|
"""
|
|
pixel_offset = reference_temporal_pixel_offset(
|
|
offset_latents, temporal_scale_factor
|
|
)
|
|
if pixel_offset == 0:
|
|
return positions
|
|
if positions.ndim < 3 or positions.shape[1] < 1:
|
|
raise ValueError(
|
|
"Reference temporal positions must have shape [batch, axes, ...] "
|
|
f"with a temporal axis; got {tuple(positions.shape)}"
|
|
)
|
|
shifted = positions.clone()
|
|
shifted[:, 0, ...] += torch.as_tensor(
|
|
pixel_offset,
|
|
device=shifted.device,
|
|
dtype=shifted.dtype,
|
|
)
|
|
return shifted
|