From 06c0870bbb2543a6633c7c9110f92424bda8bce4 Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 10:22:35 +0800 Subject: [PATCH] Add SCAIL-2 animation inference pipeline (Phase 4, ltx-pipelines) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two-stage distilled inference pipeline that animates a character from a driving video, wiring the Phase 1-3 SCAIL-2 conditioning into a runnable CLI. - scail_animation.py: ScailAnimationPipeline (mirrors distilled.py) plus a pure, testable build_scail_conditionings that assembles VideoConditionByDrivingLatent + VideoConditionByMaskChannels (driving appended, mask on the trailing driving tokens), and a load_masks helper. main() + scail_animation_arg_parser add --driving-video / --mask-path / --mode / --driving-strength on top of the standard two-stage distilled parser. - LTXModelConfigurator now reads config `mask_conditioning_channels`, so a SCAIL-trained checkpoint whose config declares it builds the widened patchify_proj automatically — no runtime widening wrapper needed at inference. - CLAUDE.md pipeline table row. Verified on CPU (verify_phase4_pipeline.py): the module imports, the CLI parses the SCAIL flags, build_scail_conditionings grows the sequence and places the mask channels on the driving tokens (target stays zero), driving-only leaves cond_channels None, and the configurator honors mask_conditioning_channels (patchify_proj widened from config). End-to-end runs still need a GPU and a SCAIL-trained checkpoint. Co-Authored-By: Claude Opus 4.8 --- docs/plan.md | 8 +- docs/tasks.md | 6 +- .../model/transformer/model_configurator.py | 2 + packages/ltx-pipelines/CLAUDE.md | 1 + .../src/ltx_pipelines/scail_animation.py | 349 ++++++++++++++++++ .../src/ltx_pipelines/utils/args.py | 37 ++ 6 files changed, 400 insertions(+), 3 deletions(-) create mode 100644 packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py diff --git a/docs/plan.md b/docs/plan.md index 1161ce7..ad197c9 100644 --- a/docs/plan.md +++ b/docs/plan.md @@ -52,8 +52,12 @@ - `configs/scail_animation_lora.yaml` + docs。CPU 單元驗證通過（`verify_phase3_trainer.py`）。 - **本機無 GPU/Linux/checkpoint → 未跑實機訓練**；dataset 前處理（driving latents + 語意 mask）與 validation runner 接線未做。 -### Phase 4 — Pipeline + CLI 包裝 ⬜ 未開始 -- 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。 +### Phase 4 — Pipeline + CLI 包裝 ✅ 已完成（程式碼路徑；實跑需 GPU） +- `LTXModelConfigurator` 讀 `mask_conditioning_channels` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime wrapper）。 +- `scail_animation.py`：`ScailAnimationPipeline`（仿 `distilled.py` 兩階段）+ 可測純函式 `build_scail_conditionings`（driving + mask 條件組裝）+ `load_masks` + `main()`。 +- `utils/args.py` `scail_animation_arg_parser`（`--driving-video/--mask-path/--mode/--driving-strength`）。 +- CPU 驗證通過（`verify_phase4_pipeline.py`）；`ltx-pipelines/CLAUDE.md` 表格 row。 +- **實機端到端需 GPU + SCAIL-trained checkpoint**（config 含 `mask_conditioning_channels` + 訓練好的 patchify_proj + LoRA）。 ## Phase 1 簡化取捨（記錄，Phase 2 需回頭處理） - (a) driving 時間座標直接複製 target 的（token-wise），故 driving 需與 target 同 F/H/W。 diff --git a/docs/tasks.md b/docs/tasks.md index fac85e4..bb21bfb 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -48,7 +48,11 @@ | # | 任務 | 狀態 | 備註 | |---|---|---|---| -| 4.1 | `scail_animation.py` pipeline + arg parser | ⬜ | 仿 `lipdub.py` | +| 4.1 | configurator 讀 `mask_conditioning_channels` | ✅ | `LTXModelConfigurator` 兩分支 `config.get("mask_conditioning_channels",0)` → SCAIL checkpoint 自描述、載入自動加寬（不需 runtime widen wrapper） | +| 4.2 | `ScailAnimationPipeline` + `build_scail_conditionings` | ✅ | `scail_animation.py`：仿 `distilled.py` 兩階段，注入 SCAIL driving+mask conditioning；`build_scail_conditionings` 為可測純函式；`load_masks` helper | +| 4.3 | `scail_animation_arg_parser` + `main()` | ✅ | `utils/args.py`：在 `default_2_stage_distilled_arg_parser` 上加 `--driving-video/--mask-path/--mode/--driving-strength` | +| 4.4 | CPU 驗證 + docs | ✅ | `verify_phase4_pipeline.py`：import、CLI round-trip、conditioning assembly(driving append + mask 在尾端)、configurator 自描述加寬；`ltx-pipelines/CLAUDE.md` pipeline 表格 row | +| 4.5 | 實機端到端跑通 | ⬜ | 需 GPU + SCAIL-trained checkpoint（含 `mask_conditioning_channels` config + 訓練好的 patchify_proj + LoRA） | ## 決議紀錄 - **範圍**：先只做 Phase 1（推論期 PoC）。Phase 2+ 待 Phase 1 驗證後再討論。 diff --git a/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py index fb47cef..59159ad 100644 --- a/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py +++ b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py @@ -69,6 +69,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]): caption_projection=caption_projection, audio_caption_projection=audio_caption_projection, cross_attention_adaln=config.get("cross_attention_adaln", False), + mask_conditioning_channels=config.get("mask_conditioning_channels", 0), ) @@ -120,6 +121,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]): apply_gated_attention=config.get("apply_gated_attention", False), caption_projection=caption_projection, cross_attention_adaln=config.get("cross_attention_adaln", False), + mask_conditioning_channels=config.get("mask_conditioning_channels", 0), ) diff --git a/packages/ltx-pipelines/CLAUDE.md b/packages/ltx-pipelines/CLAUDE.md index 2aa573d..39a004b 100644 --- a/packages/ltx-pipelines/CLAUDE.md +++ b/packages/ltx-pipelines/CLAUDE.md @@ -25,6 +25,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for | `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference | | `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control | | `LipDubPipeline` | `lipdub.py` | 2 | Distilled only | Euler | Lip dubbing with IC-LoRA + audio ref conditioning | +| `ScailAnimationPipeline` | `scail_animation.py` | 2 | Distilled + SCAIL LoRA | Euler | SCAIL-2 character animation from a driving video (driving latent concat + ΔW RoPE + in-context mask channels) | | `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration | ## Guidance diff --git a/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py b/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py new file mode 100644 index 0000000..d167e7c --- /dev/null +++ b/packages/ltx-pipelines/src/ltx_pipelines/scail_animation.py @@ -0,0 +1,349 @@ +"""SCAIL-2 character-animation inference pipeline. + +Two-stage distilled video generation that animates a character from a *driving* +video: the driving latent is concatenated into the token sequence with a RoPE +width offset (ΔW) and in-context mask channels route motion per character. This +mirrors the distilled pipeline (``distilled.py``) and lip-dub pipeline +(``lipdub.py``) structure, adding the SCAIL conditioning assembly. + +Requirements: +- A SCAIL-trained checkpoint whose config declares ``mask_conditioning_channels`` + (so ``LTXModelConfigurator`` builds the widened ``patchify_proj`` automatically) + plus the SCAIL LoRA adapter. +- A driving video (same target resolution) and a semantic mask tensor + ``[K+1, F_pix, H_pix, W_pix]`` (channel 0 = environment switch, 1..K = binding + slots), saved as a ``.pt`` file. + +The SCAIL-specific conditioning assembly is factored into +``build_scail_conditionings`` so it can be unit-tested without the heavy models. +""" + +from __future__ import annotations + +import logging +from collections.abc import Iterator + +import torch + +from ltx_core.components.noisers import GaussianNoiser +from ltx_core.conditioning import ( + ConditioningItem, + DrivingMode, + VideoConditionByDrivingLatent, + VideoConditionByMaskChannels, +) +from ltx_core.loader import LoraPathStrengthAndSDOps +from ltx_core.loader.registry import Registry +from ltx_core.model.transformer.compiling import CompilationConfig +from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number +from ltx_core.quantization import QuantizationPolicy +from ltx_core.types import Audio, VideoPixelShape +from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy +from ltx_pipelines.utils.args import ImageConditioningInput, resolve_cli_params, scail_animation_arg_parser +from ltx_pipelines.utils.blocks import ( + DiffusionStage, + ImageConditioner, + PromptEncoder, + VideoDecoder, + VideoUpsampler, +) +from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS +from ltx_pipelines.utils.denoisers import SimpleDenoiser +from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device +from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess +from ltx_pipelines.utils.types import ModalitySpec, OffloadMode + + +def build_scail_conditionings( + driving_latent: torch.Tensor, + masks: torch.Tensor | None, + *, + mode: DrivingMode = DrivingMode.ANIMATION, + strength: float = 1.0, + width_offset: float | None = None, +) -> list[ConditioningItem]: + """Assemble the SCAIL driving + mask-channel conditioning items (order matters). + + The driving item is applied first (it appends the driving tokens at the tail); the mask item is + applied second, writing its channels onto those trailing driving tokens (the noisy target keeps a + zero mask). ``masks`` may be ``None`` to run driving-only (e.g. a model without mask channels). + + Args: + driving_latent: Driving video latent ``[B, C, F, H, W]`` (same F/H/W as the target). + masks: Semantic masks ``[B, K+1, F_pix, H_pix, W_pix]`` or ``None``. + mode: SCAIL mode (animation/replacement). + strength: Driving conditioning strength (1.0 keeps it clean/frozen). + width_offset: ΔW in RoPE pixel-space width units (None = target pixel width). + """ + conditionings: list[ConditioningItem] = [ + VideoConditionByDrivingLatent( + latent=driving_latent, + mode=mode, + width_offset=width_offset, + strength=strength, + ) + ] + if masks is not None: + conditionings.append(VideoConditionByMaskChannels(masks=masks)) + return conditionings + + +def load_masks(mask_path: str, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Load a semantic mask tensor from a ``.pt`` file and shape it to ``[B, K+1, F_pix, H, W]``.""" + masks = torch.load(mask_path, map_location=device, weights_only=True) + if isinstance(masks, dict): + masks = masks["mask"] + if masks.dim() == 4: # [K+1, F, H, W] -> add batch + masks = masks.unsqueeze(0) + return masks.to(device=device, dtype=dtype) + + +class ScailAnimationPipeline: + """Two-stage distilled SCAIL-2 character animation (driving video + in-context mask channels).""" + + def __init__( + self, + distilled_checkpoint_path: str, + gemma_root: str, + spatial_upsampler_path: str, + scail_lora: LoraPathStrengthAndSDOps, + device: torch.device | None = None, + quantization: QuantizationPolicy | None = None, + registry: Registry | None = None, + compilation_config: CompilationConfig | None = None, + offload_mode: OffloadMode = OffloadMode.NONE, + alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM, + ) -> None: + self.device = device or get_device() + self.dtype = torch.bfloat16 + + self.prompt_encoder = PromptEncoder( + distilled_checkpoint_path, + gemma_root, + self.dtype, + self.device, + registry=registry, + offload_mode=offload_mode, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.image_conditioner = ImageConditioner( + distilled_checkpoint_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.stage = DiffusionStage.from_checkpoint( + distilled_checkpoint_path, + self.dtype, + self.device, + loras=(scail_lora,), + quantization=quantization, + registry=registry, + compilation_config=compilation_config, + offload_mode=offload_mode, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.upsampler = VideoUpsampler( + distilled_checkpoint_path, + spatial_upsampler_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + self.video_decoder = VideoDecoder( + distilled_checkpoint_path, + self.dtype, + self.device, + registry=registry, + alloc_trim_strategy=alloc_trim_strategy, + ) + + def _encode_driving_latent( + self, + driving_video_path: str, + height: int, + width: int, + num_frames: int, + video_encoder: VideoEncoder, + tiling_config: TilingConfig | None, + ) -> torch.Tensor: + """Decode + encode the driving video into a VAE latent at the target resolution.""" + frame_gen = decode_video_by_frame(path=driving_video_path, frame_cap=num_frames, device=self.device) + video = video_preprocess(frame_gen, height, width, self.dtype, self.device) + if tiling_config is not None: + return video_encoder.tiled_encode(video, tiling_config) + return video_encoder(video) + + def _video_conditionings( + self, + images: list[ImageConditioningInput], + driving_video_path: str, + masks: torch.Tensor | None, + height: int, + width: int, + num_frames: int, + mode: DrivingMode, + driving_strength: float, + video_encoder: VideoEncoder, + tiling_config: TilingConfig | None, + ) -> list[ConditioningItem]: + conditionings = combined_image_conditionings( + images=images, + height=height, + width=width, + video_encoder=video_encoder, + dtype=self.dtype, + device=self.device, + ) + driving_latent = self._encode_driving_latent( + driving_video_path, height, width, num_frames, video_encoder, tiling_config + ) + conditionings.extend( + build_scail_conditionings(driving_latent, masks, mode=mode, strength=driving_strength) + ) + return conditionings + + @torch.inference_mode() + def __call__( # noqa: PLR0913 + self, + prompt: str, + seed: int, + height: int, + width: int, + num_frames: int, + frame_rate: float, + driving_video_path: str, + mask_path: str | None, + images: list[ImageConditioningInput] | None = None, + mode: DrivingMode = DrivingMode.ANIMATION, + driving_strength: float = 1.0, + enhance_prompt: bool = False, + tiling_config: TilingConfig | None = None, + stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS, + stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS, + ) -> tuple[Iterator[torch.Tensor], Audio | None]: + assert_resolution(height=height, width=width, is_two_stage=True) + images = images or [] + + generator = torch.Generator(device=self.device).manual_seed(seed) + noiser = GaussianNoiser(generator=generator) + encode_tiling = TilingConfig.default() + + (ctx_p,) = self.prompt_encoder( + [prompt], + enhance_first_prompt=enhance_prompt, + enhance_prompt_image=images[0][0] if len(images) > 0 else None, + enhance_prompt_seed=seed, + ) + video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding + + masks = load_masks(mask_path, self.device, self.dtype) if mask_path is not None else None + + def build_video_conditionings(out_height: int, out_width: int) -> list[ConditioningItem]: + return self.image_conditioner( + lambda enc: self._video_conditionings( + images=images, + driving_video_path=driving_video_path, + masks=masks, + height=out_height, + width=out_width, + num_frames=num_frames, + mode=mode, + driving_strength=driving_strength, + video_encoder=enc, + tiling_config=encode_tiling, + ) + ) + + # Stage 1: low resolution. + stage_1_sigmas_t = stage_1_sigmas.to(dtype=torch.float32, device=self.device) + stage_1_conditionings = build_video_conditionings(height // 2, width // 2) + video_state, _ = self.stage( + denoiser=SimpleDenoiser(video_context, audio_context), + sigmas=stage_1_sigmas_t, + noiser=noiser, + width=width // 2, + height=height // 2, + frames=num_frames, + fps=frame_rate, + video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings), + audio=ModalitySpec(context=audio_context), + ) + + # Stage 2: upsample + refine. + upscaled = self.upsampler(video_state.latent[:1]) + stage_2_sigmas_t = stage_2_sigmas.to(dtype=torch.float32, device=self.device) + stage_2_conditionings = build_video_conditionings(height, width) + video_state, _ = self.stage( + denoiser=SimpleDenoiser(video_context, audio_context), + sigmas=stage_2_sigmas_t, + noiser=noiser, + width=width, + height=height, + frames=num_frames, + fps=frame_rate, + video=ModalitySpec( + context=video_context, + conditionings=stage_2_conditionings, + noise_scale=stage_2_sigmas_t[0].item(), + initial_latent=upscaled, + ), + audio=ModalitySpec(context=audio_context), + ) + + decoded_video = self.video_decoder(video_state.latent, tiling_config, generator) + return decoded_video, None + + +@torch.inference_mode() +def main() -> None: + logging.basicConfig(level=logging.INFO) + params = resolve_cli_params(distilled=True) + parser = scail_animation_arg_parser(params=params) + args = parser.parse_args() + + if not args.lora or len(args.lora) != 1: + raise ValueError("SCAIL animation requires exactly one --lora (the SCAIL adapter).") + + pipeline = ScailAnimationPipeline( + distilled_checkpoint_path=args.distilled_checkpoint_path, + gemma_root=args.gemma_root, + spatial_upsampler_path=args.spatial_upsampler_path, + scail_lora=args.lora[0], + quantization=args.quantization, + compilation_config=args.compile, + offload_mode=args.offload_mode, + ) + tiling_config = TilingConfig.default() + output_shape = VideoPixelShape( + batch=1, frames=args.num_frames, width=args.width, height=args.height, fps=args.frame_rate + ) + video_chunks_number = get_video_chunks_number(output_shape.frames, tiling_config) + video, _audio = pipeline( + prompt=args.prompt, + seed=args.seed, + height=args.height, + width=args.width, + num_frames=args.num_frames, + frame_rate=args.frame_rate, + driving_video_path=args.driving_video, + mask_path=args.mask_path, + images=args.images if hasattr(args, "images") else [], + mode=DrivingMode(args.mode), + driving_strength=args.driving_strength, + enhance_prompt=args.enhance_prompt, + tiling_config=tiling_config, + ) + encode_video( + video=video, + fps=int(args.frame_rate), + audio=None, + output_path=args.output_path, + video_chunks_number=video_chunks_number, + ) + + +if __name__ == "__main__": + main() diff --git a/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py b/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py index 405da47..5041cf0 100644 --- a/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py +++ b/packages/ltx-pipelines/src/ltx_pipelines/utils/args.py @@ -809,3 +809,40 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS ), ) return parser + + +def scail_animation_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser: + """Argument parser for the SCAIL-2 character-animation pipeline (distilled, two-stage). + + Extends the standard two-stage distilled parser (checkpoint paths, --prompt, --num-frames, + --frame-rate, --height, --width, --image, --lora, --spatial-upsampler-path) with the SCAIL + driving/mask conditioning inputs. The single --lora is the SCAIL adapter. + """ + parser = default_2_stage_distilled_arg_parser(params=params) + parser.add_argument( + "--driving-video", + type=resolve_existing_path, + required=True, + help="Driving video file whose motion is transferred to the animated character.", + ) + parser.add_argument( + "--mask-path", + type=resolve_existing_path, + default=None, + help="Optional .pt file with semantic masks [K+1, F_pix, H, W] (ch0 = environment switch, " + "1..K = character binding slots). Omit for a driving-only (mask-free) model.", + ) + parser.add_argument( + "--mode", + type=str, + choices=["animation", "replacement"], + default="animation", + help="SCAIL conditioning mode (default: animation).", + ) + parser.add_argument( + "--driving-strength", + type=float, + default=1.0, + help="Driving conditioning strength; 1.0 keeps the driving latent clean/frozen (default: 1.0).", + ) + return parser