Add SCAIL-2 in-context mask channels (Phase 2, plumbing + zero-init surgery)
Port mechanism 2 of SCAIL-2 (arXiv:2606.10804) to LTX-2: extra per-token in-context conditioning channels (1 environment switch + K=6 character binding slots) concatenated onto the latent before the first projection. Faithful temporal encoding for LTX's VAE (temporal factor 8): each latent frame stacks its 8 pixel sub-frames along the channel dim, giving 8*(K+1)=56 channels (vs the paper's 4*(K+1)=28 on Wan 2.1). Plumbing (backward compatible; cond_channels=None leaves every existing pipeline unchanged): - LatentState/Modality gain an optional cond_channels field (patchified [B,T,C]). - LTXModel(mask_conditioning_channels=0) config-gates a widened patchify_proj; TransformerArgsPreprocessor concatenates cond_channels (or zero-pads) before it. - tools.clear_conditioning trims it; token-appending conditioning items (reference video/audio, driving, keyframe) extend it via extend_cond_channels. New: - VideoConditionByMaskChannels: encodes (K+1) pixel masks -> 56 channels, written onto the trailing driving tokens (noisy target stays all-zero, per the paper). - widen_patchify_proj_for_mask_channels: zero-init checkpoint surgery so a converted model reproduces the base output exactly until finetuned. Verified (CPU, random weights): backward compat, zero-init widened forward is bit-identical to baseline for any cond_channels, and the driving+mask pipeline forwards without crashing with correct placement/clipping. Visual quality requires Phase 3 finetuning; Replacement-mode z_ref height shift still deferred. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+7
-6
@@ -37,12 +37,13 @@
|
||||
- 用現有 checkpoint 驗證資料流正確性（序列長度、座標偏移、attention mask、denoise_mask、不 wrap）。
|
||||
- **PoC 僅驗證 plumbing**；LTX-2 未經此訓練，畫面不會是正確動畫。
|
||||
|
||||
### Phase 2 — In-context mask channel（機制 2，checkpoint 手術）⬜ 未開始
|
||||
- `model.py:158` `patchify_proj` 由 `Linear(128, inner)` 加寬到 `Linear(128+28, inner)`。
|
||||
- `LatentState` 增加 optional conditioning-channel 欄位，跟著 patchify/concat/clear 流動。
|
||||
- 投影前 concat（`transformer_args.py:209`）。
|
||||
- 新輸入欄位 **zero-init**，載入舊 checkpoint 行為不變；寫 checkpoint 轉換 script。
|
||||
- 新增產生 28-channel mask（環境開關 + 角色槽）的 conditioning item。
|
||||
### Phase 2 — In-context mask channel（機制 2，checkpoint 手術）✅ 已完成（plumbing）
|
||||
- **決策**：時間編碼採忠實堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6），非 Wan 的 28。
|
||||
- `LTXModel(mask_conditioning_channels=0)` config-gated 加寬 `patchify_proj`；預設不變。
|
||||
- `LatentState`/`Modality` 加 `cond_channels` 欄位，跟著 clone/clear/append 流動；投影前在 `_apply_patchify_proj` concat（None 補零）。
|
||||
- 新輸入欄位 **zero-init**：`widen_patchify_proj_for_mask_channels` 轉換舊 checkpoint，行為不變、待微調才生效。
|
||||
- `VideoConditionByMaskChannels`：(K+1) 語意 mask → 空間下採樣 + 時間 8× 堆疊 → 寫入尾端 driving token（target 保持零，符合論文）。
|
||||
- 驗證通過（`verify_mask_channels.py`）：向後相容、zero-init 加寬 forward == baseline、mask pipeline 不 crash。**僅驗證 plumbing，畫質需 Phase 3 微調。**
|
||||
|
||||
### Phase 3 — 訓練整合（ltx-trainer）⬜ 未開始
|
||||
- dataset 產出 (target, driving, mask)，接上 Phase 1/2 conditioning，設微調 loss 與凍結策略。
|
||||
|
||||
+10
-6
@@ -15,14 +15,18 @@
|
||||
|
||||
## Phase 2 — In-context mask channel（checkpoint 手術）
|
||||
|
||||
> **決策**：mask channel 時間編碼採**忠實堆疊**——每 latent 幀對應 8 個 pixel 子幀沿 channel 堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6）。非 Wan 的 4×(K+1)=28。
|
||||
|
||||
| # | 任務 | 狀態 | 備註 |
|
||||
|---|---|---|---|
|
||||
| 2.1 | `patchify_proj` 加寬 `128 → 128+28` | ⬜ | `model.py:158` `_init_video`；`proj_out` 不動 |
|
||||
| 2.2 | `LatentState` 帶額外 conditioning-channel 欄位 | ⬜ | 同步改 `tools.py` patchify/unpatchify/clear、`Modality` |
|
||||
| 2.3 | 投影前 concat mask channel 到 `x` | ⬜ | `transformer_args.py:209` |
|
||||
| 2.4 | checkpoint zero-init 轉換 script | ⬜ | 新輸入欄位 zero-init，載入舊權重行為不變 |
|
||||
| 2.5 | 產生 28-channel mask 的 conditioning item | ⬜ | 環境開關 + K=6 角色綁定槽，`4(K+1)=28` |
|
||||
| 2.6 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | Phase 1 暫緩項 |
|
||||
| 2.1 | `patchify_proj` 加寬（config-gated） | ✅ | `LTXModel(mask_conditioning_channels=0)`；>0 時 `patchify_proj=Linear(in+mask, inner)`。預設不變 |
|
||||
| 2.2 | `LatentState`/`Modality` 帶 `cond_channels` 欄位 | ✅ | `types.py`、`modality.py` 加 `cond_channels: Tensor\|None=None`（patchified [B,T,C]）；`tools.clear_conditioning` 裁切；`helpers.modality_from_latent_state` 帶入 |
|
||||
| 2.3 | 投影前 concat/zero-pad cond_channels | ✅ | `transformer_args.py` `_apply_patchify_proj`：寬度不足時用 cond_channels 補齊，None 則補零 |
|
||||
| 2.4 | checkpoint zero-init 加寬轉換 | ✅ | `mask_channels_checkpoint.py` `widen_patchify_proj_for_mask_channels`：尾端補零欄，載入舊權重行為不變 |
|
||||
| 2.5 | `VideoConditionByMaskChannels`（56ch） | ✅ | `mask_channels_cond.py`：1 環境開關 + K=6 綁定槽 → 空間下採樣 + 時間 8× 堆疊(causal 首幀複製)=56ch，寫入尾端 driving token，target 保持零 |
|
||||
| 2.6 | token-append conditioning 延伸 cond_channels | ✅ | `cond_channels.py` `extend_cond_channels`；reference_video/reference_audio/driving/keyframe 皆接上 |
|
||||
| 2.7 | Phase 2 驗證（免訓練） | ✅ | `verify_mask_channels.py`：mask=0 向後相容；zero-init 加寬 forward == baseline（任意 cond_channels）；driving+mask pipeline forward 不 crash、cond_channels 形狀/placement 正確、`clear_conditioning` 剝除 |
|
||||
| 2.8 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | 仍暫緩（Phase 1 divergence，需獨立 reference token group） |
|
||||
|
||||
## Phase 3 — 訓練整合（ltx-trainer）
|
||||
|
||||
|
||||
@@ -10,7 +10,9 @@ from ltx_core.conditioning.types import (
|
||||
VideoConditionByKeyframeIndex,
|
||||
VideoConditionByLatentIndex,
|
||||
VideoConditionByMask,
|
||||
VideoConditionByMaskChannels,
|
||||
VideoConditionByReferenceLatent,
|
||||
encode_mask_channels,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
@@ -23,5 +25,7 @@ __all__ = [
|
||||
"VideoConditionByKeyframeIndex",
|
||||
"VideoConditionByLatentIndex",
|
||||
"VideoConditionByMask",
|
||||
"VideoConditionByMaskChannels",
|
||||
"VideoConditionByReferenceLatent",
|
||||
"encode_mask_channels",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Helpers for in-context conditioning channels (``LatentState.cond_channels``).
|
||||
|
||||
Conditioning channels are extra per-token input features (e.g. SCAIL-2 mask
|
||||
channels) that ride alongside the latent in patchified token space ``(B, T, C)``
|
||||
and are concatenated onto the latent before the model's first projection. When
|
||||
present they must stay length-aligned with the token sequence, so any
|
||||
conditioning item that *appends* tokens must also extend the channels.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.types import LatentState
|
||||
|
||||
|
||||
def extend_cond_channels(latent_state: LatentState, num_new_tokens: int) -> torch.Tensor | None:
|
||||
"""Return ``cond_channels`` extended with ``num_new_tokens`` zero rows, or ``None``.
|
||||
|
||||
Appended tokens (reference / driving / keyframe) carry no in-context signal by default, so they
|
||||
are padded with zeros to preserve the ``T``-alignment invariant. Returns ``None`` unchanged when
|
||||
the state has no conditioning channels (the standard case), keeping non-SCAIL pipelines untouched.
|
||||
"""
|
||||
cond_channels = latent_state.cond_channels
|
||||
if cond_channels is None:
|
||||
return None
|
||||
zeros = cond_channels.new_zeros(cond_channels.shape[0], num_new_tokens, cond_channels.shape[2])
|
||||
return torch.cat([cond_channels, zeros], dim=1)
|
||||
@@ -4,6 +4,7 @@ from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningI
|
||||
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_channels_cond import VideoConditionByMaskChannels, encode_mask_channels
|
||||
from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
|
||||
from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
|
||||
from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
|
||||
@@ -16,5 +17,7 @@ __all__ = [
|
||||
"VideoConditionByKeyframeIndex",
|
||||
"VideoConditionByLatentIndex",
|
||||
"VideoConditionByMask",
|
||||
"VideoConditionByMaskChannels",
|
||||
"VideoConditionByReferenceLatent",
|
||||
"encode_mask_channels",
|
||||
]
|
||||
|
||||
@@ -28,6 +28,7 @@ from enum import Enum
|
||||
import torch
|
||||
|
||||
from ltx_core.components.patchifiers import get_pixel_coords
|
||||
from ltx_core.conditioning.cond_channels import extend_cond_channels
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.conditioning.mask_utils import update_attention_mask
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
@@ -157,4 +158,5 @@ class VideoConditionByDrivingLatent(ConditioningItem):
|
||||
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,
|
||||
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
|
||||
)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.components.patchifiers import get_pixel_coords
|
||||
from ltx_core.conditioning.cond_channels import extend_cond_channels
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.conditioning.mask_utils import update_attention_mask
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
@@ -81,4 +82,5 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
|
||||
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,
|
||||
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
"""SCAIL-2 in-context mask conditioning (mechanism 2) for LTX-2.
|
||||
|
||||
Encodes ``K+1`` semantic pixel-space masks (1 environment switch + ``K`` character
|
||||
binding slots) into per-token conditioning channels and writes them onto the
|
||||
sequence via :attr:`LatentState.cond_channels`. Following the paper, mask signals
|
||||
are carried by the *driving* (and reference) tokens while the noisy target keeps
|
||||
an all-zero mask -- so by default the channels are written onto the last ``N``
|
||||
tokens of the sequence (the appended driving group), which means **this item must
|
||||
be applied after the driving conditioning**.
|
||||
|
||||
Channel expansion (LTX-2 faithful port of the paper's ``4(K+1)``): each semantic
|
||||
mask is spatially downsampled to the latent grid and temporally stacked along the
|
||||
channel dimension by the VAE temporal factor ``t`` (8 for LTX-2, vs 4 for the
|
||||
paper's Wan-2.1 backbone), giving ``t*(K+1)`` channels -- ``8*7 = 56`` for the
|
||||
default ``K=6``. The model's ``patchify_proj`` must be built with a matching
|
||||
``mask_conditioning_channels`` (see :class:`ltx_core.model.transformer.model.LTXModel`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
from ltx_core.types import LatentState
|
||||
|
||||
|
||||
def encode_mask_channels(
|
||||
masks: torch.Tensor,
|
||||
temporal_factor: int,
|
||||
height_lat: int,
|
||||
width_lat: int,
|
||||
frames_lat: int,
|
||||
) -> torch.Tensor:
|
||||
"""Encode ``[B, K+1, F_pix, H_pix, W_pix]`` masks into ``[B, t*(K+1), F_lat, H_lat, W_lat]``.
|
||||
|
||||
Each semantic mask is area-downsampled to the latent spatial grid, then the pixel frames that
|
||||
map to one latent frame are stacked along the channel dimension (temporal factor ``t``). The
|
||||
causal first latent frame corresponds to a single pixel frame, which is replicated across its
|
||||
``t`` stacked channels so every latent frame yields a uniform ``t`` channels per semantic class.
|
||||
"""
|
||||
b, s, f_pix, _h_pix, _w_pix = masks.shape
|
||||
t = temporal_factor
|
||||
expected_f_pix = (frames_lat - 1) * t + 1
|
||||
if f_pix != expected_f_pix:
|
||||
raise ValueError(
|
||||
f"mask pixel frames ({f_pix}) incompatible with latent frames ({frames_lat}) at temporal "
|
||||
f"factor {t}: expected (F_lat - 1) * {t} + 1 = {expected_f_pix}."
|
||||
)
|
||||
|
||||
# Spatial downsample every (semantic, frame) mask to the latent grid.
|
||||
flat = rearrange(masks.to(dtype=torch.float32), "b s f h w -> (b s f) 1 h w")
|
||||
down = F.interpolate(flat, size=(height_lat, width_lat), mode="area")
|
||||
down = rearrange(down, "(b s f) 1 h w -> b s f h w", b=b, s=s)
|
||||
|
||||
# Temporal stacking. Latent frame 0 = pixel frame 0 (causal), replicated across t channels;
|
||||
# latent frames 1.. group t consecutive pixel frames.
|
||||
first = down[:, :, :1].repeat(1, 1, t, 1, 1).unsqueeze(2) # [B, S, 1, t, H, W]
|
||||
rest = rearrange(down[:, :, 1:], "b s (fl t) h w -> b s fl t h w", t=t) # [B, S, F_lat-1, t, H, W]
|
||||
stacked = torch.cat([first, rest], dim=2) # [B, S, F_lat, t, H, W]
|
||||
|
||||
# Fold (semantic, temporal) into a single channel axis, semantic-major: channel = s * t + ti.
|
||||
return rearrange(stacked, "b s fl t h w -> b (s t) fl h w") # [B, t*(K+1), F_lat, H_lat, W_lat]
|
||||
|
||||
|
||||
class VideoConditionByMaskChannels(ConditioningItem):
|
||||
"""Write SCAIL-2 in-context mask channels onto the driving/reference tokens.
|
||||
|
||||
Args:
|
||||
masks: Pixel-space semantic masks ``[B, K+1, F_pix, H_pix, W_pix]``. Channel 0 is the
|
||||
environment switch (whether the environment comes from the reference vs the driving
|
||||
video); channels ``1..K`` are the character binding slots (regions sharing a slot share
|
||||
motion). ``F_pix`` must equal ``(F_lat - 1) * temporal_factor + 1``.
|
||||
applies_to_last_n: Number of trailing tokens to write the mask onto. ``None`` (default) uses
|
||||
the target token count, i.e. the appended driving group when this item runs right after
|
||||
the driving conditioning. The noisy target tokens keep an all-zero mask (paper-faithful).
|
||||
"""
|
||||
|
||||
def __init__(self, masks: torch.Tensor, applies_to_last_n: int | None = None):
|
||||
self.masks = masks
|
||||
self.applies_to_last_n = applies_to_last_n
|
||||
|
||||
def apply_to(self, latent_state: LatentState, latent_tools: VideoLatentTools) -> LatentState:
|
||||
shape = latent_tools.target_shape
|
||||
cond = encode_mask_channels(
|
||||
masks=self.masks,
|
||||
temporal_factor=latent_tools.scale_factors.time,
|
||||
height_lat=shape.height,
|
||||
width_lat=shape.width,
|
||||
frames_lat=shape.frames,
|
||||
)
|
||||
tokens = latent_tools.patchifier.patchify(cond) # [B, T_region, C_mask]
|
||||
b, n_region, c_mask = tokens.shape
|
||||
|
||||
total_tokens = latent_state.latent.shape[1]
|
||||
n = self.applies_to_last_n if self.applies_to_last_n is not None else n_region
|
||||
if n != n_region:
|
||||
raise ValueError(
|
||||
f"applies_to_last_n ({n}) must equal the encoded mask token count ({n_region})."
|
||||
)
|
||||
if n > total_tokens:
|
||||
raise ValueError(
|
||||
f"cannot write {n} mask tokens onto a sequence of only {total_tokens} tokens."
|
||||
)
|
||||
|
||||
cond_channels = latent_state.cond_channels
|
||||
if cond_channels is None:
|
||||
cond_channels = tokens.new_zeros(b, total_tokens, c_mask)
|
||||
else:
|
||||
if cond_channels.shape[2] != c_mask:
|
||||
raise ValueError(
|
||||
f"existing cond_channels width ({cond_channels.shape[2]}) != mask channels ({c_mask})."
|
||||
)
|
||||
cond_channels = cond_channels.clone()
|
||||
|
||||
start = total_tokens - n
|
||||
cond_channels[:, start : start + n] = tokens.to(dtype=cond_channels.dtype)
|
||||
return replace(latent_state, cond_channels=cond_channels)
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.conditioning.cond_channels import extend_cond_channels
|
||||
from ltx_core.conditioning.mask_utils import update_attention_mask
|
||||
from ltx_core.tools import LatentTools
|
||||
from ltx_core.types import LatentState
|
||||
@@ -56,4 +57,5 @@ class AudioConditionByReferenceLatent:
|
||||
positions=torch.cat([latent_state.positions, self.positions], dim=2),
|
||||
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
|
||||
attention_mask=new_attention_mask,
|
||||
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
|
||||
)
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.components.patchifiers import get_pixel_coords
|
||||
from ltx_core.conditioning.cond_channels import extend_cond_channels
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.conditioning.mask_utils import update_attention_mask
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
@@ -99,4 +100,5 @@ class VideoConditionByReferenceLatent(ConditioningItem):
|
||||
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,
|
||||
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Checkpoint surgery to enable SCAIL-2 in-context mask channels on a trained LTX-2 model.
|
||||
|
||||
Widening ``LTXModel.mask_conditioning_channels`` from 0 to ``C`` grows the video
|
||||
``patchify_proj`` weight from ``[inner, in]`` to ``[inner, in + C]``. This helper
|
||||
appends ``C`` **zero** input columns to that weight so a converted checkpoint is
|
||||
numerically identical to the original until the new columns are finetuned -- the
|
||||
extra channels contribute nothing at load time.
|
||||
|
||||
Usage::
|
||||
|
||||
sd = load_state_dict(path)
|
||||
widen_patchify_proj_for_mask_channels(sd, mask_channels=56)
|
||||
model = LTXModel(..., mask_conditioning_channels=56)
|
||||
model.load_state_dict(sd)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
_VIDEO_PROJ_SUFFIX = "patchify_proj.weight"
|
||||
_AUDIO_PROJ_SUFFIX = "audio_patchify_proj.weight"
|
||||
|
||||
|
||||
def widen_patchify_proj_for_mask_channels(
|
||||
state_dict: dict[str, torch.Tensor],
|
||||
mask_channels: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""In place, append ``mask_channels`` zero input columns to the video ``patchify_proj`` weight.
|
||||
|
||||
Only the video projection is touched (audio ``audio_patchify_proj`` is left alone). The bias,
|
||||
if present, is unchanged. Returns the same dict for convenience. Idempotency is the caller's
|
||||
responsibility -- calling twice widens twice.
|
||||
"""
|
||||
if mask_channels < 0:
|
||||
raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
|
||||
if mask_channels == 0:
|
||||
return state_dict
|
||||
|
||||
target_keys = [
|
||||
key
|
||||
for key in state_dict
|
||||
if key.endswith(_VIDEO_PROJ_SUFFIX) and not key.endswith(_AUDIO_PROJ_SUFFIX)
|
||||
]
|
||||
if not target_keys:
|
||||
raise KeyError(
|
||||
f"No '{_VIDEO_PROJ_SUFFIX}' weight found in the state dict; cannot widen for mask channels."
|
||||
)
|
||||
|
||||
for key in target_keys:
|
||||
weight = state_dict[key] # [inner_dim, in_features]
|
||||
if weight.ndim != 2:
|
||||
raise ValueError(f"Expected 2-D weight for '{key}', got shape {tuple(weight.shape)}")
|
||||
pad = weight.new_zeros(weight.shape[0], mask_channels)
|
||||
state_dict[key] = torch.cat([weight, pad], dim=1)
|
||||
|
||||
return state_dict
|
||||
@@ -38,6 +38,9 @@ class Modality:
|
||||
attention. ``None`` means unrestricted (full) attention between
|
||||
all tokens. Built incrementally by conditioning items; see
|
||||
:class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
|
||||
cond_channels: Optional in-context conditioning channels, shape ``(B, T, C_cond)``, aligned token-for-token
|
||||
with ``latent``. Concatenated onto ``latent`` before the first projection (see
|
||||
:meth:`TransformerArgsPreprocessor.prepare`); never noised. ``None`` for standard models.
|
||||
"""
|
||||
|
||||
latent: (
|
||||
@@ -53,6 +56,7 @@ class Modality:
|
||||
enabled: bool = True
|
||||
context_mask: torch.Tensor | None = None
|
||||
attention_mask: torch.Tensor | None = None
|
||||
cond_channels: torch.Tensor | None = None
|
||||
|
||||
def split(self, sizes: list[int]) -> list[Modality]:
|
||||
"""Split along the batch dimension into chunks of the given sizes."""
|
||||
|
||||
@@ -73,6 +73,7 @@ class LTXModel(torch.nn.Module):
|
||||
caption_projection: torch.nn.Module | None = None,
|
||||
audio_caption_projection: torch.nn.Module | None = None,
|
||||
cross_attention_adaln: bool = False,
|
||||
mask_conditioning_channels: int = 0,
|
||||
):
|
||||
super().__init__()
|
||||
# Log the attention backends this transformer is built with. Reading the resolved
|
||||
@@ -86,6 +87,10 @@ class LTXModel(torch.nn.Module):
|
||||
)
|
||||
self._enable_gradient_checkpointing = False
|
||||
self.cross_attention_adaln = cross_attention_adaln
|
||||
# Extra per-token input channels (e.g. SCAIL-2 in-context mask channels) concatenated onto the
|
||||
# video latent before ``patchify_proj``. 0 keeps the standard input width; >0 widens the first
|
||||
# projection. Fed via ``Modality.cond_channels`` (see TransformerArgsPreprocessor.prepare).
|
||||
self.mask_conditioning_channels = mask_conditioning_channels
|
||||
self.use_middle_indices_grid = use_middle_indices_grid
|
||||
self.rope_type = rope_type
|
||||
self.double_precision_rope = double_precision_rope
|
||||
@@ -154,8 +159,11 @@ class LTXModel(torch.nn.Module):
|
||||
caption_projection: torch.nn.Module | None = None,
|
||||
) -> None:
|
||||
"""Initialize video-specific components."""
|
||||
# Video input components
|
||||
self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
|
||||
# Video input components. When ``mask_conditioning_channels > 0`` the first projection is
|
||||
# widened to also accept the in-context conditioning channels concatenated onto the latent.
|
||||
self.patchify_proj = torch.nn.Linear(
|
||||
in_channels + self.mask_conditioning_channels, self.inner_dim, bias=True
|
||||
)
|
||||
if caption_projection is not None:
|
||||
self.caption_projection = caption_projection
|
||||
|
||||
|
||||
@@ -201,12 +201,46 @@ class TransformerArgsPreprocessor:
|
||||
)
|
||||
return pe
|
||||
|
||||
def _apply_patchify_proj(self, modality: Modality) -> torch.Tensor:
|
||||
"""Project patchified latents, concatenating any in-context conditioning channels first.
|
||||
|
||||
When ``patchify_proj`` expects more input features than the latent provides (a model built
|
||||
with ``mask_conditioning_channels > 0``), the extra width is filled by ``modality.cond_channels``
|
||||
(per-token, never noised). If those channels are absent they default to zeros, so a widened
|
||||
model still runs and, with zero-initialized new projection columns, reproduces the base model's
|
||||
output exactly.
|
||||
"""
|
||||
latent = modality.latent
|
||||
expected = self.patchify_proj.in_features
|
||||
actual = latent.shape[-1]
|
||||
if expected != actual:
|
||||
missing = expected - actual
|
||||
if missing < 0:
|
||||
raise ValueError(
|
||||
f"patchify_proj expects {expected} input features but the latent already has {actual}."
|
||||
)
|
||||
cond_channels = modality.cond_channels
|
||||
if cond_channels is None:
|
||||
cond_channels = latent.new_zeros(latent.shape[0], latent.shape[1], missing)
|
||||
elif cond_channels.shape[-1] != missing:
|
||||
raise ValueError(
|
||||
f"cond_channels has {cond_channels.shape[-1]} channels but patchify_proj needs {missing} "
|
||||
f"extra input features (latent {actual} + cond {cond_channels.shape[-1]} != {expected})."
|
||||
)
|
||||
elif cond_channels.shape[1] != latent.shape[1]:
|
||||
raise ValueError(
|
||||
f"cond_channels token length {cond_channels.shape[1]} must match latent token length "
|
||||
f"{latent.shape[1]}."
|
||||
)
|
||||
latent = torch.cat([latent, cond_channels.to(dtype=latent.dtype)], dim=-1)
|
||||
return self.patchify_proj(latent)
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
modality: Modality,
|
||||
cross_modality: Modality | None = None, # noqa: ARG002
|
||||
) -> TransformerArgs:
|
||||
x = self.patchify_proj(modality.latent)
|
||||
x = self._apply_patchify_proj(modality)
|
||||
batch_size = x.shape[0]
|
||||
timestep, embedded_timestep = self._prepare_timestep(
|
||||
modality.timesteps, self.adaln, batch_size, modality.latent.dtype
|
||||
|
||||
@@ -75,6 +75,9 @@ class LatentTools(Protocol):
|
||||
clean_latent = latent_state.clean_latent[:, :num_tokens]
|
||||
denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
|
||||
positions = latent_state.positions[:, :, :num_tokens]
|
||||
cond_channels = (
|
||||
latent_state.cond_channels[:, :num_tokens] if latent_state.cond_channels is not None else None
|
||||
)
|
||||
|
||||
return LatentState(
|
||||
latent=latent,
|
||||
@@ -82,6 +85,7 @@ class LatentTools(Protocol):
|
||||
positions=positions,
|
||||
clean_latent=clean_latent,
|
||||
attention_mask=None,
|
||||
cond_channels=cond_channels,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -193,6 +193,11 @@ class LatentState:
|
||||
clean_latent: Initial state of the latent before denoising, may include conditioning latents.
|
||||
attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention,
|
||||
0 = no attention. None means full attention everywhere. Built incrementally by conditioning items.
|
||||
cond_channels: Optional in-context conditioning channels in patchified token space, shape (B, T, C_cond),
|
||||
aligned token-for-token with ``latent``. These extra per-token feature channels (e.g. SCAIL-2 mask
|
||||
channels) are concatenated onto the latent right before the model's first projection; they are never
|
||||
noised or denoised. ``None`` means no conditioning channels (standard models). When present, the token
|
||||
length T must always match ``latent``; token-appending conditioning items extend it with zeros.
|
||||
"""
|
||||
|
||||
latent: torch.Tensor
|
||||
@@ -200,6 +205,7 @@ class LatentState:
|
||||
positions: torch.Tensor
|
||||
clean_latent: torch.Tensor
|
||||
attention_mask: torch.Tensor | None = None
|
||||
cond_channels: torch.Tensor | None = None
|
||||
|
||||
def clone(self) -> "LatentState":
|
||||
return LatentState(
|
||||
@@ -208,4 +214,5 @@ class LatentState:
|
||||
positions=self.positions.clone(),
|
||||
clean_latent=self.clean_latent.clone(),
|
||||
attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None,
|
||||
cond_channels=self.cond_channels.clone() if self.cond_channels is not None else None,
|
||||
)
|
||||
|
||||
@@ -269,6 +269,7 @@ def modality_from_latent_state(
|
||||
context=context,
|
||||
context_mask=None,
|
||||
attention_mask=state.attention_mask,
|
||||
cond_channels=state.cond_channels,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user