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,]."""