diff --git a/docs/plan.md b/docs/plan.md index 6c57ff7..226c236 100644 --- a/docs/plan.md +++ b/docs/plan.md @@ -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 與凍結策略。 diff --git a/docs/tasks.md b/docs/tasks.md index 2fa0d7f..836abee 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -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) diff --git a/packages/ltx-core/src/ltx_core/conditioning/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/__init__.py index f229a1a..c3a40b2 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/__init__.py @@ -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", ] diff --git a/packages/ltx-core/src/ltx_core/conditioning/cond_channels.py b/packages/ltx-core/src/ltx_core/conditioning/cond_channels.py new file mode 100644 index 0000000..8c6ab48 --- /dev/null +++ b/packages/ltx-core/src/ltx_core/conditioning/cond_channels.py @@ -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) diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py index e94564c..2ea56de 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py @@ -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", ] diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py index 4fb7bc0..52a81c2 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py @@ -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]), ) diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py index ba009cd..5d1cb3c 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py @@ -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]), ) diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/mask_channels_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/mask_channels_cond.py new file mode 100644 index 0000000..257bc7f --- /dev/null +++ b/packages/ltx-core/src/ltx_core/conditioning/types/mask_channels_cond.py @@ -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) diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py index 6c63db5..2ec7ec1 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py @@ -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]), ) diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py index 51e3463..fea8d7b 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py @@ -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]), ) diff --git a/packages/ltx-core/src/ltx_core/model/transformer/mask_channels_checkpoint.py b/packages/ltx-core/src/ltx_core/model/transformer/mask_channels_checkpoint.py new file mode 100644 index 0000000..5d813a3 --- /dev/null +++ b/packages/ltx-core/src/ltx_core/model/transformer/mask_channels_checkpoint.py @@ -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 diff --git a/packages/ltx-core/src/ltx_core/model/transformer/modality.py b/packages/ltx-core/src/ltx_core/model/transformer/modality.py index d4d9018..e8e5bc8 100644 --- a/packages/ltx-core/src/ltx_core/model/transformer/modality.py +++ b/packages/ltx-core/src/ltx_core/model/transformer/modality.py @@ -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.""" diff --git a/packages/ltx-core/src/ltx_core/model/transformer/model.py b/packages/ltx-core/src/ltx_core/model/transformer/model.py index 3c2bcba..47a212b 100644 --- a/packages/ltx-core/src/ltx_core/model/transformer/model.py +++ b/packages/ltx-core/src/ltx_core/model/transformer/model.py @@ -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 diff --git a/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py b/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py index 86203b3..81e5d97 100644 --- a/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py +++ b/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py @@ -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 diff --git a/packages/ltx-core/src/ltx_core/tools.py b/packages/ltx-core/src/ltx_core/tools.py index ec1696e..a9ca219 100644 --- a/packages/ltx-core/src/ltx_core/tools.py +++ b/packages/ltx-core/src/ltx_core/tools.py @@ -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, ) diff --git a/packages/ltx-core/src/ltx_core/types.py b/packages/ltx-core/src/ltx_core/types.py index c9dac29..55d5024 100644 --- a/packages/ltx-core/src/ltx_core/types.py +++ b/packages/ltx-core/src/ltx_core/types.py @@ -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, ) diff --git a/packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py b/packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py index 7a6a70e..35ac440 100644 --- a/packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py +++ b/packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py @@ -269,6 +269,7 @@ def modality_from_latent_state( context=context, context_mask=None, attention_mask=state.attention_mask, + cond_channels=state.cond_channels, )