Add SCAIL-2 driving-latent conditioning (Phase 1, inference PoC)

Port mechanisms 1+3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: concatenate a
driving video latent directly into the DiT token sequence with a width-axis
RoPE offset (ΔW) so driving coords stay detached from the target video.

- New VideoConditionByDrivingLatent + DrivingMode in ltx-core conditioning,
  modeled on VideoConditionByReferenceLatent (patchify -> positions -> append
  -> attention mask). Applies ΔW width shift, aligns time to the target, and
  guards against RoPE wrap (max_pos) and target/driving shape mismatch.
- Export both from conditioning packages.
- docs/plan.md and docs/tasks.md track the phased port.

Inference-only: no weight changes. Mechanism 2 (in-context mask channels,
patchify_proj widening) and training are deferred to Phase 2+. Validated via
plumbing checks and a real (random-weight) transformer forward smoke run;
visual quality is not validated (requires Phase 2/3 finetuning).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-09 09:06:35 +08:00
parent 9377758131
commit baa6646fd1
5 changed files with 270 additions and 0 deletions
@@ -5,6 +5,8 @@ from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.types import (
AudioConditionByReferenceLatent,
ConditioningItemAttentionStrengthWrapper,
DrivingMode,
VideoConditionByDrivingLatent,
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
VideoConditionByMask,
@@ -16,6 +18,8 @@ __all__ = [
"ConditioningError",
"ConditioningItem",
"ConditioningItemAttentionStrengthWrapper",
"DrivingMode",
"VideoConditionByDrivingLatent",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
"VideoConditionByMask",
@@ -1,6 +1,7 @@
"""Conditioning type implementations."""
from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
from ltx_core.conditioning.types.driving_video_cond import DrivingMode, VideoConditionByDrivingLatent
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
@@ -10,6 +11,8 @@ from ltx_core.conditioning.types.reference_video_cond import VideoConditionByRef
__all__ = [
"AudioConditionByReferenceLatent",
"ConditioningItemAttentionStrengthWrapper",
"DrivingMode",
"VideoConditionByDrivingLatent",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
"VideoConditionByMask",
@@ -0,0 +1,160 @@
"""Driving-video conditioning for SCAIL-2-style end-to-end character animation.
Ports mechanism 1 + 3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: the *driving*
video latent is concatenated directly into the DiT token sequence (no skeleton /
pose intermediate), and carries a fixed spatial offset ``width_offset`` (the
paper's ΔW) on the RoPE width axis so its coordinates stay detached from the
main video tokens. This is the inference-only PoC path -- it reuses the existing
frozen-reference-token machinery and does not touch the transformer or RoPE.
This mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
(same patchify -> positions -> append -> attention-mask flow); the only new
behaviour is the mode-aware coordinate assignment.
Scope note (Phase 1): SCAIL-2's in-context mask channels (mechanism 2) and the
reference-latent height shift ΔH_ref of Replacement Mode require model-weight
surgery / a separate reference token group and are intentionally out of scope
here. Because LTX handles the reference image as a frame-0 in-place replacement
(not a separate token group), the driving-token placement is identical for both
:class:`DrivingMode` values in Phase 1 -- ``mode`` is stored for forward
compatibility and to document intent, but does not yet alter the driving
coordinates. See the plan for the deferred Phase 2 work.
"""
from __future__ import annotations
from enum import Enum
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState, VideoLatentShape
# Default normalization ceiling for the RoPE width axis. Must stay in sync with
# the model's ``positional_embedding_max_pos[2]`` (see rope.py:precompute_freqs_cis
# default ``max_pos=[20, 2048, 2048]`` and model.py `_init_` default). Driving
# width coordinates that reach or exceed this value would wrap under RoPE.
DEFAULT_MAX_WIDTH_POSITION = 2048
class DrivingMode(Enum):
"""SCAIL-2 conditioning mode. See module docstring for the Phase 1 caveat."""
ANIMATION = "animation"
REPLACEMENT = "replacement"
class VideoConditionByDrivingLatent(ConditioningItem):
"""Append driving-video tokens with a width-axis RoPE offset (SCAIL-2 ΔW).
The driving tokens are appended after the target sequence as clean latents
(placeholder zeros in the noisy latent), kept frozen (``denoise_mask =
1 - strength``), temporally aligned to the target, and shifted along the
width axis by ``width_offset`` so they occupy ``[ΔW, ΔW + Wv)`` while the
target stays at ``[0, Wv)``.
Args:
latent: Driving video latents ``[B, C, F, H, W]``. Must match the target
shape (same F/H/W) so tokens align frame-for-frame with the target.
mode: SCAIL-2 mode (reserved for Phase 2; see module docstring).
width_offset: ΔW in RoPE pixel-space width units. ``None`` (default) uses
the target's pixel width (``target_shape.width * scale_factors.width``),
placing the driving tokens immediately to the right of the target.
strength: 1.0 keeps the driving latent fully clean (frozen); 0.0 would
denoise it. Default 1.0.
max_width_position: RoPE width normalization ceiling; validation raises if
the shifted driving coordinates would reach it. Keep in sync with the
model's ``positional_embedding_max_pos[2]``.
"""
def __init__(
self,
latent: torch.Tensor,
mode: DrivingMode = DrivingMode.ANIMATION,
width_offset: float | None = None,
strength: float = 1.0,
max_width_position: int = DEFAULT_MAX_WIDTH_POSITION,
):
self.latent = latent
self.mode = mode
self.width_offset = width_offset
self.strength = strength
self.max_width_position = max_width_position
def apply_to(
self,
latent_state: LatentState,
latent_tools: VideoLatentTools,
) -> LatentState:
"""Append driving tokens with target-aligned time and a ΔW width shift."""
tokens = latent_tools.patchifier.patchify(self.latent)
num_target_tokens = latent_tools.patchifier.get_token_count(latent_tools.target_shape)
if tokens.shape[1] != num_target_tokens:
raise ValueError(
"VideoConditionByDrivingLatent expects the driving latent to match the target shape "
f"(same F/H/W): got {tokens.shape[1]} driving tokens vs {num_target_tokens} target tokens. "
"Resize/resample the driving video to the target resolution and frame count."
)
# Compute the driving tokens' own pixel-space coordinates (same flow as
# the base reference conditioning and create_initial_state).
latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
device=self.latent.device,
)
positions = get_pixel_coords(
latent_coords=latent_coords,
scale_factors=latent_tools.scale_factors,
causal_fix=latent_tools.causal_fix,
).to(dtype=torch.float32)
# Temporal alignment: copy the target's time coordinates so the driving
# tokens sit on exactly the same time grid as z_t (robust to causal_fix /
# fps nuances). Token order is a flattened (f h w) grid identical to the
# target's, so a token-wise copy is frame-aligned.
positions[:, 0:1, :] = latent_state.positions[:, 0:1, :num_target_tokens].to(dtype=torch.float32)
# ΔW: shift the driving tokens along the width axis so they stay spatially
# detached from the target tokens. Default offset = target pixel width,
# giving target=[0, Wv), driving=[Wv, 2*Wv).
width_offset = self.width_offset
if width_offset is None:
width_offset = float(latent_tools.target_shape.width * latent_tools.scale_factors.width)
positions[:, 2, ...] = positions[:, 2, ...] + width_offset
max_width = positions[:, 2, ...].max().item()
if max_width >= self.max_width_position:
raise ValueError(
f"Driving width coordinate {max_width:.1f} reaches the RoPE ceiling "
f"{self.max_width_position} and would wrap. Reduce width_offset "
f"(currently {width_offset:.1f}) or lower the output width."
)
denoise_mask = torch.full(
size=(*tokens.shape[:2], 1),
fill_value=1.0 - self.strength,
device=self.latent.device,
dtype=self.latent.dtype,
)
new_attention_mask = update_attention_mask(
latent_state=latent_state,
attention_mask=None,
num_noisy_tokens=num_target_tokens,
num_new_tokens=tokens.shape[1],
batch_size=tokens.shape[0],
device=self.latent.device,
dtype=self.latent.dtype,
)
return LatentState(
latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
)