Add SCAIL-2 training integration (Phase 3, ltx-trainer FlexibleStrategy)
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 <noreply@anthropic.com>
This commit is contained in:
+6
-2
@@ -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。
|
||||
|
||||
+11
-2
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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"
|
||||
@@ -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) |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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. "
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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,]."""
|
||||
|
||||
Reference in New Issue
Block a user