From baa6646fd1fc3ad029eae9ed73bcb33e18196c11 Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 09:06:35 +0800 Subject: [PATCH 1/7] Add SCAIL-2 driving-latent conditioning (Phase 1, inference PoC) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/plan.md | 61 +++++++ docs/tasks.md | 42 +++++ .../src/ltx_core/conditioning/__init__.py | 4 + .../ltx_core/conditioning/types/__init__.py | 3 + .../conditioning/types/driving_video_cond.py | 160 ++++++++++++++++++ 5 files changed, 270 insertions(+) create mode 100644 docs/plan.md create mode 100644 docs/tasks.md create mode 100644 packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py diff --git a/docs/plan.md b/docs/plan.md new file mode 100644 index 0000000..6c57ff7 --- /dev/null +++ b/docs/plan.md @@ -0,0 +1,61 @@ +# SCAIL-2 → LTX-2 移植計畫 + +將 SCAIL-2（arXiv:2606.10804，`zai-org/SCAIL-2`）的端到端角色動畫手法移植到 LTX-2。 + +> **注意**：SCAIL-2 原始實作建構於 **Wan 2.1**，非 LTX-2。座標慣例與架構需翻譯到 LTX 的資料流。 + +## 三個核心機制 + +1. **Driving latent 直接串接** — 把驅動影片 latent 直接接進 DiT token 序列（不經骨架/pose 中介），用 width 軸座標偏移 ΔW 讓 driving 座標跟主 video 座標分開。 +2. **In-context mask channel** — 疊加額外輸入 channel（1 個環境開關 + K=6 個角色綁定槽，展開為 `4(K+1)=28` channel）到模型輸入，讓模型知道背景/角色對應。 +3. **Mode-specific RoPE** — Animation Mode 與 Replacement Mode 用不同的座標指派規則。 + +## LTX-2 對應落點 + +| 機制 | LTX-2 落點 | 改權重？ | +|---|---|---| +| 1. Driving 串接 + ΔW | 新 `ConditioningItem`（clone `reference_video_cond.py`），改 `positions` 偏移 | 否 | +| 3. Mode-specific RoPE | 上述 item 加 `DrivingMode` enum | 否 | +| 2. In-context mask channel | 加寬 `patchify_proj` 輸入 channel + `LatentState` 帶額外 channel + 投影前 concat | **是**（需微調） | + +**關鍵洞察**：LTX 的 `ConditioningItem.apply_to(latent_state, latent_tools) -> LatentState` 就是「串接進序列」的天然注入點，機制 1+3 **完全不用動 transformer / rope.py**。RoPE 由 `LatentState.positions` `[B,3,T,2]`（pixel 座標，axis1=(time,h,w)）驅動；ΔW = 對 `positions[:,2]`（width）加常數。width 正規化上界 `max_pos[2]=2048`，超過會 wrap。 + +## SCAIL-2 座標規則（論文） + +序列 `[z_ref; z_t; z_driv]`，driving 永遠在 width 軸帶固定偏移 ΔW： + +| | Animation | Replacement | +|---|---|---| +| z_ref | T=0, H=[0,Hv), W=[0,Wv) | T=0, **H=[ΔH_ref, ΔH_ref+Hv)**, W=[0,Wv) | +| z_t | T=[1,Tv], H=[0,Hv), W=[0,Wv) | T=[0,Tv−1], H=[0,Hv), W=[0,Wv) | +| z_driv | T=[1,Tv], H=[0,Hv), **W=[ΔW, ΔW+Wv)** | T=[0,Tv−1], H=[0,Hv), **W=[ΔW, ΔW+Wv)** | + +## 分階段計畫 + +### Phase 1 — Driving 串接 item（機制 1+3，推論期 PoC）✅ 已完成 +- 純推論、不改權重、不動 `patchify_proj`、不碰 ltx-trainer。 +- 用現有 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 3 — 訓練整合（ltx-trainer）⬜ 未開始 +- dataset 產出 (target, driving, mask)，接上 Phase 1/2 conditioning，設微調 loss 與凍結策略。 + +### Phase 4 — Pipeline + CLI 包裝 ⬜ 未開始 +- 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。 + +## Phase 1 簡化取捨（記錄，Phase 2 需回頭處理） +- (a) driving 時間座標直接複製 target 的（token-wise），故 driving 需與 target 同 F/H/W。 +- (b) 單一 frozen driving group 的 `attention_mask` 維持 None（= 全連接，target 完全看得到 driving），與 reference cond 一致。 +- (c) **ANIMATION 與 REPLACEMENT 在 Phase 1 產生相同 driving 座標** — mode 差異（z_ref 的 ΔH_ref 高度位移、target 時間原點、mask channel）屬 Phase 2，enum 先保留佔位。 + +## 參考 +- 論文：arXiv:2606.10804 — *SCAIL-2: Unifying Controlled Character Animation with End-to-end In-Context Conditioning* +- 官方實作：`zai-org/SCAIL-2`（GitHub / HuggingFace），建構於 Wan 2.1 +- 藍本檔案：`packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py` diff --git a/docs/tasks.md b/docs/tasks.md new file mode 100644 index 0000000..2fa0d7f --- /dev/null +++ b/docs/tasks.md @@ -0,0 +1,42 @@ +# SCAIL-2 → LTX-2 移植任務追蹤 + +狀態圖例：✅ 完成 ｜ 🔄 進行中 ｜ ⬜ 未開始 ｜ ⏸️ 暫緩 + +相關計畫見 [`plan.md`](./plan.md)。 + +## Phase 1 — Driving 串接 item（推論期 PoC） + +| # | 任務 | 狀態 | 產出 / 備註 | +|---|---|---|---| +| 1.1 | 新增 `VideoConditionByDrivingLatent` + `DrivingMode` | ✅ | `packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py`，以 `reference_video_cond.py` 為藍本，實作 ΔW width 偏移、時間對齊 target、`max_pos` 防呆、token 數防呆 | +| 1.2 | 匯出新 conditioning 類別 | ✅ | `conditioning/types/__init__.py`、`conditioning/__init__.py` | +| 1.3 | 免-GPU 資料流驗證腳本 | ✅ | 序列長度、ΔW 不重疊、時間對齊、frozen denoise_mask、attention_mask=None、`clear_conditioning` 剝除、超界/token 數防呆 — 全通過；ruff clean | +| 1.4 | （選配）端到端 smoke run | ✅ | 本機 CPU-only、無 checkpoint，改用小型真實 `LTXModel`（隨機權重）跑 pipeline 實走路徑：`create_noised_state`（含 driving cond）→ `modality_from_latent_state` → **真 transformer forward**（seq 160=target 80+driving 80）→ `clear_conditioning`（→80）。ANIMATION/REPLACEMENT 皆通過：不 crash、輸出 finite、driving frozen 且正確剝除。**僅驗證整合，不評估畫質** | + +## Phase 2 — In-context mask channel（checkpoint 手術） + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 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 暫緩項 | + +## Phase 3 — 訓練整合（ltx-trainer） + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 3.1 | dataset 產出 (target, driving, mask) | ⬜ | | +| 3.2 | 接上 conditioning + 微調 loss / 凍結策略 | ⬜ | | + +## Phase 4 — Pipeline + CLI + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 4.1 | `scail_animation.py` pipeline + arg parser | ⬜ | 仿 `lipdub.py` | + +## 決議紀錄 +- **範圍**：先只做 Phase 1（推論期 PoC）。Phase 2+ 待 Phase 1 驗證後再討論。 +- **架構前提**：SCAIL-2 建構於 Wan 2.1，本移植為跨架構移植。 diff --git a/packages/ltx-core/src/ltx_core/conditioning/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/__init__.py index 8d3eb9f..f229a1a 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/__init__.py @@ -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", 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 bd8d962..e94564c 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py @@ -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", 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 new file mode 100644 index 0000000..4fb7bc0 --- /dev/null +++ b/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py @@ -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, + ) -- 2.34.1 From 110adc781edbaea5ad22bba5c41102b18091ef2e Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 09:50:27 +0800 Subject: [PATCH 2/7] 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 --- docs/plan.md | 13 +- docs/tasks.md | 16 ++- .../src/ltx_core/conditioning/__init__.py | 4 + .../ltx_core/conditioning/cond_channels.py | 28 ++++ .../ltx_core/conditioning/types/__init__.py | 3 + .../conditioning/types/driving_video_cond.py | 2 + .../conditioning/types/keyframe_cond.py | 2 + .../conditioning/types/mask_channels_cond.py | 122 ++++++++++++++++++ .../types/reference_audio_cond.py | 2 + .../types/reference_video_cond.py | 2 + .../transformer/mask_channels_checkpoint.py | 57 ++++++++ .../ltx_core/model/transformer/modality.py | 4 + .../src/ltx_core/model/transformer/model.py | 12 +- .../model/transformer/transformer_args.py | 36 +++++- packages/ltx-core/src/ltx_core/tools.py | 4 + packages/ltx-core/src/ltx_core/types.py | 7 + .../src/ltx_pipelines/utils/helpers.py | 1 + 17 files changed, 300 insertions(+), 15 deletions(-) create mode 100644 packages/ltx-core/src/ltx_core/conditioning/cond_channels.py create mode 100644 packages/ltx-core/src/ltx_core/conditioning/types/mask_channels_cond.py create mode 100644 packages/ltx-core/src/ltx_core/model/transformer/mask_channels_checkpoint.py 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, ) -- 2.34.1 From e03cc62548ef91beac3b2f153a9de95d9950c2b1 Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 10:09:40 +0800 Subject: [PATCH 3/7] Add SCAIL-2 training integration (Phase 3, ltx-trainer FlexibleStrategy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire SCAIL-2 driving + in-context mask conditioning into the trainer via the unified FlexibleStrategy, so the widened patchify_proj (Phase 2) can be trained. - flexible.py: new DrivingConditionConfig (driving-latent concat with a RoPE width offset ΔW) and MaskChannelsConditionConfig (semantic masks -> per-token channels), added to the condition union and get_data_sources. Driving is prepended (cond-first, target stays at the tail for loss slicing); mask channels are written onto the driving tokens via Modality.cond_channels, reusing ltx-core encode_mask_channels. The noisy target keeps a zero mask. - model_loader.load_transformer gains mask_conditioning_channels, widening the video patchify_proj with zero-init columns via a new live-module helper (widen_module_patchify_proj_for_mask_channels). ModelConfig exposes the field. - trainer unfreezes patchify_proj in LoRA mode when mask channels are active (the new input columns are new base params LoRA cannot reach). - configs/scail_animation_lora.yaml plus README / training-modes table rows. Verified on CPU (verify_phase3_trainer.py): config round-trips, prepare_training _inputs builds cond_channels [B,T,56] with the mask on the driving tokens, a tiny widened model forwards and compute_loss returns a finite [B] loss, and the widen helper is output-preserving at zero init. Real training needs Linux+GPU+checkpoint; dataset preprocessing (driving latents + semantic masks) and validation-runner wiring are left for later. Co-Authored-By: Claude Opus 4.8 --- docs/plan.md | 8 +- docs/tasks.md | 13 +- .../transformer/mask_channels_checkpoint.py | 35 ++++ packages/ltx-trainer/configs/README.md | 5 + .../configs/scail_animation_lora.yaml | 132 ++++++++++++++ packages/ltx-trainer/docs/training-modes.md | 1 + .../ltx-trainer/src/ltx_trainer/config.py | 8 + .../src/ltx_trainer/model_loader.py | 14 +- .../ltx-trainer/src/ltx_trainer/trainer.py | 17 ++ .../training_strategies/flexible.py | 172 +++++++++++++++++- 10 files changed, 398 insertions(+), 7 deletions(-) create mode 100644 packages/ltx-trainer/configs/scail_animation_lora.yaml diff --git a/docs/plan.md b/docs/plan.md index 226c236..1161ce7 100644 --- a/docs/plan.md +++ b/docs/plan.md @@ -45,8 +45,12 @@ - `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 與凍結策略。 +### Phase 3 — 訓練整合（ltx-trainer）✅ 已完成（程式碼路徑；實訓需 GPU） +- **決策**：完整整合到 `FlexibleStrategy`；訓練 = LoRA + 解凍 `patchify_proj`（新 mask 欄位無法純 LoRA 訓練）。 +- 新 `DrivingConditionConfig` + `MaskChannelsConditionConfig`（`flexible.py`）；`_apply_driving_condition`(cond-first concat + ΔW) + `_build_mask_channels`(重用 `encode_mask_channels`) → `Modality.cond_channels`。 +- `load_transformer(mask_conditioning_channels=)` 用 `widen_module_patchify_proj_for_mask_channels` 加寬；`ModelConfig.mask_conditioning_channels`；trainer 在 LoRA 模式解凍 patchify_proj。 +- `configs/scail_animation_lora.yaml` + docs。CPU 單元驗證通過（`verify_phase3_trainer.py`）。 +- **本機無 GPU/Linux/checkpoint → 未跑實機訓練**；dataset 前處理（driving latents + 語意 mask）與 validation runner 接線未做。 ### Phase 4 — Pipeline + CLI 包裝 ⬜ 未開始 - 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。 diff --git a/docs/tasks.md b/docs/tasks.md index 836abee..fac85e4 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -30,10 +30,19 @@ ## Phase 3 — 訓練整合（ltx-trainer） +> **決策**：完整整合到 `FlexibleStrategy`；訓練方式 = **LoRA + 解凍 patchify_proj**（新 mask 欄位無法純 LoRA 訓練）。**本機無 GPU/Linux/checkpoint，只做到 CPU 單元驗證**，實訓需在 GPU 機器跑。 + | # | 任務 | 狀態 | 備註 | |---|---|---|---| -| 3.1 | dataset 產出 (target, driving, mask) | ⬜ | | -| 3.2 | 接上 conditioning + 微調 loss / 凍結策略 | ⬜ | | +| 3.1 | 新增 Driving/MaskChannels ConditionConfig | ✅ | `flexible.py`：`DrivingConditionConfig`(latents_dir/mode/width_offset) + `MaskChannelsConditionConfig`(mask_dir/num_slots)，加入 union + `get_data_sources` | +| 3.2 | strategy 接線 driving concat + cond_channels | ✅ | `_apply_driving_condition`(cond-first concat + ΔW width 偏移) + `_build_mask_channels`(重用 `encode_mask_channels`，寫前 N driving token) → `Modality.cond_channels` | +| 3.3 | model_loader + ModelConfig 支援加寬 | ✅ | `widen_module_patchify_proj_for_mask_channels`(live module zero-init) + `load_transformer(mask_conditioning_channels=)` + `ModelConfig.mask_conditioning_channels` | +| 3.4 | LoRA 模式解凍 patchify_proj | ✅ | `trainer._unfreeze_patchify_proj`：mask_channels>0 時把 video patchify_proj 設 trainable，讓新欄位隨 LoRA 一起訓 | +| 3.5 | 範例 config + docs | ✅ | `configs/scail_animation_lora.yaml`；`configs/README.md` 與 `docs/training-modes.md` 表格 row | +| 3.6 | CPU 單元驗證 | ✅ | `verify_phase3_trainer.py`：config round-trip、prepare_training_inputs 建 cond_channels[B,T,56]、driving 前置/mask placement、widened model forward + compute_loss finite、widen helper zero-init 等價 | +| 3.7 | dataset 前處理（產 driving latents + 語意 mask） | ⬜ | 需 process_dataset 產出 `driving_latents/`（同 target 形狀）與 `char_masks/`(mask=[K+1,F_pix,H,W])；語意 mask 需分割模型，屬資料工程，未做 | +| 3.8 | validation runner 接 driving/mask | ⬜ | 驗證期取樣尚未接 SCAIL 條件（config 內 validation 先停用），與 Phase 4 一起 | +| 3.9 | 實機訓練跑通 | ⬜ | 需 Linux + GPU + checkpoint，本機無法 | ## Phase 4 — Pipeline + CLI 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 index 5d813a3..d845535 100644 --- 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 @@ -22,6 +22,41 @@ _VIDEO_PROJ_SUFFIX = "patchify_proj.weight" _AUDIO_PROJ_SUFFIX = "audio_patchify_proj.weight" +def widen_module_patchify_proj_for_mask_channels(model: torch.nn.Module, mask_channels: int) -> torch.nn.Module: + """In place, widen a live ``LTXModel``'s video ``patchify_proj`` to accept ``mask_channels`` extra inputs. + + Replaces ``model.patchify_proj`` with a wider ``Linear`` whose original input columns are copied and + whose ``mask_channels`` new columns are zero (so the model reproduces its pre-widening output until + those columns are trained). Also sets ``model.mask_conditioning_channels`` for metadata/consistency. + Idempotent guard: raises if the model is already widened to a different value. + """ + if mask_channels < 0: + raise ValueError(f"mask_channels must be non-negative, got {mask_channels}") + + proj = getattr(model, "patchify_proj", None) + if not isinstance(proj, torch.nn.Linear): + raise AttributeError("model has no linear 'patchify_proj' to widen") + + already = int(getattr(model, "mask_conditioning_channels", 0)) + if mask_channels in (0, already): + model.mask_conditioning_channels = mask_channels + return model + if already != 0: + raise ValueError(f"patchify_proj already widened for {already} mask channels, refusing to re-widen") + + new_in = proj.in_features + mask_channels + wider = torch.nn.Linear(new_in, proj.out_features, bias=proj.bias is not None) + wider = wider.to(device=proj.weight.device, dtype=proj.weight.dtype) + with torch.no_grad(): + wider.weight.zero_() + wider.weight[:, : proj.in_features].copy_(proj.weight) + if proj.bias is not None: + wider.bias.copy_(proj.bias) + model.patchify_proj = wider + model.mask_conditioning_channels = mask_channels + return model + + def widen_patchify_proj_for_mask_channels( state_dict: dict[str, torch.Tensor], mask_channels: int, diff --git a/packages/ltx-trainer/configs/README.md b/packages/ltx-trainer/configs/README.md index 52dd414..899438b 100644 --- a/packages/ltx-trainer/configs/README.md +++ b/packages/ltx-trainer/configs/README.md @@ -23,5 +23,10 @@ adjust paths, dataset, and hyperparameters. | **Audio Inpainting** | — | Generated | `mask` | [`audio_inpainting_lora.yaml`](./audio_inpainting_lora.yaml) | | **A2A IC-LoRA** | — | Generated | `reference` | [`a2a_ic_lora.yaml`](./a2a_ic_lora.yaml) | | **AV2AV IC-LoRA** | Generated | Generated | `reference` (both) | [`av2av_ic_lora.yaml`](./av2av_ic_lora.yaml) | +| **SCAIL Animation** | Generated | — | `driving` + `mask_channels` | [`scail_animation_lora.yaml`](./scail_animation_lora.yaml) | The [`accelerate/`](./accelerate) directory holds the Accelerate launch configs (FSDP, DDP) for multi-GPU training. + +> **SCAIL Animation** (SCAIL-2 character animation) also sets `model.mask_conditioning_channels: 56` to widen the +> video `patchify_proj` for the in-context mask channels; in LoRA mode that projection is unfrozen so the new columns +> train. See [Training Modes Guide](../docs/training-modes.md). diff --git a/packages/ltx-trainer/configs/scail_animation_lora.yaml b/packages/ltx-trainer/configs/scail_animation_lora.yaml new file mode 100644 index 0000000..4824637 --- /dev/null +++ b/packages/ltx-trainer/configs/scail_animation_lora.yaml @@ -0,0 +1,132 @@ +# ============================================================================= +# LTX-2 SCAIL-2 Character Animation (LoRA + mask channels) Training Configuration +# ============================================================================= +# +# Trains SCAIL-2-style end-to-end character animation: a driving video latent is +# concatenated into the token sequence with a RoPE width offset (ΔW), and +# in-context mask channels (1 environment switch + K binding slots) are attached +# to the driving tokens to route motion per character. +# +# This combines LoRA on the attention/FFN blocks with an *unfrozen* widened +# patchify_proj (its new mask-channel input columns cannot be reached by LoRA and +# are trained directly). Set `model.mask_conditioning_channels` to the channel +# count = temporal_factor * (K + 1) = 8 * (6 + 1) = 56 for the default K=6. +# +# Dataset structure: +# preprocessed_data_root/ +# ├── latents/ # Target video latents (what the model generates) +# ├── conditions/ # Text embeddings for each video +# ├── driving_latents/ # Driving video latents (same F/H/W as target) +# └── char_masks/ # Semantic masks per sample, "mask" = [K+1, F_pix, H_pix, W_pix] +# # channel 0 = environment switch, 1..K = character binding slots +# +# ============================================================================= + +model: + model_path: "path/to/ltx-2-model.safetensors" + text_encoder_path: "path/to/gemma-text-encoder" + training_mode: "lora" + # Widen the video patchify_proj by 56 zero-init input columns (8 * (K+1), K=6). + # In LoRA mode the trainer additionally unfreezes patchify_proj so these train. + mask_conditioning_channels: 56 + load_checkpoint: null + +lora: + rank: 32 + alpha: 32 + dropout: 0.0 + target_modules: + - "attn1.to_k" + - "attn1.to_q" + - "attn1.to_v" + - "attn1.to_out.0" + - "attn2.to_k" + - "attn2.to_q" + - "attn2.to_v" + - "attn2.to_out.0" + - "ff.net.0.proj" + - "ff.net.2" + +training_strategy: + name: "flexible" + video: + is_generated: true + latents_dir: "latents" + conditions: + # SCAIL-2 driving conditioning: concatenate driving latents with a RoPE width + # offset so they stay spatially detached from the target tokens. + - type: driving + latents_dir: "driving_latents" + mode: "animation" # "animation" | "replacement" + width_offset: null # null = target pixel width (driving sits just to the right) + probability: 1.0 + # In-context mask channels attached to the driving tokens (target stays zero-mask). + - type: mask_channels + mask_dir: "char_masks" + num_slots: 6 # K binding slots; channels = 8 * (K + 1) = 56 (match model.mask_conditioning_channels) + +optimization: + learning_rate: 2e-4 + steps: 3000 + batch_size: 1 + gradient_accumulation_steps: 1 + max_grad_norm: 1.0 + optimizer_type: "adamw" + scheduler_type: "linear" + scheduler_params: { } + enable_gradient_checkpointing: true + +acceleration: + mixed_precision_mode: "bf16" + quantization: null + load_text_encoder_in_8bit: false + offload_optimizer_during_validation: false + +data: + preprocessed_data_root: "/path/to/preprocessed/data" + num_dataloader_workers: 2 + +validation: + # NOTE: SCAIL driving + mask-channel validation conditions are not yet wired into + # the validation runner (Phase 3 training path only). Keep validation minimal / + # disabled until the inference pipeline (Phase 4) lands. + samples: + - prompt: >- + A person performing an energetic dance routine, matching the motion of the driving + performer, with crisp footwork and expressive arm movements in a bright studio. + conditions: [] + negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted" + video_dims: [ 512, 512, 81 ] + frame_rate: 25.0 + seed: 42 + inference_steps: 30 + interval: null # disabled: driving/mask validation lands in Phase 4 + guidance_scale: 4.0 + stg_scale: 1.0 + stg_blocks: [29] + stg_mode: "stg_v" + generate_audio: false + skip_initial_validation: true + +checkpoints: + interval: 250 + keep_last_n: 3 + precision: "bfloat16" + +flow_matching: + timestep_sampling_mode: "shifted_logit_normal" + timestep_sampling_params: { } + +hub: + push_to_hub: false + hub_model_id: null + +wandb: + enabled: false + project: "ltx-2-trainer" + entity: null + tags: [ "ltx2", "scail-2", "character-animation" ] + log_validation_videos: true + +seed: 42 +output_dir: "outputs/scail_animation_lora" diff --git a/packages/ltx-trainer/docs/training-modes.md b/packages/ltx-trainer/docs/training-modes.md index 2469c77..8776cdd 100644 --- a/packages/ltx-trainer/docs/training-modes.md +++ b/packages/ltx-trainer/docs/training-modes.md @@ -40,6 +40,7 @@ Before diving into individual modes, here are the core ideas behind the flexible | **Audio Inpainting** | — | Generated | `mask` | [`audio_inpainting_lora`](../configs/audio_inpainting_lora.yaml) | | **A2A IC-LoRA** | — | Generated | `reference` | [`a2a_ic_lora`](../configs/a2a_ic_lora.yaml) | | **AV2AV IC-LoRA** | Generated | Generated | `reference` (both) | [`av2av_ic_lora`](../configs/av2av_ic_lora.yaml) | +| **SCAIL Animation** | Generated | — | `driving` + `mask_channels` | [`scail_animation_lora`](../configs/scail_animation_lora.yaml) | --- diff --git a/packages/ltx-trainer/src/ltx_trainer/config.py b/packages/ltx-trainer/src/ltx_trainer/config.py index c4d3009..ac065f6 100644 --- a/packages/ltx-trainer/src/ltx_trainer/config.py +++ b/packages/ltx-trainer/src/ltx_trainer/config.py @@ -252,6 +252,14 @@ class ModelConfig(ConfigBaseModel): description="Training mode - either LoRA fine-tuning or full model fine-tuning", ) + mask_conditioning_channels: int = Field( + default=0, + ge=0, + description="SCAIL-2 in-context mask channels. If > 0, the video patchify_proj is widened by this " + "many zero-init input columns at load time. Use with a 'driving' + 'mask_channels' condition and, " + "in LoRA mode, patchify_proj is additionally unfrozen so the new columns can train. 0 disables.", + ) + load_checkpoint: str | Path | None = Field( default=None, description="Path to a checkpoint file or directory to load from. " diff --git a/packages/ltx-trainer/src/ltx_trainer/model_loader.py b/packages/ltx-trainer/src/ltx_trainer/model_loader.py index f6aeba6..28426b3 100644 --- a/packages/ltx-trainer/src/ltx_trainer/model_loader.py +++ b/packages/ltx-trainer/src/ltx_trainer/model_loader.py @@ -50,27 +50,39 @@ def load_transformer( checkpoint_path: str | Path, device: Device = "cpu", dtype: torch.dtype = torch.bfloat16, + mask_conditioning_channels: int = 0, ) -> "LTXModel": """Load the LTX transformer model. Args: checkpoint_path: Path to the safetensors checkpoint file device: Device to load model on dtype: Data type for model weights + mask_conditioning_channels: If > 0, widen the video ``patchify_proj`` by this many zero-init + input columns after loading (SCAIL-2 in-context mask channels). The converted model + reproduces the base output exactly until those columns are trained. Returns: Loaded LTXModel transformer """ from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.transformer.mask_channels_checkpoint import ( + widen_module_patchify_proj_for_mask_channels, + ) from ltx_core.model.transformer.model_configurator import ( LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator, ) - return SingleGPUModelBuilder( + model = SingleGPUModelBuilder( model_path=str(checkpoint_path), model_class_configurator=LTXModelConfigurator, model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, ).build(device=_to_torch_device(device), dtype=dtype) + if mask_conditioning_channels > 0: + widen_module_patchify_proj_for_mask_channels(model, mask_conditioning_channels) + + return model + def load_video_vae_encoder( checkpoint_path: str | Path, diff --git a/packages/ltx-trainer/src/ltx_trainer/trainer.py b/packages/ltx-trainer/src/ltx_trainer/trainer.py index 8eb3fbc..79b58d1 100644 --- a/packages/ltx-trainer/src/ltx_trainer/trainer.py +++ b/packages/ltx-trainer/src/ltx_trainer/trainer.py @@ -399,6 +399,7 @@ class LtxvTrainer: checkpoint_path=self._config.model.model_path, device="cpu", dtype=torch.bfloat16, + mask_conditioning_channels=self._config.model.mask_conditioning_channels, ) # DDP-safe: LOCAL_RANK is set by accelerate before trainer init. Loading on bare @@ -440,9 +441,25 @@ class LtxvTrainer: else: raise ValueError(f"Unknown training mode: {self._config.model.training_mode}") + # SCAIL-2: the widened patchify_proj has new mask-channel input columns that LoRA cannot reach + # (they are new base parameters, not a low-rank delta on an existing weight). Unfreeze the whole + # patchify_proj so those columns train alongside the LoRA adapters. Harmless in full mode (already + # trainable). Placed before trainable-param collection so the params below pick it up. + if self._config.model.mask_conditioning_channels > 0: + self._unfreeze_patchify_proj() + self._trainable_params = [p for p in self._transformer.parameters() if p.requires_grad] logger.debug(f"Trainable params count: {sum(p.numel() for p in self._trainable_params):,}") + def _unfreeze_patchify_proj(self) -> None: + """Make the video ``patchify_proj`` parameters trainable (for SCAIL-2 mask-channel columns).""" + count = 0 + for name, param in self._transformer.named_parameters(): + if "patchify_proj" in name and "audio_patchify_proj" not in name: + param.requires_grad_(True) + count += param.numel() + logger.info(f"Unfroze video patchify_proj for mask-channel training ({count:,} params)") + def _init_timestep_sampler(self) -> None: """Initialize the timestep sampler based on the config.""" sampler_cls = SAMPLERS[self._config.flow_matching.timestep_sampling_mode] diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py index 82d0d9a..874096f 100644 --- a/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py +++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py @@ -15,6 +15,7 @@ import torch from pydantic import BaseModel, ConfigDict, Field, model_validator from torch import Tensor +from ltx_core.conditioning import encode_mask_channels from ltx_core.model.transformer.modality import Modality from ltx_trainer.timestep_samplers import TimestepSampler from ltx_trainer.training_strategies.base_strategy import ( @@ -107,6 +108,45 @@ class ReferenceConditionConfig(BaseModel): probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition") +class DrivingConditionConfig(BaseModel): + """SCAIL-2 driving-video conditioning (concatenation with a RoPE width offset). + Driving latents are concatenated to the sequence like a reference, but their RoPE width + coordinates are shifted by ``width_offset`` (the paper's ΔW) so they stay spatially detached + from the target tokens. Driving tokens are clean (timestep=0) and excluded from loss. + """ + + model_config = ConfigDict(extra="forbid") + + type: Literal["driving"] = "driving" + latents_dir: str = Field(..., description="Directory for driving-video latents (same F/H/W as target)") + mode: Literal["animation", "replacement"] = Field( + default="animation", + description="SCAIL mode. Reserved for the reference/height-shift distinction; driving-token " + "placement is currently identical for both (see ltx-core VideoConditionByDrivingLatent).", + ) + width_offset: float | None = Field( + default=None, + description="ΔW in RoPE pixel-space width units. None = target pixel width (driving sits just to the right).", + ) + probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition") + + +class MaskChannelsConditionConfig(BaseModel): + """SCAIL-2 in-context mask channels attached to the driving tokens. + Encodes ``num_slots + 1`` semantic pixel masks (1 environment switch + K binding slots) into + ``temporal_factor * (num_slots + 1)`` per-token conditioning channels (56 for K=6 on LTX-2) and + writes them onto the driving tokens via ``Modality.cond_channels``. The noisy target keeps an + all-zero mask (paper-faithful). Requires a ``driving`` condition and a model built with a matching + ``mask_conditioning_channels``. + """ + + model_config = ConfigDict(extra="forbid") + + type: Literal["mask_channels"] = "mask_channels" + mask_dir: str = Field(..., description="Directory of semantic masks [K+1, F_pix, H_pix, W_pix] per sample") + num_slots: int = Field(default=6, ge=1, description="Number of character binding slots K (channels = t*(K+1))") + + # Discriminated union for condition configs ConditionConfig = Annotated[ Union[ @@ -116,6 +156,8 @@ ConditionConfig = Annotated[ SpatialCropConditionConfig, MaskConditionConfig, ReferenceConditionConfig, + DrivingConditionConfig, + MaskChannelsConditionConfig, ], Field(discriminator="type"), ] @@ -200,9 +242,9 @@ class FlexibleStrategyConfig(TrainingStrategyConfigBase): if modality_config is None: continue for cond in modality_config.conditions: - if isinstance(cond, ReferenceConditionConfig): + if isinstance(cond, (ReferenceConditionConfig, DrivingConditionConfig)): sources[cond.latents_dir] = cond.latents_dir - elif isinstance(cond, MaskConditionConfig): + elif isinstance(cond, (MaskConditionConfig, MaskChannelsConditionConfig)): sources[cond.mask_dir] = cond.mask_dir return sources @@ -428,6 +470,37 @@ class FlexibleStrategy(TrainingStrategy): modality_key=modality_key, ) + # Step 5b: Apply SCAIL-2 driving conditioning (concatenation with a RoPE width offset). + # Driving tokens are prepended (cond-first, like reference) so the target stays at the tail + # for loss slicing; ``driving_token_count`` marks how many leading tokens are driving. + driving_token_count = 0 + for cond in modality_config.conditions: + if isinstance(cond, DrivingConditionConfig) and modality_key == "video": + noisy_latents, positions, timesteps, loss_mask, driving_token_count = self._apply_driving_condition( + noisy_latents=noisy_latents, + positions=positions, + timesteps=timesteps, + loss_mask=loss_mask, + target_width=data.width, + batch=batch, + config=cond, + ) + + # Step 5c: Build SCAIL-2 in-context mask channels on the driving tokens (target stays zero). + cond_channels = None + for cond in modality_config.conditions: + if isinstance(cond, MaskChannelsConditionConfig) and modality_key == "video": + cond_channels = self._build_mask_channels( + config=cond, + batch=batch, + total_tokens=noisy_latents.shape[1], + driving_token_count=driving_token_count, + target_frames=data.num_frames, + target_height=data.height, + target_width=data.width, + device=device, + ) + # Step 6: Build Modality modality = Modality( enabled=True, @@ -437,6 +510,7 @@ class FlexibleStrategy(TrainingStrategy): positions=positions, context=prompt_embeds, context_mask=prompt_attention_mask, + cond_channels=cond_channels, ) return ModalityProcessingResult( @@ -669,6 +743,100 @@ class FlexibleStrategy(TrainingStrategy): return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets + def _apply_driving_condition( + self, + noisy_latents: Tensor, + positions: Tensor, + timesteps: Tensor, + loss_mask: Tensor | None, + target_width: int, + batch: dict[str, Any], + config: DrivingConditionConfig, + ) -> tuple[Tensor, Tensor, Tensor, Tensor | None, int]: + """Prepend SCAIL-2 driving latents with a RoPE width offset (ΔW). + Driving latents share the target's frame/height/width, so their time and height coordinates + already match the target (same ``_get_video_positions`` grid); only the width axis is shifted + by ``width_offset`` so the driving tokens stay spatially detached. Driving tokens are clean + (timestep=0) and excluded from loss. Returns the driving token count for mask placement. + The apply/skip decision is batch-wide (concatenation changes the sequence length) but drawn + from the torch RNG for reproducibility, mirroring ``_apply_reference_condition``. + """ + if torch.rand((), device=noisy_latents.device).item() >= config.probability: + return noisy_latents, positions, timesteps, loss_mask, 0 + + cond = self._patchify_latent_data(batch[config.latents_dir], "video") + drv_latents = cond.latents + batch_size, drv_seq_len, _ = drv_latents.shape + device = drv_latents.device + dtype = drv_latents.dtype + + drv_positions = self._get_video_positions( + num_frames=cond.num_frames, + height=cond.height, + width=cond.width, + batch_size=batch_size, + fps=cond.fps, + device=device, + ).clone() + width_offset = ( + config.width_offset + if config.width_offset is not None + else float(target_width * VIDEO_SCALE_FACTORS.width) + ) + drv_positions[:, 2, ...] = drv_positions[:, 2, ...] + width_offset + + drv_timesteps = torch.zeros(batch_size, drv_seq_len, device=device, dtype=dtype) + drv_loss_mask = torch.zeros(batch_size, drv_seq_len, dtype=torch.bool, device=device) + + combined_latents = torch.cat([drv_latents, noisy_latents], dim=1) + combined_positions = torch.cat([drv_positions, positions], dim=2) + combined_timesteps = torch.cat([drv_timesteps, timesteps], dim=1) + combined_loss_mask = torch.cat([drv_loss_mask, loss_mask], dim=1) if loss_mask is not None else None + + return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, drv_seq_len + + def _build_mask_channels( + self, + config: MaskChannelsConditionConfig, + batch: dict[str, Any], + total_tokens: int, + driving_token_count: int, + target_frames: int, + target_height: int, + target_width: int, + device: torch.device, + ) -> Tensor: + """Encode semantic masks into per-token channels written onto the leading driving tokens. + The mask tensor at ``batch[mask_dir]["mask"]`` is expected as ``[B, K+1, F_pix, H_pix, W_pix]`` + (channel 0 = environment switch, ``1..K`` = binding slots). It is encoded to + ``temporal_factor * (K+1)`` per-token channels (56 for K=6) via ``encode_mask_channels`` and + placed on the driving tokens; the target tokens keep an all-zero mask (paper-faithful). + """ + masks = batch[config.mask_dir]["mask"].to(device=device) + encoded = encode_mask_channels( + masks=masks, + temporal_factor=VIDEO_SCALE_FACTORS.time, + height_lat=target_height, + width_lat=target_width, + frames_lat=target_frames, + ) + tokens = self._video_patchifier.patchify(encoded) # [B, N, C_mask] + batch_size, n_region, c_mask = tokens.shape + + cond_channels = tokens.new_zeros(batch_size, total_tokens, c_mask) + if driving_token_count == 0: + # No driving tokens (e.g. driving skipped by its probability): fall back to writing the + # mask onto the leading target tokens so the channel width still matches the model. + cond_channels[:, :n_region] = tokens + else: + if n_region != driving_token_count: + raise ValueError( + f"mask token count ({n_region}) must equal the driving token count " + f"({driving_token_count}); driving and mask must describe the same grid." + ) + cond_channels[:, :driving_token_count] = tokens + return cond_channels + @staticmethod def _compute_modality_loss(pred: Tensor, targets: Tensor, loss_mask: Tensor) -> Tensor: """Compute per-element MSE loss for a single modality. Returns [B,].""" -- 2.34.1 From 06c0870bbb2543a6633c7c9110f92424bda8bce4 Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 10:22:35 +0800 Subject: [PATCH 4/7] Add SCAIL-2 animation inference pipeline (Phase 4, ltx-pipelines) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two-stage distilled inference pipeline that animates a character from a driving video, wiring the Phase 1-3 SCAIL-2 conditioning into a runnable CLI. - scail_animation.py: ScailAnimationPipeline (mirrors distilled.py) plus a pure, testable build_scail_conditionings that assembles VideoConditionByDrivingLatent + VideoConditionByMaskChannels (driving appended, mask on the trailing driving tokens), and a load_masks helper. main() + scail_animation_arg_parser add --driving-video / --mask-path / --mode / --driving-strength on top of the standard two-stage distilled parser. - LTXModelConfigurator now reads config `mask_conditioning_channels`, so a SCAIL-trained checkpoint whose config declares it builds the widened patchify_proj automatically — no runtime widening wrapper needed at inference. - CLAUDE.md pipeline table row. Verified on CPU (verify_phase4_pipeline.py): the module imports, the CLI parses the SCAIL flags, build_scail_conditionings grows the sequence and places the mask channels on the driving tokens (target stays zero), driving-only leaves cond_channels None, and the configurator honors mask_conditioning_channels (patchify_proj widened from config). End-to-end runs still need a GPU and a SCAIL-trained checkpoint. Co-Authored-By: Claude Opus 4.8 --- docs/plan.md | 8 +- docs/tasks.md | 6 +- .../model/transformer/model_configurator.py | 2 + packages/ltx-pipelines/CLAUDE.md | 1 + .../src/ltx_pipelines/scail_animation.py | 349 ++++++++++++++++++ .../src/ltx_pipelines/utils/args.py | 37 ++ 6 files changed, 400 insertions(+), 3 deletions(-) create mode 100644 packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py diff --git a/docs/plan.md b/docs/plan.md index 1161ce7..ad197c9 100644 --- a/docs/plan.md +++ b/docs/plan.md @@ -52,8 +52,12 @@ - `configs/scail_animation_lora.yaml` + docs。CPU 單元驗證通過（`verify_phase3_trainer.py`）。 - **本機無 GPU/Linux/checkpoint → 未跑實機訓練**；dataset 前處理（driving latents + 語意 mask）與 validation runner 接線未做。 -### Phase 4 — Pipeline + CLI 包裝 ⬜ 未開始 -- 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。 +### Phase 4 — Pipeline + CLI 包裝 ✅ 已完成（程式碼路徑；實跑需 GPU） +- `LTXModelConfigurator` 讀 `mask_conditioning_channels` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime wrapper）。 +- `scail_animation.py`：`ScailAnimationPipeline`（仿 `distilled.py` 兩階段）+ 可測純函式 `build_scail_conditionings`（driving + mask 條件組裝）+ `load_masks` + `main()`。 +- `utils/args.py` `scail_animation_arg_parser`（`--driving-video/--mask-path/--mode/--driving-strength`）。 +- CPU 驗證通過（`verify_phase4_pipeline.py`）；`ltx-pipelines/CLAUDE.md` 表格 row。 +- **實機端到端需 GPU + SCAIL-trained checkpoint**（config 含 `mask_conditioning_channels` + 訓練好的 patchify_proj + LoRA）。 ## Phase 1 簡化取捨（記錄，Phase 2 需回頭處理） - (a) driving 時間座標直接複製 target 的（token-wise），故 driving 需與 target 同 F/H/W。 diff --git a/docs/tasks.md b/docs/tasks.md index fac85e4..bb21bfb 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -48,7 +48,11 @@ | # | 任務 | 狀態 | 備註 | |---|---|---|---| -| 4.1 | `scail_animation.py` pipeline + arg parser | ⬜ | 仿 `lipdub.py` | +| 4.1 | configurator 讀 `mask_conditioning_channels` | ✅ | `LTXModelConfigurator` 兩分支 `config.get("mask_conditioning_channels",0)` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime widen wrapper） | +| 4.2 | `ScailAnimationPipeline` + `build_scail_conditionings` | ✅ | `scail_animation.py`：仿 `distilled.py` 兩階段，注入 SCAIL driving+mask conditioning；`build_scail_conditionings` 為可測純函式；`load_masks` helper | +| 4.3 | `scail_animation_arg_parser` + `main()` | ✅ | `utils/args.py`：在 `default_2_stage_distilled_arg_parser` 上加 `--driving-video/--mask-path/--mode/--driving-strength` | +| 4.4 | CPU 驗證 + docs | ✅ | `verify_phase4_pipeline.py`：import、CLI round-trip、conditioning assembly(driving append + mask 在尾端)、configurator 自描述加寬；`ltx-pipelines/CLAUDE.md` pipeline 表格 row | +| 4.5 | 實機端到端跑通 | ⬜ | 需 GPU + SCAIL-trained checkpoint（含 `mask_conditioning_channels` config + 訓練好的 patchify_proj + LoRA） | ## 決議紀錄 - **範圍**：先只做 Phase 1（推論期 PoC）。Phase 2+ 待 Phase 1 驗證後再討論。 diff --git a/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py index fb47cef..59159ad 100644 --- a/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py +++ b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py @@ -69,6 +69,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]): caption_projection=caption_projection, audio_caption_projection=audio_caption_projection, cross_attention_adaln=config.get("cross_attention_adaln", False), + mask_conditioning_channels=config.get("mask_conditioning_channels", 0), ) @@ -120,6 +121,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]): apply_gated_attention=config.get("apply_gated_attention", False), caption_projection=caption_projection, cross_attention_adaln=config.get("cross_attention_adaln", False), + mask_conditioning_channels=config.get("mask_conditioning_channels", 0), ) diff --git a/packages/ltx-pipelines/CLAUDE.md b/packages/ltx-pipelines/CLAUDE.md index 2aa573d..39a004b 100644 --- a/packages/ltx-pipelines/CLAUDE.md +++ b/packages/ltx-pipelines/CLAUDE.md @@ -25,6 +25,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for | `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference | | `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control | | `LipDubPipeline` | `lipdub.py` | 2 | Distilled only | Euler | Lip dubbing with IC-LoRA + audio ref conditioning | +| `ScailAnimationPipeline` | `scail_animation.py` | 2 | Distilled + SCAIL LoRA | Euler | SCAIL-2 character animation from a driving video (driving latent concat + ΔW RoPE + in-context mask channels) | | `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration | ## Guidance diff --git a/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py b/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py new file mode 100644 index 0000000..d167e7c --- /dev/null +++ b/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py @@ -0,0 +1,349 @@ +"""SCAIL-2 character-animation inference pipeline. + +Two-stage distilled video generation that animates a character from a *driving* +video: the driving latent is concatenated into the token sequence with a RoPE +width offset (ΔW) and in-context mask channels route motion per character. This +mirrors the distilled pipeline (``distilled.py``) and lip-dub pipeline +(``lipdub.py``) structure, adding the SCAIL conditioning assembly. + +Requirements: +- A SCAIL-trained checkpoint whose config declares ``mask_conditioning_channels`` + (so ``LTXModelConfigurator`` builds the widened ``patchify_proj`` automatically) + plus the SCAIL LoRA adapter. +- A driving video (same target resolution) and a semantic mask tensor + ``[K+1, F_pix, H_pix, W_pix]`` (channel 0 = environment switch, 1..K = binding + slots), saved as a ``.pt`` file. + +The SCAIL-specific conditioning assembly is factored into +``build_scail_conditionings`` so it can be unit-tested without the heavy models. +""" + +from __future__ import annotations + +import logging +from collections.abc import Iterator + +import torch + +from ltx_core.components.noisers import GaussianNoiser +from ltx_core.conditioning import ( + ConditioningItem, + DrivingMode, + VideoConditionByDrivingLatent, + VideoConditionByMaskChannels, +) +from ltx_core.loader import LoraPathStrengthAndSDOps +from ltx_core.loader.registry import Registry +from ltx_core.model.transformer.compiling import CompilationConfig +from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number +from ltx_core.quantization import QuantizationPolicy +from ltx_core.types import Audio, VideoPixelShape +from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy +from ltx_pipelines.utils.args import ImageConditioningInput, resolve_cli_params, scail_animation_arg_parser +from ltx_pipelines.utils.blocks import ( + DiffusionStage, + ImageConditioner, + PromptEncoder, + VideoDecoder, + VideoUpsampler, +) +from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS +from ltx_pipelines.utils.denoisers import SimpleDenoiser +from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device +from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess +from ltx_pipelines.utils.types import ModalitySpec, OffloadMode + + +def build_scail_conditionings( + driving_latent: torch.Tensor, + masks: torch.Tensor | None, + *, + mode: DrivingMode = DrivingMode.ANIMATION, + strength: float = 1.0, + width_offset: float | None = None, +) -> list[ConditioningItem]: + """Assemble the SCAIL driving + mask-channel conditioning items (order matters). + + The driving item is applied first (it appends the driving tokens at the tail); the mask item is + applied second, writing its channels onto those trailing driving tokens (the noisy target keeps a + zero mask). ``masks`` may be ``None`` to run driving-only (e.g. a model without mask channels). + + Args: + driving_latent: Driving video latent ``[B, C, F, H, W]`` (same F/H/W as the target). + masks: Semantic masks ``[B, K+1, F_pix, H_pix, W_pix]`` or ``None``. + mode: SCAIL mode (animation/replacement). + strength: Driving conditioning strength (1.0 keeps it clean/frozen). + width_offset: ΔW in RoPE pixel-space width units (None = target pixel width). + """ + conditionings: list[ConditioningItem] = [ + VideoConditionByDrivingLatent( + latent=driving_latent, + mode=mode, + width_offset=width_offset, + strength=strength, + ) + ] + if masks is not None: + conditionings.append(VideoConditionByMaskChannels(masks=masks)) + return conditionings + + +def load_masks(mask_path: str, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Load a semantic mask tensor from a ``.pt`` file and shape it to ``[B, K+1, F_pix, H, W]``.""" + masks = torch.load(mask_path, map_location=device, weights_only=True) + if isinstance(masks, dict): + masks = masks["mask"] + if masks.dim() == 4: # [K+1, F, H, W] -> add batch + masks = masks.unsqueeze(0) + return masks.to(device=device, dtype=dtype) + + +class ScailAnimationPipeline: + """Two-stage distilled SCAIL-2 character animation (driving video + in-context mask channels).""" + + def __init__( + self, + distilled_checkpoint_path: str, + gemma_root: str, + spatial_upsampler_path: str, + scail_lora: LoraPathStrengthAndSDOps, + device: torch.device | None = None, + quantization: QuantizationPolicy | None = None, + registry: Registry | None = None, + compilation_config: CompilationConfig | None = None, + offload_mode: OffloadMode = OffloadMode.NONE, + alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM, + ) -> None: + self.device = device or get_device() + self.dtype = torch.bfloat16 + + self.prompt_encoder = PromptEncoder( + distilled_checkpoint_path, + gemma_root, + self.dtype, + self.device, + registry=registry, + offload_mode=offload_mode, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.image_conditioner = ImageConditioner( + distilled_checkpoint_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.stage = DiffusionStage.from_checkpoint( + distilled_checkpoint_path, + self.dtype, + self.device, + loras=(scail_lora,), + quantization=quantization, + registry=registry, + compilation_config=compilation_config, + offload_mode=offload_mode, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.upsampler = VideoUpsampler( + distilled_checkpoint_path, + spatial_upsampler_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.video_decoder = VideoDecoder( + distilled_checkpoint_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + + def _encode_driving_latent( + self, + driving_video_path: str, + height: int, + width: int, + num_frames: int, + video_encoder: VideoEncoder, + tiling_config: TilingConfig | None, + ) -> torch.Tensor: + """Decode + encode the driving video into a VAE latent at the target resolution.""" + frame_gen = decode_video_by_frame(path=driving_video_path, frame_cap=num_frames, device=self.device) + video = video_preprocess(frame_gen, height, width, self.dtype, self.device) + if tiling_config is not None: + return video_encoder.tiled_encode(video, tiling_config) + return video_encoder(video) + + def _video_conditionings( + self, + images: list[ImageConditioningInput], + driving_video_path: str, + masks: torch.Tensor | None, + height: int, + width: int, + num_frames: int, + mode: DrivingMode, + driving_strength: float, + video_encoder: VideoEncoder, + tiling_config: TilingConfig | None, + ) -> list[ConditioningItem]: + conditionings = combined_image_conditionings( + images=images, + height=height, + width=width, + video_encoder=video_encoder, + dtype=self.dtype, + device=self.device, + ) + driving_latent = self._encode_driving_latent( + driving_video_path, height, width, num_frames, video_encoder, tiling_config + ) + conditionings.extend( + build_scail_conditionings(driving_latent, masks, mode=mode, strength=driving_strength) + ) + return conditionings + + @torch.inference_mode() + def __call__( # noqa: PLR0913 + self, + prompt: str, + seed: int, + height: int, + width: int, + num_frames: int, + frame_rate: float, + driving_video_path: str, + mask_path: str | None, + images: list[ImageConditioningInput] | None = None, + mode: DrivingMode = DrivingMode.ANIMATION, + driving_strength: float = 1.0, + enhance_prompt: bool = False, + tiling_config: TilingConfig | None = None, + stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS, + stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS, + ) -> tuple[Iterator[torch.Tensor], Audio | None]: + assert_resolution(height=height, width=width, is_two_stage=True) + images = images or [] + + generator = torch.Generator(device=self.device).manual_seed(seed) + noiser = GaussianNoiser(generator=generator) + encode_tiling = TilingConfig.default() + + (ctx_p,) = self.prompt_encoder( + [prompt], + enhance_first_prompt=enhance_prompt, + enhance_prompt_image=images[0][0] if len(images) > 0 else None, + enhance_prompt_seed=seed, + ) + video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding + + masks = load_masks(mask_path, self.device, self.dtype) if mask_path is not None else None + + def build_video_conditionings(out_height: int, out_width: int) -> list[ConditioningItem]: + return self.image_conditioner( + lambda enc: self._video_conditionings( + images=images, + driving_video_path=driving_video_path, + masks=masks, + height=out_height, + width=out_width, + num_frames=num_frames, + mode=mode, + driving_strength=driving_strength, + video_encoder=enc, + tiling_config=encode_tiling, + ) + ) + + # Stage 1: low resolution. + stage_1_sigmas_t = stage_1_sigmas.to(dtype=torch.float32, device=self.device) + stage_1_conditionings = build_video_conditionings(height // 2, width // 2) + video_state, _ = self.stage( + denoiser=SimpleDenoiser(video_context, audio_context), + sigmas=stage_1_sigmas_t, + noiser=noiser, + width=width // 2, + height=height // 2, + frames=num_frames, + fps=frame_rate, + video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings), + audio=ModalitySpec(context=audio_context), + ) + + # Stage 2: upsample + refine. + upscaled = self.upsampler(video_state.latent[:1]) + stage_2_sigmas_t = stage_2_sigmas.to(dtype=torch.float32, device=self.device) + stage_2_conditionings = build_video_conditionings(height, width) + video_state, _ = self.stage( + denoiser=SimpleDenoiser(video_context, audio_context), + sigmas=stage_2_sigmas_t, + noiser=noiser, + width=width, + height=height, + frames=num_frames, + fps=frame_rate, + video=ModalitySpec( + context=video_context, + conditionings=stage_2_conditionings, + noise_scale=stage_2_sigmas_t[0].item(), + initial_latent=upscaled, + ), + audio=ModalitySpec(context=audio_context), + ) + + decoded_video = self.video_decoder(video_state.latent, tiling_config, generator) + return decoded_video, None + + +@torch.inference_mode() +def main() -> None: + logging.basicConfig(level=logging.INFO) + params = resolve_cli_params(distilled=True) + parser = scail_animation_arg_parser(params=params) + args = parser.parse_args() + + if not args.lora or len(args.lora) != 1: + raise ValueError("SCAIL animation requires exactly one --lora (the SCAIL adapter).") + + pipeline = ScailAnimationPipeline( + distilled_checkpoint_path=args.distilled_checkpoint_path, + gemma_root=args.gemma_root, + spatial_upsampler_path=args.spatial_upsampler_path, + scail_lora=args.lora[0], + quantization=args.quantization, + compilation_config=args.compile, + offload_mode=args.offload_mode, + ) + tiling_config = TilingConfig.default() + output_shape = VideoPixelShape( + batch=1, frames=args.num_frames, width=args.width, height=args.height, fps=args.frame_rate + ) + video_chunks_number = get_video_chunks_number(output_shape.frames, tiling_config) + video, _audio = pipeline( + prompt=args.prompt, + seed=args.seed, + height=args.height, + width=args.width, + num_frames=args.num_frames, + frame_rate=args.frame_rate, + driving_video_path=args.driving_video, + mask_path=args.mask_path, + images=args.images if hasattr(args, "images") else [], + mode=DrivingMode(args.mode), + driving_strength=args.driving_strength, + enhance_prompt=args.enhance_prompt, + tiling_config=tiling_config, + ) + encode_video( + video=video, + fps=int(args.frame_rate), + audio=None, + output_path=args.output_path, + video_chunks_number=video_chunks_number, + ) + + +if __name__ == "__main__": + main() diff --git a/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py b/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py index 405da47..5041cf0 100644 --- a/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py +++ b/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py @@ -809,3 +809,40 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS ), ) return parser + + +def scail_animation_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser: + """Argument parser for the SCAIL-2 character-animation pipeline (distilled, two-stage). + + Extends the standard two-stage distilled parser (checkpoint paths, --prompt, --num-frames, + --frame-rate, --height, --width, --image, --lora, --spatial-upsampler-path) with the SCAIL + driving/mask conditioning inputs. The single --lora is the SCAIL adapter. + """ + parser = default_2_stage_distilled_arg_parser(params=params) + parser.add_argument( + "--driving-video", + type=resolve_existing_path, + required=True, + help="Driving video file whose motion is transferred to the animated character.", + ) + parser.add_argument( + "--mask-path", + type=resolve_existing_path, + default=None, + help="Optional .pt file with semantic masks [K+1, F_pix, H, W] (ch0 = environment switch, " + "1..K = character binding slots). Omit for a driving-only (mask-free) model.", + ) + parser.add_argument( + "--mode", + type=str, + choices=["animation", "replacement"], + default="animation", + help="SCAIL conditioning mode (default: animation).", + ) + parser.add_argument( + "--driving-strength", + type=float, + default=1.0, + help="Driving conditioning strength; 1.0 keeps the driving latent clean/frozen (default: 1.0).", + ) + return parser -- 2.34.1 From a598f89d99bfba5462195054bc5d77597ac4620b Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 21:00:45 +0800 Subject: [PATCH 5/7] Fix ltx-trainer ruff target-version and a hidden lint finding The [tool.ruff] target-version was accidentally set to the package version "1.1.7", which made ruff fail to parse the whole package's pyproject and silently skip linting. Set it to "py310" to match requires-python >=3.10. With ruff working again it flagged a too-many-branches finding in the Phase 3 SCAIL wiring: extract the driving + mask-channel loops from _process_modality into a new _apply_scail_conditions helper. Behavior is unchanged (Phase 3 CPU verification still passes); full `ruff check .` on the trainer now passes. Co-Authored-By: Claude Opus 4.8 --- packages/ltx-trainer/pyproject.toml | 2 +- .../training_strategies/flexible.py | 92 +++++++++++++------ 2 files changed, 63 insertions(+), 31 deletions(-) diff --git a/packages/ltx-trainer/pyproject.toml b/packages/ltx-trainer/pyproject.toml index ce956d6..6c57086 100644 --- a/packages/ltx-trainer/pyproject.toml +++ b/packages/ltx-trainer/pyproject.toml @@ -54,7 +54,7 @@ build-backend = "hatchling.build" [tool.ruff] -target-version = "1.1.7" +target-version = "py310" line-length = 120 # Restrict isort first-party detection to src/ so stray dirs (e.g. wandb/ run output) # next to pyproject.toml don't get classified as first-party packages. See ruff#10519. diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py index 874096f..867cb42 100644 --- a/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py +++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py @@ -470,36 +470,18 @@ class FlexibleStrategy(TrainingStrategy): modality_key=modality_key, ) - # Step 5b: Apply SCAIL-2 driving conditioning (concatenation with a RoPE width offset). - # Driving tokens are prepended (cond-first, like reference) so the target stays at the tail - # for loss slicing; ``driving_token_count`` marks how many leading tokens are driving. - driving_token_count = 0 - for cond in modality_config.conditions: - if isinstance(cond, DrivingConditionConfig) and modality_key == "video": - noisy_latents, positions, timesteps, loss_mask, driving_token_count = self._apply_driving_condition( - noisy_latents=noisy_latents, - positions=positions, - timesteps=timesteps, - loss_mask=loss_mask, - target_width=data.width, - batch=batch, - config=cond, - ) - - # Step 5c: Build SCAIL-2 in-context mask channels on the driving tokens (target stays zero). - cond_channels = None - for cond in modality_config.conditions: - if isinstance(cond, MaskChannelsConditionConfig) and modality_key == "video": - cond_channels = self._build_mask_channels( - config=cond, - batch=batch, - total_tokens=noisy_latents.shape[1], - driving_token_count=driving_token_count, - target_frames=data.num_frames, - target_height=data.height, - target_width=data.width, - device=device, - ) + # Step 5b: Apply SCAIL-2 driving + in-context mask conditioning (video only). + noisy_latents, positions, timesteps, loss_mask, cond_channels = self._apply_scail_conditions( + modality_config=modality_config, + modality_key=modality_key, + noisy_latents=noisy_latents, + positions=positions, + timesteps=timesteps, + loss_mask=loss_mask, + data=data, + batch=batch, + device=device, + ) # Step 6: Build Modality modality = Modality( @@ -743,6 +725,56 @@ class FlexibleStrategy(TrainingStrategy): return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets + def _apply_scail_conditions( + self, + modality_config: ModalityConfig, + modality_key: str, + noisy_latents: Tensor, + positions: Tensor, + timesteps: Tensor, + loss_mask: Tensor | None, + data: LatentData, + batch: dict[str, Any], + device: torch.device, + ) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]: + """Apply SCAIL-2 driving concatenation then in-context mask channels (video only). + Driving tokens are prepended (cond-first, like reference) so the target stays at the tail for + loss slicing; the mask channels are then written onto those leading driving tokens (the noisy + target keeps a zero mask). Returns the possibly-extended sequence tensors plus ``cond_channels`` + (``None`` when no mask condition is present). + """ + if modality_key != "video": + return noisy_latents, positions, timesteps, loss_mask, None + + driving_token_count = 0 + for cond in modality_config.conditions: + if isinstance(cond, DrivingConditionConfig): + noisy_latents, positions, timesteps, loss_mask, driving_token_count = self._apply_driving_condition( + noisy_latents=noisy_latents, + positions=positions, + timesteps=timesteps, + loss_mask=loss_mask, + target_width=data.width, + batch=batch, + config=cond, + ) + + cond_channels = None + for cond in modality_config.conditions: + if isinstance(cond, MaskChannelsConditionConfig): + cond_channels = self._build_mask_channels( + config=cond, + batch=batch, + total_tokens=noisy_latents.shape[1], + driving_token_count=driving_token_count, + target_frames=data.num_frames, + target_height=data.height, + target_width=data.width, + device=device, + ) + + return noisy_latents, positions, timesteps, loss_mask, cond_channels + def _apply_driving_condition( self, noisy_latents: Tensor, -- 2.34.1 From b69eedbd54a39d62104234f3050d50b595b36395 Mon Sep 17 00:00:00 2001 From: indigo Date: Sun, 12 Jul 2026 09:37:27 +0800 Subject: [PATCH 6/7] Add Modal test environment for SCAIL-2 training/inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A free-tier-friendly Modal app to validate the SCAIL-2 code path on real Linux/GPU without a local GPU or the 19B checkpoint. - modal/checks.py: device-aware (CPU/GPU auto) consolidation of the Phase 1-4 plumbing checks using tiny random-init models + synthetic data (no checkpoint, no dataset): driving concat, zero-init patchify_proj widening (output-preserving), a FlexibleStrategy driving+mask training step + loss, and inference conditioning assembly. Passes on CPU locally. - modal/app.py: builds the workspace via `uv sync` (skips CUDA-only ltx-kernels; attention falls back to SDPA). Functions: verify (CPU, ~free), smoke (T4, cents), train (A10G/A100, paid — runs the real trainer against a checkpoint + data on the scail-data Volume). - modal/README.md: free-tier setup (modal setup), run commands, cost table, the scale-up path, and the expected preprocessed dataset layout. Co-Authored-By: Claude Opus 4.8 --- modal/README.md | 101 ++++++++++++++++++++++++++++++ modal/app.py | 84 +++++++++++++++++++++++++ modal/checks.py | 159 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 344 insertions(+) create mode 100644 modal/README.md create mode 100644 modal/app.py create mode 100644 modal/checks.py diff --git a/modal/README.md b/modal/README.md new file mode 100644 index 0000000..e39d635 --- /dev/null +++ b/modal/README.md @@ -0,0 +1,101 @@ +# SCAIL-2 on Modal — free testing environment + +Test the SCAIL-2 LTX-2 training/inference **code path** on Modal without a local +GPU. The free path uses tiny random-init models + synthetic data (no checkpoint, +no dataset), so it validates that the SCAIL wiring runs on real Linux/GPU — not +model quality. + +> ⚠️ The real 19B LTX-2 model is **not** free to train/run. The `verify`/`smoke` +> targets are near-free; `train` on a real checkpoint uses a paid GPU. + +## 1. One-time setup + +1. Sign up at [modal.com](https://modal.com) (the free Starter plan includes a + monthly credit allowance — enough for many `verify`/`smoke` runs). +2. Install and authenticate: + ```bash + pip install modal + modal setup # opens a browser to link your account / token + ``` + +## 2. Free / near-free checks + +From the repo root: + +```bash +# CPU: SCAIL-2 Phase 1-4 plumbing (driving concat, mask channels, zero-init widen, +# a FlexibleStrategy training step + loss, inference conditioning assembly). +modal run modal/app.py::verify + +# Same checks on a T4 GPU (validates the CUDA path). ~cents. +modal run modal/app.py::smoke +``` + +The first run builds the image (installs the `ltx-core`/`ltx-pipelines`/`ltx-trainer` +workspace via `uv sync`). `ltx-kernels` (CUDA-compiled) is intentionally skipped — +attention falls back to PyTorch SDPA, so no CUDA toolchain is required. + +Expected tail: +``` +[OK] Phase 1 driving concat: seq 48 -> 96 +[OK] Phase 2 zero-init widen: in_features=72, output preserved +[OK] Phase 3 training step: cond_channels (1, 96, 56), loss ... +[OK] Phase 4 build_scail_conditionings: [driving, mask_channels] +All SCAIL-2 checks passed on cuda # (or cpu) +``` + +## 3. Scaling up to real training (paid) + +`train` runs `packages/ltx-trainer/scripts/train.py` against a real checkpoint. +You must supply the weights + data via the persistent `scail-data` Volume. + +1. Create/populate the Volume (checkpoint, Gemma encoder, preprocessed latents): + ```bash + modal volume create scail-data # if not auto-created + modal volume put scail-data /local/ltx-2-model.safetensors /model/ltx-2.safetensors + modal volume put scail-data /local/gemma /model/gemma + modal volume put scail-data /local/preprocessed /data/preprocessed + ``` +2. Copy `configs/scail_animation_lora.yaml`, and point its paths at the mounted + Volume (everything lands under `/data` in the container): + ```yaml + model: + model_path: "/data/model/ltx-2.safetensors" + text_encoder_path: "/data/model/gemma" + mask_conditioning_channels: 56 # widens patchify_proj at load + data: + preprocessed_data_root: "/data/preprocessed" + ``` + (Put your edited config on the Volume too, or bake it into the repo.) +3. Launch (pick a GPU big enough for the model — the 19B needs an A100): + ```bash + # edit gpu="A10G" -> "A100" in modal/app.py::train for the full model + modal run modal/app.py::train --config-rel /data/scail_animation_lora.yaml + ``` + +### Dataset preprocessing (not yet automated for SCAIL) + +The SCAIL training config expects, under `preprocessed_data_root/`: + +``` +latents/ # target video latents +conditions/ # text embeddings +driving_latents/ # driving video latents (same F/H/W as target) +char_masks/ # per-sample "mask" = [K+1, F_pix, H, W] (ch0 env switch, 1..K binding slots) +``` + +`latents/`, `conditions/`, and `driving_latents/` come from the existing +`packages/ltx-trainer/scripts/process_dataset.py` (run it once per video set). +`char_masks/` still needs a segmentation step (e.g. SAM) to produce the semantic +masks — that preprocessing is not implemented yet (see `docs/tasks.md` 3.7). + +## Cost notes + +| Target | GPU | Rough cost | Use | +|---------|-------|-----------|-----| +| `verify`| none | ~free | validate code path on CPU | +| `smoke` | T4 | cents | validate CUDA path | +| `train` | A10G/A100 | paid | real fine-tuning (needs checkpoint + data) | + +Free credits are best spent on `verify`/`smoke` to catch integration issues before +committing a paid GPU to a real run. Watch usage in the Modal dashboard. diff --git a/modal/app.py b/modal/app.py new file mode 100644 index 0000000..0ab20f6 --- /dev/null +++ b/modal/app.py @@ -0,0 +1,84 @@ +"""Modal app for testing the SCAIL-2 LTX-2 training/inference path. + +Free-tier friendly: `verify` runs on CPU and `smoke` on a cheap T4, both using +tiny random-init models + synthetic data (no checkpoint, no dataset) so a full run +costs pennies. `train` is the scale-up entrypoint that runs the real trainer once +you upload a checkpoint + preprocessed data to the `scail-data` Volume -- that one +uses a real GPU and is NOT free. + +Setup (once): + pip install modal && modal setup + +Run: + modal run modal/app.py::verify # CPU plumbing check (~free) + modal run modal/app.py::smoke # same checks on a T4 GPU (cents) + modal run modal/app.py::train --config-rel configs/scail_animation_lora.yaml +""" + +import subprocess +from pathlib import Path + +import modal + +REPO = Path(__file__).parent.parent + +# Build the workspace once into the image. ltx-kernels (CUDA-compiled) is excluded +# from the workspace and not needed -- attention falls back to SDPA. +image = ( + modal.Image.debian_slim(python_version="3.11") + .apt_install("git") + .pip_install("uv") + .env({"UV_LINK_MODE": "copy", "UV_PROJECT_ENVIRONMENT": "/root/LTX-2/.venv"}) + .add_local_dir( + str(REPO), + "/root/LTX-2", + copy=True, + ignore=[".git", ".venv", "**/__pycache__", "**/*.pyc", "outputs", "wandb", "**/.pytest_cache"], + ) + .run_commands("cd /root/LTX-2 && uv sync --package ltx-trainer") +) + +app = modal.App("scail-ltx2", image=image) + +# Persistent volume for the (large) checkpoint + preprocessed dataset used by `train`. +data_volume = modal.Volume.from_name("scail-data", create_if_missing=True) + +_UV_PY = ["uv", "run", "--package", "ltx-trainer", "python"] + + +def _run(args: list[str]) -> None: + subprocess.run([*_UV_PY, *args], cwd="/root/LTX-2", check=True) + + +@app.function(timeout=1800) +def verify() -> None: + """Run the SCAIL-2 plumbing checks on CPU (near-free).""" + _run(["modal/checks.py"]) + + +@app.function(gpu="T4", timeout=1800) +def smoke() -> None: + """Run the same checks on a T4 GPU to validate the CUDA path (cents).""" + _run(["modal/checks.py"]) + + +@app.function(gpu="A10G", timeout=60 * 60 * 6, volumes={"/data": data_volume}) +def train(config_rel: str = "configs/scail_animation_lora.yaml") -> None: + """Run the real SCAIL trainer. NOT free -- needs a real checkpoint + data. + + Upload your base checkpoint, Gemma text encoder, and preprocessed data to the + ``scail-data`` Volume (mounted at /data), and point the config's paths at /data. + For the full 19B model use a bigger GPU (e.g. gpu="A100") and expect real cost. + """ + _run(["packages/ltx-trainer/scripts/train.py", config_rel]) + data_volume.commit() + + +@app.local_entrypoint() +def main(target: str = "verify", config_rel: str = "configs/scail_animation_lora.yaml") -> None: + if target == "smoke": + smoke.remote() + elif target == "train": + train.remote(config_rel) + else: + verify.remote() diff --git a/modal/checks.py b/modal/checks.py new file mode 100644 index 0000000..686e8c9 --- /dev/null +++ b/modal/checks.py @@ -0,0 +1,159 @@ +# ruff: noqa: T201 +"""Device-aware SCAIL-2 smoke checks (CPU or GPU), runnable anywhere. + +Consolidates the Phase 1-4 plumbing checks into one script that auto-selects CUDA +when available. It uses tiny, randomly-initialised models and synthetic data, so it +needs **no checkpoint and no dataset** -- ideal for a near-free Modal test that +validates the SCAIL training/inference code path end to end in a real Linux/GPU +environment before spending credits on the full 19B model. + +Run locally: uv run --package ltx-trainer python modal/checks.py +On Modal: modal run modal/app.py::verify (CPU) + modal run modal/app.py::smoke (T4 GPU) +""" + +from __future__ import annotations + +from dataclasses import replace + +import torch + +from ltx_core.components.noisers import GaussianNoiser +from ltx_core.components.patchifiers import VideoLatentPatchifier +from ltx_core.conditioning import DrivingMode, VideoConditionByDrivingLatent +from ltx_core.model.transformer.mask_channels_checkpoint import ( + widen_module_patchify_proj_for_mask_channels, +) +from ltx_core.model.transformer.model import LTXModel, LTXModelType +from ltx_core.tools import VideoLatentTools +from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape +from ltx_pipelines.scail_animation import build_scail_conditionings +from ltx_pipelines.utils.helpers import create_noised_state, modality_from_latent_state +from ltx_trainer.timestep_samplers import UniformTimestepSampler +from ltx_trainer.training_strategies.flexible import ( + DrivingConditionConfig, + FlexibleStrategy, + FlexibleStrategyConfig, + MaskChannelsConditionConfig, + ModalityConfig, +) + +# --- config --------------------------------------------------------------------- +B, C, F, H, W = 1, 16, 3, 4, 4 +HEADS, HEAD_DIM, LAYERS, INNER = 4, 16, 2, 64 +CTX_LEN, K = 8, 6 +SCALE = SpatioTemporalScaleFactors.default() +MASK_CH = SCALE.time * (K + 1) # 56 +F_PIX = (F - 1) * SCALE.time + 1 +N = F * H * W + + +def _tiny_model(mask_channels: int, device: torch.device, dtype: torch.dtype) -> LTXModel: + model = LTXModel( + model_type=LTXModelType.VideoOnly, + num_attention_heads=HEADS, + attention_head_dim=HEAD_DIM, + in_channels=C, + out_channels=C, + num_layers=LAYERS, + cross_attention_dim=INNER, + mask_conditioning_channels=mask_channels, + ) + for p in model.parameters(): + torch.nn.init.normal_(p, std=0.02) + return model.to(device=device, dtype=dtype).eval() + + +def main() -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + print(f"== SCAIL-2 checks on device={device} " + f"({torch.cuda.get_device_name(0) if device.type == 'cuda' else 'cpu'}) ==") + torch.manual_seed(0) + + tools = VideoLatentTools( + patchifier=VideoLatentPatchifier(patch_size=1), + target_shape=VideoLatentShape(batch=B, channels=C, frames=F, height=H, width=W), + fps=24.0, + scale_factors=SCALE, + ) + + # Phase 1: driving-latent concat + ΔW RoPE (inference conditioning path). + driving = torch.randn(B, C, F, H, W, device=device, dtype=dtype) + noiser = GaussianNoiser(generator=torch.Generator(device=device).manual_seed(1)) + state = create_noised_state( + tools, [VideoConditionByDrivingLatent(driving, mode=DrivingMode.ANIMATION)], noiser, dtype, device, 1.0 + ) + assert state.latent.shape[1] == 2 * N + print(f"[OK] Phase 1 driving concat: seq {N} -> {state.latent.shape[1]}") + + # Phase 2: widen patchify_proj (zero-init) is output-preserving. + base = _tiny_model(0, device, dtype) + ctx = torch.randn(B, CTX_LEN, INNER, device=device, dtype=dtype) + sigma = torch.ones(B, device=device, dtype=dtype) + mod = modality_from_latent_state(state, context=ctx, sigma=sigma) + with torch.inference_mode(): + out0, _ = base(video=mod, audio=None, perturbations=None) + widen_module_patchify_proj_for_mask_channels(base, MASK_CH) + mod_cc = replace(mod, cond_channels=torch.randn(B, 2 * N, MASK_CH, device=device, dtype=dtype)) + with torch.inference_mode(): + out1, _ = base(video=mod_cc, audio=None, perturbations=None) + assert torch.allclose(out0, out1, atol=1e-4), "zero-init widening changed output" + print(f"[OK] Phase 2 zero-init widen: in_features={base.patchify_proj.in_features}, output preserved") + + # Phase 3: FlexibleStrategy driving + mask-channels training step. + cfg = FlexibleStrategyConfig( + name="flexible", + video=ModalityConfig( + is_generated=True, + latents_dir="video_latents", + conditions=[ + DrivingConditionConfig(type="driving", latents_dir="driving_latents"), + MaskChannelsConditionConfig(type="mask_channels", mask_dir="char_masks", num_slots=K), + ], + ), + ) + strategy = FlexibleStrategy(cfg) + + def latents() -> dict: + return { + "latents": torch.randn(B, C, F, H, W, device=device), + "num_frames": torch.tensor([F]), + "height": torch.tensor([H]), + "width": torch.tensor([W]), + "fps": torch.tensor([24.0]), + } + + mask_pix = torch.rand(B, K + 1, F_PIX, H * SCALE.height, W * SCALE.width, device=device) + batch = { + "video_latents": latents(), + "driving_latents": latents(), + "char_masks": {"mask": (mask_pix > 0.5).float()}, + "conditions": { + "video_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device), + "audio_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device), + "prompt_attention_mask": torch.ones(B, CTX_LEN, device=device), + }, + } + inputs = strategy.prepare_training_inputs(batch, UniformTimestepSampler(0.0, 1.0)) + assert inputs.video.cond_channels.shape == (B, 2 * N, MASK_CH) + model = _tiny_model(MASK_CH, device, dtype) + with torch.inference_mode(): + vpred, _ = model(video=inputs.video, audio=None, perturbations=None) + loss = strategy.compute_loss(vpred, None, inputs) + assert loss.shape == (B,) + assert torch.isfinite(loss).all() + cc_shape = tuple(inputs.video.cond_channels.shape) + print(f"[OK] Phase 3 training step: cond_channels {cc_shape}, loss {loss.item():.4f}") + + # Phase 4: inference conditioning assembly. + masks = batch["char_masks"]["mask"] + conds = build_scail_conditionings(driving, masks, mode=DrivingMode.ANIMATION) + assert len(conds) == 2 + print("[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]") + + print("\nAll SCAIL-2 checks passed on", device) + + +if __name__ == "__main__": + main() -- 2.34.1 From 5594d49c7644e33b1a3824c4a0ded407a2c8e166 Mon Sep 17 00:00:00 2001 From: indigo Date: Mon, 13 Jul 2026 09:49:25 +0800 Subject: [PATCH 7/7] Add char_masks preprocessing script for SCAIL-2 training (task 3.7) scripts/process_char_masks.py turns per-sample character label-map videos/images (integer pixel labels: 0 = environment, 1..K = characters -> binding slots) into the pixel-space semantic-mask tensors the SCAIL training/inference path consumes: {"mask": [K+1, F_pix, H_pix, W_pix]} (ch0 = environment switch, ch1..K = slots). - Aligns to the target video's latent grid read from the saved latent metadata (F_pix=(F-1)*8+1, H*32, W*32), so char_masks/ lines up file-for-file with latents/ / driving_latents/ for PrecomputedDataset. - Nearest-neighbour resize so integer labels are never blended; labels > K are dropped with a warning; ch0 filled uniformly with --environment-switch. - Reuses process_videos.py helpers (naming, atomic save, VAE factors) and matches its typer CLI conventions. Verified on CPU: a synthetic 2-character label map (plus an out-of-range id) produces mask (7,17,128,128) with ch0 uniform, slots placed correctly, id>K dropped, and feeds encode_mask_channels to the 8*(K+1)=56 channels. README + docs/tasks.md 3.7 updated (upstream label-map generation via SAM/tracking is dataset-specific and still out of scope). Co-Authored-By: Claude Opus 4.8 --- docs/tasks.md | 2 +- modal/README.md | 14 +- .../ltx-trainer/scripts/process_char_masks.py | 217 ++++++++++++++++++ 3 files changed, 230 insertions(+), 3 deletions(-) create mode 100644 packages/ltx-trainer/scripts/process_char_masks.py diff --git a/docs/tasks.md b/docs/tasks.md index bb21bfb..3ed1a12 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -40,7 +40,7 @@ | 3.4 | LoRA 模式解凍 patchify_proj | ✅ | `trainer._unfreeze_patchify_proj`：mask_channels>0 時把 video patchify_proj 設 trainable，讓新欄位隨 LoRA 一起訓 | | 3.5 | 範例 config + docs | ✅ | `configs/scail_animation_lora.yaml`；`configs/README.md` 與 `docs/training-modes.md` 表格 row | | 3.6 | CPU 單元驗證 | ✅ | `verify_phase3_trainer.py`：config round-trip、prepare_training_inputs 建 cond_channels[B,T,56]、driving 前置/mask placement、widened model forward + compute_loss finite、widen helper zero-init 等價 | -| 3.7 | dataset 前處理（產 driving latents + 語意 mask） | ⬜ | 需 process_dataset 產出 `driving_latents/`（同 target 形狀）與 `char_masks/`(mask=[K+1,F_pix,H,W])；語意 mask 需分割模型，屬資料工程，未做 | +| 3.7 | dataset 前處理（產 driving latents + 語意 mask） | 🟡 | `char_masks/` 前處理已做：`scripts/process_char_masks.py`（label-map 影片/圖 → `[K+1,F_pix,H,W]`，對齊 target latent，nearest 保留整數 label，ch0 環境開關）。`driving_latents/` 沿用既有 `process_videos.py`（driving 影片走 video latent 路徑，同 target 形狀）。**仍缺**：從原始影片產生 label-map 的分割/追蹤步驟（SAM 等，資料集特定，未含） | | 3.8 | validation runner 接 driving/mask | ⬜ | 驗證期取樣尚未接 SCAIL 條件（config 內 validation 先停用），與 Phase 4 一起 | | 3.9 | 實機訓練跑通 | ⬜ | 需 Linux + GPU + checkpoint，本機無法 | diff --git a/modal/README.md b/modal/README.md index e39d635..c2dd0fb 100644 --- a/modal/README.md +++ b/modal/README.md @@ -86,8 +86,18 @@ char_masks/ # per-sample "mask" = [K+1, F_pix, H, W] (ch0 env switch, 1. `latents/`, `conditions/`, and `driving_latents/` come from the existing `packages/ltx-trainer/scripts/process_dataset.py` (run it once per video set). -`char_masks/` still needs a segmentation step (e.g. SAM) to produce the semantic -masks — that preprocessing is not implemented yet (see `docs/tasks.md` 3.7). +`char_masks/` is produced by `packages/ltx-trainer/scripts/process_char_masks.py` +from per-sample **label-map** videos/images (integer pixel labels: `0` = +environment, `1..K` = characters → binding slots): + +```bash +python packages/ltx-trainer/scripts/process_char_masks.py dataset.csv \ + --mask-column char_labels --latents-dir ./latents \ + --output-dir ./char_masks --num-slots 6 --main-media-column media_path +``` + +You still need a segmentation/tracking model (e.g. SAM) to *produce* those label +maps from raw video — that upstream step is dataset-specific and not included. ## Cost notes diff --git a/packages/ltx-trainer/scripts/process_char_masks.py b/packages/ltx-trainer/scripts/process_char_masks.py new file mode 100644 index 0000000..26002d3 --- /dev/null +++ b/packages/ltx-trainer/scripts/process_char_masks.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 + +"""Preprocess SCAIL-2 character binding-slot masks into per-sample ``.pt`` tensors. + +Each sample's input is a **label-map** video or image whose integer pixel values index +characters: ``0`` = environment/background, ``k`` in ``1..K`` = character *k* (assigned to +binding slot *k*). This script aligns the label map to the target video's latent grid (read +from the saved video-latent metadata), and emits the pixel-space semantic-mask tensor that +``ltx_core.conditioning.encode_mask_channels`` (and the trainer's ``mask_channels`` condition) +consumes: + + {"mask": tensor[K+1, F_pix, H_pix, W_pix]} # ch0 = environment switch, ch1..K = binding slots + +where ``F_pix = (latent_frames - 1) * 8 + 1``, ``H_pix = latent_h * 32``, ``W_pix = latent_w * 32`` +(the SCAIL encoder then spatially downsamples + temporally stacks these to the 8*(K+1)=56 channels). + +Label maps are resized with **nearest-neighbour** interpolation so integer labels are never +blended. The environment-switch channel (ch0) is filled uniformly with ``--environment-switch`` +(``0.0`` = derive the environment from the reference image, ``1.0`` = from the driving video), +matching the paper's single-bit environment signal. + +Output naming mirrors the target latents (same relative path), so ``char_masks/`` lines up +file-for-file with ``latents/`` / ``driving_latents/`` for the trainer's ``PrecomputedDataset``. + +Standalone usage:: + + python scripts/process_char_masks.py dataset.csv \\ + --mask-column char_labels --latents-dir ./latents --output-dir ./char_masks --num-slots 6 +""" + +from pathlib import Path + +import numpy as np +import torch +import typer +from PIL import Image + +# Sibling scripts (resolved via scripts/ on sys.path), reused so naming + alignment match the other stages. +from process_videos import ( + IMAGE_FILE_EXTENSIONS, + VAE_SPATIAL_FACTOR, + VAE_TEMPORAL_FACTOR, + _atomic_save, + _load_paths_from_dataset, + _output_relative, +) +from rich.console import Console +from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn + +from ltx_trainer import logger +from ltx_trainer.video_utils import read_video + +app = typer.Typer( + pretty_exceptions_enable=False, + no_args_is_help=True, + help="Preprocess SCAIL-2 character label maps into [K+1, F_pix, H_pix, W_pix] mask tensors.", +) + + +def _load_label_frames(mask_file: Path, pixel_f: int) -> torch.Tensor: + """Load a label-map video/image as integer labels ``[F_pix, H, W]`` (float-typed integers). + + Images are tiled across ``pixel_f`` frames. Videos are read via the shared ``read_video`` helper + (which returns ``[F, C, H, W]`` in ``[0, 1]``); labels are recovered as ``round(value * 255)``, + so label maps must be stored as small-integer grayscale (label ``k`` -> pixel value ``k``). + """ + if mask_file.suffix.lower() in IMAGE_FILE_EXTENSIONS: + arr = np.array(Image.open(mask_file).convert("L")) # exact integer labels [H, W] (writable copy) + labels = torch.from_numpy(arr).float() + return labels.unsqueeze(0).expand(pixel_f, -1, -1).contiguous() + + frames, _ = read_video(str(mask_file), max_frames=pixel_f) # [F, C, H, W] in [0, 1] + labels = frames[:, 0].mul(255.0).round() # channel 0 -> integer labels [F, H, W] + if labels.shape[0] < pixel_f: + # Pad by repeating the last frame so every latent frame has a label map. + pad = labels[-1:].expand(pixel_f - labels.shape[0], -1, -1) + labels = torch.cat([labels, pad], dim=0) + return labels[:pixel_f] + + +def _labels_to_slot_masks(labels: torch.Tensor, num_slots: int, environment_switch: float) -> torch.Tensor: + """Convert integer label frames ``[F, H, W]`` to ``[K+1, F, H, W]`` (ch0 env switch, ch1..K slots).""" + f_pix, h, w = labels.shape + out = torch.zeros(num_slots + 1, f_pix, h, w, dtype=torch.float32) + out[0] = environment_switch # uniform environment-switch channel + for k in range(1, num_slots + 1): + out[k] = (labels == k).float() + if bool(((labels > num_slots) & (labels > 0)).any().item()): + logger.warning( + f"Label map contains ids > num_slots ({num_slots}); those pixels are dropped (treated as environment)." + ) + return out + + +def _resize_labels(labels: torch.Tensor, pixel_h: int, pixel_w: int) -> torch.Tensor: + """Nearest-neighbour resize of integer label frames ``[F, H, W]`` to ``[F, pixel_h, pixel_w]``.""" + if labels.shape[1:] == (pixel_h, pixel_w): + return labels + return torch.nn.functional.interpolate( + labels.unsqueeze(1), size=(pixel_h, pixel_w), mode="nearest" + ).squeeze(1) + + +def compute_char_masks( + dataset_file: str | Path, + mask_column: str, + latents_dir: str, + output_dir: str, + num_slots: int = 6, + environment_switch: float = 0.0, + main_media_column: str | None = None, + overwrite: bool = False, +) -> None: + """Preprocess character label maps into ``[K+1, F_pix, H_pix, W_pix]`` mask tensors. + + Args: + dataset_file: Metadata file (CSV/JSON/JSONL) with a column of label-map paths. + mask_column: Column containing the per-sample label-map video/image paths. + latents_dir: Directory of target video latents (read for spatial/temporal alignment). + output_dir: Directory to write ``char_masks`` ``.pt`` files. + num_slots: Number of binding slots K (output has ``K+1`` channels). + environment_switch: Uniform value for ch0 (0.0 = env from reference, 1.0 = from driving video). + main_media_column: Column used for output naming (defaults to ``mask_column``); set it to the + target-video column so masks align with ``latents/`` when label maps live elsewhere. + overwrite: Recompute even if the output already exists. + """ + dataset_path = Path(dataset_file) + data_root = dataset_path.parent + latents_path = Path(latents_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + naming_column = main_media_column or mask_column + mask_paths = _load_paths_from_dataset(dataset_path, mask_column) + naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths + + console = Console() + success = 0 + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + console=console, + ) as progress: + task = progress.add_task("Processing char masks", total=len(mask_paths)) + for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True): + progress.advance(task) + rel_path = _output_relative(naming_file, data_root) + latent_file = latents_path / rel_path.with_suffix(".pt") + out_file = output_path / rel_path.with_suffix(".pt") + + if not latent_file.exists(): + logger.warning(f"No target latent at {latent_file}, skipping mask {mask_file}") + continue + if not overwrite and out_file.is_file(): + continue + + meta = torch.load(latent_file, map_location="cpu", weights_only=True) + pixel_h = meta["height"] * VAE_SPATIAL_FACTOR + pixel_w = meta["width"] * VAE_SPATIAL_FACTOR + pixel_f = (meta["num_frames"] - 1) * VAE_TEMPORAL_FACTOR + 1 + + labels = _load_label_frames(mask_file, pixel_f) + labels = _resize_labels(labels, pixel_h, pixel_w) + mask = _labels_to_slot_masks(labels, num_slots, environment_switch) + + out_file.parent.mkdir(parents=True, exist_ok=True) + _atomic_save({"mask": mask.contiguous()}, out_file) + success += 1 + + logger.info(f"Char-mask preprocessing complete: {success} masks saved to {output_path}") + + +@app.command() +def main( + dataset_file: str = typer.Argument(..., help="Metadata file (CSV/JSON/JSONL) with a label-map column"), + mask_column: str = typer.Option(..., help="Column of per-sample label-map video/image paths"), + latents_dir: str = typer.Option(..., help="Directory of target video latents (for alignment)"), + output_dir: str = typer.Option(..., help="Output directory for char_masks .pt files"), + num_slots: int = typer.Option(6, help="Number of character binding slots K (channels = K+1)"), + environment_switch: float = typer.Option( + 0.0, help="Uniform ch0 value: 0.0 = environment from reference, 1.0 = from driving video" + ), + main_media_column: str | None = typer.Option( + None, help="Column for output naming (defaults to --mask-column; set to the target-video column to align)" + ), + overwrite: bool = typer.Option(False, help="Recompute even if the output already exists"), +) -> None: + """Preprocess SCAIL-2 character label maps into ``[K+1, F_pix, H_pix, W_pix]`` mask tensors. + + Example:: + + python scripts/process_char_masks.py dataset.csv \\ + --mask-column char_labels --latents-dir ./latents \\ + --output-dir ./char_masks --num-slots 6 --main-media-column media_path + """ + if not Path(dataset_file).is_file(): + raise typer.BadParameter(f"Dataset file not found: {dataset_file}") + if num_slots < 1: + raise typer.BadParameter("--num-slots must be >= 1") + + compute_char_masks( + dataset_file=dataset_file, + mask_column=mask_column, + latents_dir=latents_dir, + output_dir=output_dir, + num_slots=num_slots, + environment_switch=environment_switch, + main_media_column=main_media_column, + overwrite=overwrite, + ) + + +if __name__ == "__main__": + app() -- 2.34.1