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:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user