LTX-2/modal/checks.py

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