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,