Add Modal test environment for SCAIL-2 training/inference

A free-tier-friendly Modal app to validate the SCAIL-2 code path on real Linux/GPU without a local GPU or the 19B checkpoint. - modal/checks.py: device-aware (CPU/GPU auto) consolidation of the Phase 1-4 plumbing checks using tiny random-init models + synthetic data (no checkpoint, no dataset): driving concat, zero-init patchify_proj widening (output-preserving), a FlexibleStrategy driving+mask training step + loss, and inference conditioning assembly. Passes on CPU locally. - modal/app.py: builds the workspace via `uv sync` (skips CUDA-only ltx-kernels; attention falls back to SDPA). Functions: verify (CPU, ~free), smoke (T4, cents), train (A10G/A100, paid — runs the real trainer against a checkpoint + data on the scail-data Volume). - modal/README.md: free-tier setup (modal setup), run commands, cost table, the scale-up path, and the expected preprocessed dataset layout. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-12 09:37:27 +08:00
parent a598f89d99
commit b69eedbd54
3 changed files with 344 additions and 0 deletions
@@ -0,0 +1,101 @@
+# SCAIL-2 on Modal — free testing environment
+
+Test the SCAIL-2 LTX-2 training/inference **code path** on Modal without a local
+GPU. The free path uses tiny random-init models + synthetic data (no checkpoint,
+no dataset), so it validates that the SCAIL wiring runs on real Linux/GPU — not
+model quality.
+
+> ⚠️ The real 19B LTX-2 model is **not** free to train/run. The `verify`/`smoke`
+> targets are near-free; `train` on a real checkpoint uses a paid GPU.
+
+## 1. One-time setup
+
+1. Sign up at [modal.com](https://modal.com) (the free Starter plan includes a
+   monthly credit allowance — enough for many `verify`/`smoke` runs).
+2. Install and authenticate:
+   ```bash
+   pip install modal
+   modal setup          # opens a browser to link your account / token
+   ```
+
+## 2. Free / near-free checks
+
+From the repo root:
+
+```bash
+# CPU: SCAIL-2 Phase 1-4 plumbing (driving concat, mask channels, zero-init widen,
+# a FlexibleStrategy training step + loss, inference conditioning assembly).
+modal run modal/app.py::verify
+
+# Same checks on a T4 GPU (validates the CUDA path). ~cents.
+modal run modal/app.py::smoke
+```
+
+The first run builds the image (installs the `ltx-core`/`ltx-pipelines`/`ltx-trainer`
+workspace via `uv sync`). `ltx-kernels` (CUDA-compiled) is intentionally skipped —
+attention falls back to PyTorch SDPA, so no CUDA toolchain is required.
+
+Expected tail:
+```
+[OK] Phase 1 driving concat: seq 48 -> 96
+[OK] Phase 2 zero-init widen: in_features=72, output preserved
+[OK] Phase 3 training step: cond_channels (1, 96, 56), loss ...
+[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]
+All SCAIL-2 checks passed on cuda   # (or cpu)
+```
+
+## 3. Scaling up to real training (paid)
+
+`train` runs `packages/ltx-trainer/scripts/train.py` against a real checkpoint.
+You must supply the weights + data via the persistent `scail-data` Volume.
+
+1. Create/populate the Volume (checkpoint, Gemma encoder, preprocessed latents):
+   ```bash
+   modal volume create scail-data          # if not auto-created
+   modal volume put scail-data /local/ltx-2-model.safetensors /model/ltx-2.safetensors
+   modal volume put scail-data /local/gemma /model/gemma
+   modal volume put scail-data /local/preprocessed /data/preprocessed
+   ```
+2. Copy `configs/scail_animation_lora.yaml`, and point its paths at the mounted
+   Volume (everything lands under `/data` in the container):
+   ```yaml
+   model:
+     model_path: "/data/model/ltx-2.safetensors"
+     text_encoder_path: "/data/model/gemma"
+     mask_conditioning_channels: 56       # widens patchify_proj at load
+   data:
+     preprocessed_data_root: "/data/preprocessed"
+   ```
+   (Put your edited config on the Volume too, or bake it into the repo.)
+3. Launch (pick a GPU big enough for the model — the 19B needs an A100):
+   ```bash
+   # edit gpu="A10G" -> "A100" in modal/app.py::train for the full model
+   modal run modal/app.py::train --config-rel /data/scail_animation_lora.yaml
+   ```
+
+### Dataset preprocessing (not yet automated for SCAIL)
+
+The SCAIL training config expects, under `preprocessed_data_root/`:
+
+```
+latents/           # target video latents
+conditions/        # text embeddings
+driving_latents/   # driving video latents (same F/H/W as target)
+char_masks/        # per-sample "mask" = [K+1, F_pix, H, W]  (ch0 env switch, 1..K binding slots)
+```
+
+`latents/`, `conditions/`, and `driving_latents/` come from the existing
+`packages/ltx-trainer/scripts/process_dataset.py` (run it once per video set).
+`char_masks/` still needs a segmentation step (e.g. SAM) to produce the semantic
+masks — that preprocessing is not implemented yet (see `docs/tasks.md` 3.7).
+
+## Cost notes
+
+| Target  | GPU   | Rough cost | Use |
+|---------|-------|-----------|-----|
+| `verify`| none  | ~free      | validate code path on CPU |
+| `smoke` | T4    | cents      | validate CUDA path |
+| `train` | A10G/A100 | paid   | real fine-tuning (needs checkpoint + data) |
+
+Free credits are best spent on `verify`/`smoke` to catch integration issues before
+committing a paid GPU to a real run. Watch usage in the Modal dashboard.
@@ -0,0 +1,84 @@
+"""Modal app for testing the SCAIL-2 LTX-2 training/inference path.
+
+Free-tier friendly: `verify` runs on CPU and `smoke` on a cheap T4, both using
+tiny random-init models + synthetic data (no checkpoint, no dataset) so a full run
+costs pennies. `train` is the scale-up entrypoint that runs the real trainer once
+you upload a checkpoint + preprocessed data to the `scail-data` Volume -- that one
+uses a real GPU and is NOT free.
+
+Setup (once):
+    pip install modal && modal setup
+
+Run:
+    modal run modal/app.py::verify     # CPU plumbing check (~free)
+    modal run modal/app.py::smoke      # same checks on a T4 GPU (cents)
+    modal run modal/app.py::train --config-rel configs/scail_animation_lora.yaml
+"""
+
+import subprocess
+from pathlib import Path
+
+import modal
+
+REPO = Path(__file__).parent.parent
+
+# Build the workspace once into the image. ltx-kernels (CUDA-compiled) is excluded
+# from the workspace and not needed -- attention falls back to SDPA.
+image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .apt_install("git")
+    .pip_install("uv")
+    .env({"UV_LINK_MODE": "copy", "UV_PROJECT_ENVIRONMENT": "/root/LTX-2/.venv"})
+    .add_local_dir(
+        str(REPO),
+        "/root/LTX-2",
+        copy=True,
+        ignore=[".git", ".venv", "**/__pycache__", "**/*.pyc", "outputs", "wandb", "**/.pytest_cache"],
+    )
+    .run_commands("cd /root/LTX-2 && uv sync --package ltx-trainer")
+)
+
+app = modal.App("scail-ltx2", image=image)
+
+# Persistent volume for the (large) checkpoint + preprocessed dataset used by `train`.
+data_volume = modal.Volume.from_name("scail-data", create_if_missing=True)
+
+_UV_PY = ["uv", "run", "--package", "ltx-trainer", "python"]
+
+
+def _run(args: list[str]) -> None:
+    subprocess.run([*_UV_PY, *args], cwd="/root/LTX-2", check=True)
+
+
+@app.function(timeout=1800)
+def verify() -> None:
+    """Run the SCAIL-2 plumbing checks on CPU (near-free)."""
+    _run(["modal/checks.py"])
+
+
+@app.function(gpu="T4", timeout=1800)
+def smoke() -> None:
+    """Run the same checks on a T4 GPU to validate the CUDA path (cents)."""
+    _run(["modal/checks.py"])
+
+
+@app.function(gpu="A10G", timeout=60 * 60 * 6, volumes={"/data": data_volume})
+def train(config_rel: str = "configs/scail_animation_lora.yaml") -> None:
+    """Run the real SCAIL trainer. NOT free -- needs a real checkpoint + data.
+
+    Upload your base checkpoint, Gemma text encoder, and preprocessed data to the
+    ``scail-data`` Volume (mounted at /data), and point the config's paths at /data.
+    For the full 19B model use a bigger GPU (e.g. gpu="A100") and expect real cost.
+    """
+    _run(["packages/ltx-trainer/scripts/train.py", config_rel])
+    data_volume.commit()
+
+
+@app.local_entrypoint()
+def main(target: str = "verify", config_rel: str = "configs/scail_animation_lora.yaml") -> None:
+    if target == "smoke":
+        smoke.remote()
+    elif target == "train":
+        train.remote(config_rel)
+    else:
+        verify.remote()
@@ -0,0 +1,159 @@
+# 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()