# 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()