Add inference code and attention settings for sfp4 checkpoint-750

backend_snapshot/README.md

@@ -13,6 +13,17 @@ Key runtime settings:
 - `VSA_WARMUP_STEPS=0`
 - tile size: `4 x 4 x 4 = 64` video tokens
 
+Training attention semantics:
+
+- Video self-attention uses `SPARSE_FP4_OURS_P_ATTN`.
+- Cross-attention is not quantized/sparse in this backend. It falls back to
+  dense SDPA when `query_length != key_length`.
+- `force_dense` paths also use dense SDPA.
+- Q/K/V fake quantization uses FP4 with STE and no q/k mean subtraction.
+- Selected sparse tiles use group-local P quantization in the Triton kernel.
+- Dropped VSA tiles use tile-level q_mean/k_mean score plus mean_v
+  compensation.
+
 Important files:
 
 - `fastvideo/attention/backends/sparse_fp4_ours_p_attn.py`: Python attention backend, Q/K/V fake quantization, top-k block map, tile mean setup.
@@ -25,6 +36,25 @@ Important files:
 - `fastvideo/training/training_pipeline.py` and `fastvideo/training/wan_training_pipeline.py`: legacy SFT training path used by the launch script.
 - `scripts/training/run_sparse_fp4_train_v4_1n_sparse09_hpo_on_ours_p_init2050_interactive.sh`: exact Slurm wrapper for this run.
 - `scripts/training/run_sparse_fp4_train_v4_common.sh`: common SFT launch/resume script.
+- `training_attention_settings.json`: structured attention/training settings
+  for this checkpoint.
+- `scripts/inference/run_sfp4_ours_p_checkpoint_750.sh`: inference example
+  for the uploaded transformer checkpoint.
+- `fastvideo/entrypoints/cli/generate.py`, `fastvideo/entrypoints/video_generator.py`,
+  `fastvideo/pipelines/basic/wan/wan_pipeline.py`, and
+  `fastvideo/pipelines/stages/denoising.py`: `fastvideo generate` inference
+  path used by the example script.
+
+Example inference flow:
+
+```bash
+hf download yitongl/sparse_quant_exp \
+  --repo-type model \
+  --local-dir checkpoints/hf_download/sparse_quant_exp \
+  --include 'transformer/*'
+
+bash backend_snapshot/scripts/inference/run_sfp4_ours_p_checkpoint_750.sh
+```
 
 Source repo HEAD when staged:
 

backend_snapshot/examples/inference/basic/basic.py

@@ -0,0 +1,47 @@
+from fastvideo import VideoGenerator
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+    )
+
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param.num_frames = 45
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+
+
+if __name__ == "__main__":
+    main()

backend_snapshot/fastvideo/api/compat.py

@@ -0,0 +1,503 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from collections.abc import Mapping
+from copy import deepcopy
+from dataclasses import fields, is_dataclass
+from pathlib import Path
+from typing import Any
+
+from fastvideo.api.overrides import apply_overrides, parse_cli_overrides
+from fastvideo.api.parser import config_to_dict, load_raw_config, parse_config
+from fastvideo.api.schema import (
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OutputConfig,
+    RequestRuntimeConfig,
+    SamplingConfig,
+)
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.utils import shallow_asdict
+
+_EXPLICIT_REQUEST_ATTR = "_fastvideo_explicit_request"
+_INPUT_FIELD_NAMES = {field.name for field in fields(InputConfig)}
+_SAMPLING_FIELD_NAMES = {field.name for field in fields(SamplingConfig)}
+_RUNTIME_FIELD_NAMES = {field.name for field in fields(RequestRuntimeConfig)}
+_OUTPUT_FIELD_NAMES = {field.name for field in fields(OutputConfig)}
+_MISSING = object()
+_LEGACY_REQUEST_ALIASES = {
+    "neg_prompt": "negative_prompt",
+}
+_REQUEST_PIPELINE_OVERRIDE_FIELDS = frozenset({
+    "embedded_cfg_scale",
+})
+
+
+def normalize_generator_config(config: GeneratorConfig | Mapping[str, Any], ) -> GeneratorConfig:
+    if isinstance(config, GeneratorConfig):
+        return config
+    return parse_config(GeneratorConfig, config)
+
+
+def load_generator_config_from_file(
+    path: str | Path,
+    overrides: list[str] | Mapping[str, Any] | None = None,
+) -> GeneratorConfig:
+    raw = load_raw_config(path)
+    normalized_overrides = _normalize_overrides(overrides)
+
+    if _looks_like_run_or_serve_config(raw):
+        if normalized_overrides:
+            raw = apply_overrides(raw, normalized_overrides)
+        return parse_config(GeneratorConfig, raw["generator"])
+
+    if normalized_overrides:
+        adjusted = normalized_overrides
+        if all(key.startswith("generator.") for key in adjusted):
+            adjusted = {key[len("generator."):]: value for key, value in adjusted.items()}
+        raw = apply_overrides(raw, adjusted)
+
+    return parse_config(GeneratorConfig, raw)
+
+
+def legacy_from_pretrained_to_config(
+    model_path: str,
+    kwargs: Mapping[str, Any],
+) -> GeneratorConfig:
+    raw: dict[str, Any] = {"model_path": model_path}
+    engine: dict[str, Any] = {}
+    parallelism: dict[str, Any] = {}
+    offload: dict[str, Any] = {}
+    compile_config: dict[str, Any] = {}
+    pipeline: dict[str, Any] = {}
+    components: dict[str, Any] = {}
+    quantization: dict[str, Any] = {}
+    experimental: dict[str, Any] = {}
+
+    for key, value in kwargs.items():
+        if key == "revision":
+            raw["revision"] = value
+        elif key == "trust_remote_code":
+            raw["trust_remote_code"] = value
+        elif key == "num_gpus":
+            engine["num_gpus"] = value
+        elif key == "distributed_executor_backend":
+            engine["execution_backend"] = value
+        elif key in {"tp_size", "sp_size", "hsdp_replicate_dim", "hsdp_shard_dim", "dist_timeout"}:
+            parallelism[key] = value
+        elif key == "dit_cpu_offload":
+            offload["dit"] = value
+        elif key == "dit_layerwise_offload":
+            offload["dit_layerwise"] = value
+        elif key == "text_encoder_cpu_offload":
+            offload["text_encoder"] = value
+        elif key == "image_encoder_cpu_offload":
+            offload["image_encoder"] = value
+        elif key == "vae_cpu_offload":
+            offload["vae"] = value
+        elif key == "pin_cpu_memory":
+            offload["pin_cpu_memory"] = value
+        elif key == "enable_torch_compile":
+            compile_config["enabled"] = value
+        elif key == "torch_compile_kwargs":
+            compile_config["kwargs"] = deepcopy(value)
+        elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast"}:
+            engine[key] = value
+        elif key == "override_text_encoder_quant":
+            quantization["text_encoder_quant"] = value
+        elif key == "transformer_quant":
+            quantization["transformer_quant"] = value
+        elif key == "workload_type":
+            pipeline["workload_type"] = value
+        elif key == "lora_path":
+            components["lora_path"] = value
+        elif key == "override_pipeline_cls_name":
+            components["override_pipeline_cls_name"] = value
+        elif key == "override_transformer_cls_name":
+            components["override_transformer_cls_name"] = value
+        elif key == "pipeline_config":
+            if isinstance(value, str):
+                components["pipeline_config_path"] = value
+            else:
+                experimental[key] = deepcopy(value)
+        elif key == "override_text_encoder_safetensors":
+            components["text_encoder_weights"] = value
+        elif key == "init_weights_from_safetensors":
+            components["transformer_weights"] = value
+        elif key == "init_weights_from_safetensors_2":
+            components["transformer_2_weights"] = value
+        else:
+            experimental[key] = deepcopy(value)
+
+    if parallelism:
+        engine["parallelism"] = parallelism
+    if offload:
+        engine["offload"] = offload
+    if compile_config:
+        engine["compile"] = compile_config
+    if quantization:
+        engine["quantization"] = quantization
+    if engine:
+        raw["engine"] = engine
+
+    if components:
+        pipeline["components"] = components
+    if experimental:
+        pipeline["experimental"] = experimental
+    if pipeline:
+        raw["pipeline"] = pipeline
+
+    return parse_config(GeneratorConfig, raw)
+
+
+def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, Any], ) -> FastVideoArgs:
+    normalized = normalize_generator_config(config)
+    unsupported = []
+    if normalized.pipeline.profile is not None:
+        unsupported.append("pipeline.profile")
+    if normalized.pipeline.profile_version is not None:
+        unsupported.append("pipeline.profile_version")
+    if normalized.pipeline.components.config_root is not None:
+        unsupported.append("pipeline.components.config_root")
+    if normalized.pipeline.components.vae_weights is not None:
+        unsupported.append("pipeline.components.vae_weights")
+    if normalized.pipeline.components.upsampler_weights is not None:
+        unsupported.append("pipeline.components.upsampler_weights")
+    if unsupported:
+        joined = ", ".join(unsupported)
+        raise NotImplementedError(f"VideoGenerator compatibility adapter does not support {joined} yet")
+
+    engine = normalized.engine
+    kwargs: dict[str, Any] = {
+        "model_path": normalized.model_path,
+        "revision": normalized.revision,
+        "trust_remote_code": normalized.trust_remote_code,
+        "num_gpus": engine.num_gpus,
+        "distributed_executor_backend": engine.execution_backend,
+        "tp_size": engine.parallelism.tp_size,
+        "sp_size": engine.parallelism.sp_size,
+        "hsdp_replicate_dim": engine.parallelism.hsdp_replicate_dim,
+        "hsdp_shard_dim": engine.parallelism.hsdp_shard_dim,
+        "dist_timeout": engine.parallelism.dist_timeout,
+        "dit_cpu_offload": engine.offload.dit,
+        "dit_layerwise_offload": engine.offload.dit_layerwise,
+        "text_encoder_cpu_offload": engine.offload.text_encoder,
+        "image_encoder_cpu_offload": engine.offload.image_encoder,
+        "vae_cpu_offload": engine.offload.vae,
+        "pin_cpu_memory": engine.offload.pin_cpu_memory,
+        "enable_torch_compile": engine.compile.enabled,
+        "torch_compile_kwargs": deepcopy(engine.compile.kwargs),
+        "enable_stage_verification": engine.enable_stage_verification,
+        "use_fsdp_inference": engine.use_fsdp_inference,
+        "disable_autocast": engine.disable_autocast,
+    }
+    if normalized.pipeline.workload_type is not None:
+        kwargs["workload_type"] = normalized.pipeline.workload_type
+
+    quantization = engine.quantization
+    if quantization is not None and quantization.text_encoder_quant is not None:
+        kwargs["override_text_encoder_quant"] = quantization.text_encoder_quant
+    if quantization is not None and quantization.transformer_quant is not None:
+        kwargs["transformer_quant"] = quantization.transformer_quant
+
+    components = normalized.pipeline.components
+    if components.pipeline_config_path is not None:
+        kwargs["pipeline_config"] = components.pipeline_config_path
+    if components.lora_path is not None:
+        kwargs["lora_path"] = components.lora_path
+    if components.override_pipeline_cls_name is not None:
+        kwargs["override_pipeline_cls_name"] = components.override_pipeline_cls_name
+    if components.override_transformer_cls_name is not None:
+        kwargs["override_transformer_cls_name"] = components.override_transformer_cls_name
+    if components.text_encoder_weights is not None:
+        kwargs["override_text_encoder_safetensors"] = components.text_encoder_weights
+    if components.transformer_weights is not None:
+        kwargs["init_weights_from_safetensors"] = components.transformer_weights
+    if components.transformer_2_weights is not None:
+        kwargs["init_weights_from_safetensors_2"] = components.transformer_2_weights
+
+    kwargs.update(deepcopy(normalized.pipeline.profile_overrides))
+    kwargs.update(deepcopy(normalized.pipeline.experimental))
+    return FastVideoArgs.from_kwargs(**kwargs)
+
+
+def normalize_generation_request(request: GenerationRequest | Mapping[str, Any], ) -> GenerationRequest:
+    normalized = (request if isinstance(request, GenerationRequest) else parse_config(GenerationRequest, request))
+
+    if not hasattr(normalized, _EXPLICIT_REQUEST_ATTR):
+        setattr(normalized, _EXPLICIT_REQUEST_ATTR, _serialize_generation_request(normalized))
+    return normalized
+
+
+def legacy_generate_call_to_request(
+    prompt: str | None,
+    sampling_param: SamplingParam | None,
+    *,
+    mouse_cond: Any | None = None,
+    keyboard_cond: Any | None = None,
+    grid_sizes: Any | None = None,
+    legacy_kwargs: Mapping[str, Any] | None = None,
+) -> GenerationRequest:
+    raw = _sampling_param_to_request_raw(sampling_param)
+    if prompt is not None:
+        raw["prompt"] = prompt
+
+    for key, value in (legacy_kwargs or {}).items():
+        _apply_request_field(raw, key, value)
+
+    if mouse_cond is not None:
+        raw.setdefault("inputs", {})["mouse_cond"] = mouse_cond
+    if keyboard_cond is not None:
+        raw.setdefault("inputs", {})["keyboard_cond"] = keyboard_cond
+    if grid_sizes is not None:
+        raw.setdefault("inputs", {})["grid_sizes"] = grid_sizes
+
+    normalized = parse_config(GenerationRequest, raw)
+    setattr(normalized, _EXPLICIT_REQUEST_ATTR, deepcopy(raw))
+    return normalized
+
+
+def request_to_sampling_param(
+    request: GenerationRequest,
+    *,
+    model_path: str,
+) -> SamplingParam:
+    if request.plan is not None:
+        raise NotImplementedError("GenerationRequest.plan is not wired into VideoGenerator yet")
+    if request.state is not None:
+        raise NotImplementedError("GenerationRequest.state is not wired into VideoGenerator yet")
+
+    sampling_param = SamplingParam.from_pretrained(model_path)
+    updates = _explicit_request_updates(request)
+
+    for key, value in updates.items():
+        if hasattr(sampling_param, key):
+            setattr(sampling_param, key, deepcopy(value))
+        elif key in _REQUEST_PIPELINE_OVERRIDE_FIELDS or _is_supported_as_default_only(key, value):
+            continue
+        else:
+            raise ValueError(f"Request field {key!r} is not supported by sampling params for {model_path}")
+
+    sampling_param.__post_init__()
+    sampling_param.check_sampling_param()
+    return sampling_param
+
+
+def expand_request_prompt_batch(request: GenerationRequest, ) -> list[GenerationRequest]:
+    if not isinstance(request.prompt, list):
+        return [request]
+
+    requests: list[GenerationRequest] = []
+    for index, prompt in enumerate(request.prompt):
+        single_request = deepcopy(request)
+        single_request.prompt = prompt
+        _fan_out_batched_input_value(request, single_request, "image_path", index)
+        _fan_out_batched_input_value(request, single_request, "video_path", index)
+        _fan_out_explicit_request_metadata(request, single_request, index, prompt)
+        requests.append(single_request)
+    return requests
+
+
+def _looks_like_run_or_serve_config(raw: Mapping[str, Any]) -> bool:
+    return isinstance(raw.get("generator"), Mapping)
+
+
+def _normalize_overrides(overrides: list[str] | Mapping[str, Any] | None, ) -> dict[str, Any] | None:
+    if not overrides:
+        return None
+    if isinstance(overrides, list):
+        return parse_cli_overrides(overrides)
+    return dict(overrides)
+
+
+def _sampling_param_to_request_raw(sampling_param: SamplingParam | None, ) -> dict[str, Any]:
+    if sampling_param is None:
+        return {}
+
+    raw: dict[str, Any] = {}
+    for key, value in shallow_asdict(sampling_param).items():
+        if key == "prompt":
+            continue
+        _apply_request_field(raw, key, deepcopy(value))
+    return raw
+
+
+def _apply_request_field(
+    raw: dict[str, Any],
+    key: str,
+    value: Any,
+) -> None:
+    key = _LEGACY_REQUEST_ALIASES.get(key, key)
+    if key == "negative_prompt":
+        raw["negative_prompt"] = value
+        return
+    if key in _INPUT_FIELD_NAMES:
+        raw.setdefault("inputs", {})[key] = value
+        return
+    if key in _SAMPLING_FIELD_NAMES:
+        raw.setdefault("sampling", {})[key] = value
+        return
+    if key in _RUNTIME_FIELD_NAMES:
+        raw.setdefault("runtime", {})[key] = value
+        return
+    if key in _OUTPUT_FIELD_NAMES:
+        raw.setdefault("output", {})[key] = value
+        return
+    raw.setdefault("extensions", {})[key] = value
+
+
+def request_to_pipeline_overrides(request: GenerationRequest) -> dict[str, Any]:
+    overrides: dict[str, Any] = {}
+    for key, value in _explicit_request_updates(request).items():
+        if key in _REQUEST_PIPELINE_OVERRIDE_FIELDS:
+            overrides[key] = deepcopy(value)
+    return overrides
+
+
+def _explicit_request_updates(request: GenerationRequest) -> dict[str, Any]:
+    raw = getattr(request, _EXPLICIT_REQUEST_ATTR, None)
+    if raw is None:
+        raw = _serialize_generation_request(request)
+
+    return _extract_request_updates(raw)
+
+
+def _extract_request_updates(raw: Mapping[str, Any]) -> dict[str, Any]:
+    updates: dict[str, Any] = {}
+    if "negative_prompt" in raw:
+        updates["negative_prompt"] = deepcopy(raw["negative_prompt"])
+
+    for section_name in ("inputs", "sampling", "runtime", "output"):
+        section = raw.get(section_name)
+        if not isinstance(section, Mapping):
+            continue
+        for key, value in section.items():
+            updates[key] = deepcopy(value)
+
+    stage_overrides = raw.get("stage_overrides")
+    if stage_overrides:
+        updates.update(_flatten_stage_overrides(stage_overrides))
+
+    extensions = raw.get("extensions")
+    if isinstance(extensions, Mapping):
+        for key, value in extensions.items():
+            updates[key] = deepcopy(value)
+
+    return updates
+
+
+def _flatten_stage_overrides(stage_overrides: Any) -> dict[str, Any]:
+    if not isinstance(stage_overrides, Mapping):
+        raise ValueError("GenerationRequest.stage_overrides must be a mapping")
+
+    flattened: dict[str, Any] = {}
+    for stage_name, overrides in stage_overrides.items():
+        if not isinstance(overrides, Mapping):
+            raise ValueError(f"GenerationRequest.stage_overrides.{stage_name} must be a mapping")
+        for key, value in overrides.items():
+            if key in flattened and flattened[key] != value:
+                raise ValueError(f"Conflicting stage override for {key!r} across stages")
+            flattened[key] = deepcopy(value)
+    return flattened
+
+
+def _serialize_generation_request(request: GenerationRequest) -> dict[str, Any]:
+    return deepcopy(config_to_dict(request))
+
+
+def _fan_out_batched_input_value(
+    source_request: GenerationRequest,
+    target_request: GenerationRequest,
+    field_name: str,
+    index: int,
+) -> None:
+    value = getattr(source_request.inputs, field_name)
+    if not isinstance(value, list):
+        return
+    _validate_batched_input_length(source_request.prompt, value, field_name)
+    setattr(target_request.inputs, field_name, deepcopy(value[index]))
+
+
+def _fan_out_explicit_request_metadata(
+    source_request: GenerationRequest,
+    target_request: GenerationRequest,
+    index: int,
+    prompt: str,
+) -> None:
+    raw = getattr(source_request, _EXPLICIT_REQUEST_ATTR, None)
+    if raw is None:
+        return
+
+    raw = deepcopy(raw)
+    raw["prompt"] = prompt
+    inputs = raw.get("inputs")
+    if isinstance(inputs, dict):
+        for field_name in ("image_path", "video_path"):
+            value = inputs.get(field_name)
+            if isinstance(value, list):
+                _validate_batched_input_length(source_request.prompt, value, field_name)
+                inputs[field_name] = deepcopy(value[index])
+
+    setattr(target_request, _EXPLICIT_REQUEST_ATTR, raw)
+
+
+def _validate_batched_input_length(
+    prompts: str | list[str] | None,
+    values: list[Any],
+    field_name: str,
+) -> None:
+    if not isinstance(prompts, list):
+        return
+    if len(values) != len(prompts):
+        raise ValueError(f"GenerationRequest.inputs.{field_name} must have the same length as request.prompt")
+
+
+def _is_supported_as_default_only(key: str, value: Any) -> bool:
+    default_value = _DEFAULT_REQUEST_UPDATES.get(key, _MISSING)
+    return default_value is not _MISSING and _values_equal(value, default_value)
+
+
+def _collect_non_default_fields(
+    value: Any,
+    default: Any,
+) -> dict[str, Any]:
+    if not (is_dataclass(value) and is_dataclass(default)):
+        return {}
+
+    result: dict[str, Any] = {}
+    for field in fields(value):
+        current = getattr(value, field.name)
+        default_value = getattr(default, field.name)
+        if is_dataclass(current) and is_dataclass(default_value):
+            nested = _collect_non_default_fields(current, default_value)
+            if nested:
+                result[field.name] = nested
+            continue
+        if not _values_equal(current, default_value):
+            result[field.name] = deepcopy(current)
+    return result
+
+
+def _values_equal(left: Any, right: Any) -> bool:
+    if left is right:
+        return True
+    try:
+        return bool(left == right)
+    except Exception:
+        return False
+
+
+_DEFAULT_REQUEST_UPDATES = _extract_request_updates(config_to_dict(GenerationRequest()))
+
+__all__ = [
+    "generator_config_to_fastvideo_args",
+    "legacy_from_pretrained_to_config",
+    "legacy_generate_call_to_request",
+    "load_generator_config_from_file",
+    "normalize_generation_request",
+    "normalize_generator_config",
+    "request_to_pipeline_overrides",
+    "request_to_sampling_param",
+]

backend_snapshot/fastvideo/configs/pipelines/wan.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig
+from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config,
+                                               WAN2_1ControlCLIPVisionConfig)
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+    assert torch.isnan(hidden_state).sum() == 0
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack(
+        [torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
+    return prompt_embeds_tensor
+
+
+@dataclass
+class WanT2V480PConfig(PipelineConfig):
+    """Base configuration for Wan T2V 1.3B pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=WanVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Denoising stage
+    flow_shift: float | None = 3.0
+
+    # Text encoding stage
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (t5_postprocess_text, ))
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
+
+    # self-forcing params
+    warp_denoising_step: bool = True
+
+    # WanConfig-specific added parameters
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class WanT2V720PConfig(WanT2V480PConfig):
+    """Base configuration for Wan T2V 14B 720P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: float | None = 5.0
+
+
+@dataclass
+class WanI2V480PConfig(WanT2V480PConfig):
+    """Base configuration for Wan I2V 14B 480P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Precision for each component
+    image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig)
+    image_encoder_precision: str = "fp32"
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class WanI2V720PConfig(WanI2V480PConfig):
+    """Base configuration for Wan I2V 14B 720P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: float | None = 5.0
+
+
+@dataclass
+class WANV2VConfig(WanI2V480PConfig):
+    """Configuration for WAN2.1 1.3B Control pipeline."""
+
+    image_encoder_config: EncoderConfig = field(default_factory=WAN2_1ControlCLIPVisionConfig)
+    # CLIP encoder precision
+    image_encoder_precision: str = 'bf16'
+
+
+@dataclass
+class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
+    """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: float | None = 8.0
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
+
+
+@dataclass
+class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
+    flow_shift: float | None = 5.0
+    ti2v_task: bool = True
+    expand_timesteps: bool = True
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.dit_config.expand_timesteps = self.expand_timesteps
+
+
+@dataclass
+class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
+    flow_shift: float | None = 5.0
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
+
+
+@dataclass
+class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
+    flow_shift: float | None = 12.0
+    boundary_ratio: float | None = 0.875
+
+    # self-forcing params
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250])
+    warp_denoising_step: bool = True
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio
+
+
+@dataclass
+class Wan2_2_I2V_A14B_Config(WanI2V480PConfig):
+    flow_shift: float | None = 5.0
+    boundary_ratio: float | None = 0.900
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.dit_config.boundary_ratio = self.boundary_ratio
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
+    is_causal: bool = True
+    flow_shift: float | None = 5.0
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250])
+    warp_denoising_step: bool = True
+
+
+@dataclass
+class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config):
+    is_causal: bool = True
+    flow_shift: float | None = 12.0
+    boundary_ratio: float | None = 0.875
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 850, 700, 550, 350, 275, 200, 125])
+    warp_denoising_step: bool = True
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+
+# =============================================
+# ============= Matrix Game ===================
+# =============================================
+@dataclass
+class MatrixGameBaseI2V480PConfig(WanI2V480PConfig):
+    dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
+    flow_shift: float | None = 5.0
+
+
+@dataclass
+class MatrixGameI2V480PConfig(WanI2V480PConfig):
+    dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
+
+    image_encoder_config: EncoderConfig = field(default_factory=WAN2_1ControlCLIPVisionConfig)
+
+    is_causal: bool = True
+    flow_shift: float | None = 5.0
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 666, 333])
+    warp_denoising_step: bool = True
+    context_noise: int = 0
+    num_frames_per_block: int = 3
+    # sliding_window_num_frames: int = 15

backend_snapshot/fastvideo/configs/sample/base.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Any
+
+from fastvideo.logger import init_logger
+from fastvideo.utils import StoreBoolean
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SamplingParam:
+    """
+    Sampling parameters for video generation.
+    """
+    # All fields below are copied from ForwardBatch
+    data_type: str = "video"
+
+    # Image inputs
+    image_path: str | None = None
+    pil_image: Any | None = None
+
+    # Video inputs
+    video_path: str | None = None
+
+    # Action control inputs (Matrix-Game)
+    mouse_cond: Any | None = None  # Shape: (B, T, 2)
+    keyboard_cond: Any | None = None  # Shape: (B, T, K)
+    grid_sizes: Any | None = None  # Shape: (3,) [F,H,W]
+
+    # Camera control inputs (HYWorld)
+    pose: str | None = None  # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
+
+    # Camera control inputs (LingBotWorld)
+    c2ws_plucker_emb: Any | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
+
+    # Refine inputs (LongCat 480p->720p upscaling)
+    # Path-based refine (load stage1 video from disk, e.g. MP4)
+    refine_from: str | None = None  # Path to stage1 video (480p output from distill)
+    t_thresh: float = 0.5  # Threshold for timestep scheduling in refinement
+    spatial_refine_only: bool = False  # If True, only spatial (no temporal doubling)
+    num_cond_frames: int = 0  # Number of conditioning frames
+    # In-memory refine input (for two-stage pipeline where stage1 frames are already in memory)
+    # This mirrors LongCat's demo where a list of frames (e.g. np.ndarray or PIL.Image)
+    # is passed directly to the refinement pipeline instead of reloading from disk.
+    stage1_video: Any | None = None
+
+    # Text inputs
+    prompt: str | list[str] | None = None
+    negative_prompt: str | None = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+    prompt_path: str | None = None
+    output_path: str = "outputs/"
+    output_video_name: str | None = None
+
+    # Batch info
+    num_videos_per_prompt: int = 1
+    seed: int = 1024
+
+    # Original dimensions (before VAE scaling)
+    num_frames: int = 125
+    height: int = 720
+    width: int = 1280
+    height_sr: int = 1072
+    width_sr: int = 1920
+    fps: int = 24
+
+    # Denoising parameters
+    num_inference_steps: int = 50
+    num_inference_steps_sr: int = 50
+    guidance_scale: float = 1.0
+    guidance_scale_2: float | None = None
+    guidance_rescale: float = 0.0
+    boundary_ratio: float | None = None
+    sigmas: list[float] | None = None
+
+    # TeaCache parameters
+    enable_teacache: bool = False
+
+    # GEN3C camera control
+    trajectory_type: str | None = None
+    movement_distance: float | None = None
+    camera_rotation: str | None = None
+
+    # Misc
+    save_video: bool = True
+    return_frames: bool = True
+    return_trajectory_latents: bool = False  # returns all latents for each timestep
+    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
+
+    def __post_init__(self) -> None:
+        self.data_type = "video" if self.num_frames > 1 else "image"
+
+    def __getattr__(self, name: str) -> Any:
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def check_sampling_param(self) -> None:
+        if self.prompt_path and not self.prompt_path.endswith(".txt"):
+            raise ValueError("prompt_path must be a txt file")
+
+    def update(self, source_dict: dict[str, Any]) -> None:
+        for key, value in source_dict.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                logger.exception("%s has no attribute %s", type(self).__name__, key)
+
+        self.__post_init__()
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> "SamplingParam":
+        from fastvideo.registry import get_sampling_param_cls_for_name
+        sampling_cls = get_sampling_param_cls_for_name(model_path)
+        if sampling_cls is not None:
+            sampling_param: SamplingParam = sampling_cls()
+        else:
+            logger.warning("Couldn't find an optimal sampling param for %s. Using the default sampling param.",
+                           model_path)
+            sampling_param = cls()
+
+        return sampling_param
+
+    @staticmethod
+    def add_cli_args(parser: Any) -> Any:
+        """Add CLI arguments for SamplingParam fields"""
+        parser.add_argument(
+            "--prompt",
+            type=str,
+            default=SamplingParam.prompt,
+            help="Text prompt for video generation",
+        )
+        parser.add_argument(
+            "--negative-prompt",
+            type=str,
+            default=SamplingParam.negative_prompt,
+            help="Negative text prompt for video generation",
+        )
+        parser.add_argument(
+            "--prompt-path",
+            type=str,
+            default=SamplingParam.prompt_path,
+            help="Path to a text file containing the prompt",
+        )
+        parser.add_argument(
+            "--output-path",
+            type=str,
+            default=SamplingParam.output_path,
+            help="Path to save the generated video",
+        )
+        parser.add_argument(
+            "--output-video-name",
+            type=str,
+            default=SamplingParam.output_video_name,
+            help="Name of the output video",
+        )
+        parser.add_argument(
+            "--num-videos-per-prompt",
+            type=int,
+            default=SamplingParam.num_videos_per_prompt,
+            help="Number of videos to generate per prompt",
+        )
+        parser.add_argument(
+            "--seed",
+            type=int,
+            default=SamplingParam.seed,
+            help="Random seed for generation",
+        )
+        parser.add_argument(
+            "--num-frames",
+            type=int,
+            default=SamplingParam.num_frames,
+            help="Number of frames to generate",
+        )
+        parser.add_argument(
+            "--height",
+            type=int,
+            default=SamplingParam.height,
+            help="Height of generated video",
+        )
+        parser.add_argument(
+            "--width",
+            type=int,
+            default=SamplingParam.width,
+            help="Width of generated video",
+        )
+        parser.add_argument(
+            "--fps",
+            type=int,
+            default=SamplingParam.fps,
+            help="Frames per second for saved video",
+        )
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=SamplingParam.num_inference_steps,
+            help="Number of denoising steps",
+        )
+        parser.add_argument(
+            "--guidance-scale",
+            type=float,
+            default=SamplingParam.guidance_scale,
+            help="Classifier-free guidance scale",
+        )
+        parser.add_argument(
+            "--guidance-rescale",
+            type=float,
+            default=SamplingParam.guidance_rescale,
+            help="Guidance rescale factor",
+        )
+        parser.add_argument(
+            "--boundary-ratio",
+            type=float,
+            default=SamplingParam.boundary_ratio,
+            help="Boundary timestep ratio",
+        )
+        parser.add_argument(
+            "--save-video",
+            action="store_true",
+            default=SamplingParam.save_video,
+            help="Whether to save the video to disk",
+        )
+        parser.add_argument(
+            "--no-save-video",
+            action="store_false",
+            dest="save_video",
+            help="Don't save the video to disk",
+        )
+        parser.add_argument(
+            "--return-frames",
+            action="store_true",
+            default=False,
+            help="Whether to return the raw frames",
+        )
+        parser.add_argument(
+            "--image-path",
+            type=str,
+            default=SamplingParam.image_path,
+            help="Path to input image for image-to-video generation",
+        )
+        parser.add_argument(
+            "--video-path",
+            type=str,
+            default=SamplingParam.video_path,
+            help="Path to input video for video-to-video generation",
+        )
+        parser.add_argument(
+            "--refine-from",
+            type=str,
+            default=SamplingParam.refine_from,
+            help="Path to stage1 video for refinement (LongCat 480p->720p)",
+        )
+        parser.add_argument(
+            "--t-thresh",
+            type=float,
+            default=SamplingParam.t_thresh,
+            help="Threshold for timestep scheduling in refinement (default: 0.5)",
+        )
+        parser.add_argument(
+            "--spatial-refine-only",
+            action=StoreBoolean,
+            default=SamplingParam.spatial_refine_only,
+            help="Only perform spatial super-resolution (no temporal doubling)",
+        )
+        parser.add_argument(
+            "--num-cond-frames",
+            type=int,
+            default=SamplingParam.num_cond_frames,
+            help="Number of conditioning frames for refinement",
+        )
+        parser.add_argument(
+            "--moba-config-path",
+            type=str,
+            default=None,
+            help="Path to a JSON file containing V-MoBA specific configurations.",
+        )
+        parser.add_argument(
+            "--return-trajectory-latents",
+            action="store_true",
+            default=SamplingParam.return_trajectory_latents,
+            help="Whether to return the trajectory",
+        )
+        parser.add_argument(
+            "--return-trajectory-decoded",
+            action="store_true",
+            default=SamplingParam.return_trajectory_decoded,
+            help="Whether to return the decoded trajectory",
+        )
+        return parser
+
+
+@dataclass
+class CacheParams:
+    cache_type: str = "none"

backend_snapshot/fastvideo/configs/sample/wan.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class WanT2V_1_3B_SamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage
+    guidance_scale: float = 3.0
+    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+    num_inference_steps: int = 50
+
+
+@dataclass
+class WanT2V_14B_SamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage
+    guidance_scale: float = 5.0
+    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+    num_inference_steps: int = 50
+
+
+@dataclass
+class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam):
+    # Denoising stage
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 40
+
+
+@dataclass
+class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
+    # Denoising stage
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 40
+
+
+@dataclass
+class FastWanT2V480P_SamplingParam(WanT2V_1_3B_SamplingParam):
+    # DMD parameters
+    # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
+    num_inference_steps: int = 3
+    num_frames: int = 61
+    height: int = 448
+    width: int = 832
+    fps: int = 16
+
+
+# =============================================
+# ============= Wan2.1 Fun Models =============
+# =============================================
+@dataclass
+class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    guidance_scale: float = 6.0
+    num_inference_steps: int = 50
+
+
+@dataclass
+class Wan2_1_Fun_1_3B_Control_SamplingParam(SamplingParam):
+    fps: int = 16
+    num_frames: int = 49
+    height: int = 832
+    width: int = 480
+    guidance_scale: float = 6.0
+
+
+# =============================================
+# ============= Wan2.2 TI2V Models =============
+# =============================================
+@dataclass
+class Wan2_2_Base_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.2 TI2V 5B model."""
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+
+
+@dataclass
+class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam):
+    """Sampling parameters for Wan2.2 TI2V 5B model."""
+    height: int = 704
+    width: int = 1280
+    num_frames: int = 121
+    fps: int = 24
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 50
+
+
+@dataclass
+class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
+    guidance_scale: float = 4.0  # high_noise
+    guidance_scale_2: float = 3.0  # low_noise
+    num_inference_steps: int = 40
+    fps: int = 16
+    # NOTE(will): default boundary timestep is tracked by PipelineConfig, but
+    # can be overridden during sampling
+
+
+@dataclass
+class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
+    guidance_scale: float = 3.5  # high_noise
+    guidance_scale_2: float = 3.5  # low_noise
+    num_inference_steps: int = 40
+    fps: int = 16
+    # NOTE(will): default boundary timestep is tracked by PipelineConfig, but
+    # can be overridden during sampling
+
+
+@dataclass
+class Wan2_2_Fun_A14B_Control_SamplingParam(Wan2_1_Fun_1_3B_Control_SamplingParam):
+    num_frames: int = 81
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam(Wan2_1_Fun_1_3B_InP_SamplingParam):
+    pass
+
+
+@dataclass
+class SelfForcingWan2_2_T2V_A14B_480P_SamplingParam(Wan2_2_T2V_A14B_SamplingParam):
+    num_inference_steps: int = 8
+    num_frames: int = 81
+    height: int = 448
+    width: int = 832
+    fps: int = 16
+
+
+@dataclass
+class MatrixGame2_SamplingParam(SamplingParam):
+    height: int = 352
+    width: int = 640
+    num_frames: int = 57
+    fps: int = 25
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 3
+    negative_prompt: str | None = None

backend_snapshot/fastvideo/configs/wan_1.3B_t2v_pipeline.json

@@ -0,0 +1,40 @@
+{
+  "embedded_cfg_scale": 6.0,
+  "flow_shift": 3,
+  "dit_cpu_offload": true,
+  "disable_autocast": false,
+  "precision": "bf16",
+  "vae_precision": "fp32",
+  "vae_tiling": false,
+  "vae_sp": false,
+  "vae_config": {
+    "load_encoder": false,
+    "load_decoder": true,
+    "tile_sample_min_height": 256,
+    "tile_sample_min_width": 256,
+    "tile_sample_min_num_frames": 16,
+    "tile_sample_stride_height": 192,
+    "tile_sample_stride_width": 192,
+    "tile_sample_stride_num_frames": 12,
+    "blend_num_frames": 8,
+    "use_tiling": false,
+    "use_temporal_tiling": false,
+    "use_parallel_tiling": false,
+    "use_feature_cache": true
+  },
+  "dit_config": {
+    "prefix": "Wan",
+    "quant_config": null
+  },
+  "text_encoder_precisions": [
+    "fp32"
+  ],
+  "text_encoder_configs": [
+    {
+      "prefix": "t5",
+      "quant_config": null,
+      "lora_config": null
+    }
+  ],
+  "enable_torch_compile": false
+}

backend_snapshot/fastvideo/entrypoints/cli/generate.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
+
+import argparse
+import dataclasses
+import os
+from typing import cast
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample.base import SamplingParam
+from fastvideo.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.entrypoints.cli.utils import RaiseNotImplementedAction
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser
+
+logger = init_logger(__name__)
+
+
+class GenerateSubcommand(CLISubcommand):
+    """The `generate` subcommand for the FastVideo CLI"""
+
+    def __init__(self) -> None:
+        self.name = "generate"
+        super().__init__()
+        self.init_arg_names = self._get_init_arg_names()
+        self.generation_arg_names = self._get_generation_arg_names()
+
+    def _get_init_arg_names(self) -> list[str]:
+        """Get names of arguments for VideoGenerator initialization"""
+        return ["num_gpus", "tp_size", "sp_size", "model_path"]
+
+    def _get_generation_arg_names(self) -> list[str]:
+        """Get names of arguments for generate_video method"""
+        return [field.name for field in dataclasses.fields(SamplingParam)]
+
+    def cmd(self, args: argparse.Namespace) -> None:
+        excluded_args = ['subparser', 'config', 'dispatch_function']
+
+        provided_args = {}
+        for k, v in vars(args).items():
+            if (k not in excluded_args and v is not None and hasattr(args, '_provided') and k in args._provided):
+                provided_args[k] = v
+
+        if 'model_path' in vars(args) and args.model_path is not None:
+            provided_args['model_path'] = args.model_path
+
+        if 'prompt' in vars(args) and args.prompt is not None:
+            provided_args['prompt'] = args.prompt
+
+        merged_args = {**provided_args}
+
+        logger.info('CLI Args: %s', merged_args)
+
+        if 'model_path' not in merged_args or not merged_args['model_path']:
+            raise ValueError("model_path must be provided either in config file or via --model-path")
+
+        # Check if either prompt or prompt_txt is provided
+        has_prompt = 'prompt' in merged_args and merged_args['prompt']
+        has_prompt_txt = 'prompt_txt' in merged_args and merged_args['prompt_txt']
+
+        if not (has_prompt or has_prompt_txt):
+            raise ValueError("Either prompt or prompt_txt must be provided")
+
+        if has_prompt and has_prompt_txt:
+            raise ValueError("Cannot provide both 'prompt' and 'prompt_txt'. Use only one of them.")
+
+        init_args = {k: v for k, v in merged_args.items() if k not in self.generation_arg_names}
+        generation_args = {k: v for k, v in merged_args.items() if k in self.generation_arg_names}
+        generation_args.setdefault("return_frames", False)
+
+        model_path = init_args.pop('model_path')
+        prompt = generation_args.pop('prompt', None)
+
+        generator = VideoGenerator.from_pretrained(model_path=model_path, **init_args)
+
+        # Call generate_video - it handles both single and batch modes
+        generator.generate_video(prompt=prompt, **generation_args)
+
+    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
+        if args.num_gpus is not None and args.num_gpus <= 0:
+            raise ValueError("Number of gpus must be positive")
+
+        if args.config and not os.path.exists(args.config):
+            raise ValueError(f"Config file not found: {args.config}")
+
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        generate_parser = subparsers.add_parser(
+            "generate",
+            help="Run inference on a model",
+            usage="fastvideo generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]")
+
+        generate_parser.add_argument(
+            "--config",
+            type=str,
+            default='',
+            required=False,
+            help="Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional."
+        )
+
+        generate_parser = FastVideoArgs.add_cli_args(generate_parser)
+        generate_parser = SamplingParam.add_cli_args(generate_parser)
+
+        generate_parser.add_argument(
+            "--text-encoder-configs",
+            action=RaiseNotImplementedAction,
+            help="JSON array of text encoder configurations (NOT YET IMPLEMENTED)",
+        )
+
+        return cast(FlexibleArgumentParser, generate_parser)
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [GenerateSubcommand()]

backend_snapshot/fastvideo/entrypoints/video_generator.py

@@ -0,0 +1,797 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+VideoGenerator module for FastVideo.
+
+This module provides a consolidated interface for generating videos using
+diffusion models.
+"""
+
+import os
+import re
+import shutil
+import threading
+import time
+import tempfile
+import warnings
+from collections.abc import Mapping
+from copy import deepcopy
+from typing import Any
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+
+from fastvideo.api.compat import (
+    expand_request_prompt_batch,
+    generator_config_to_fastvideo_args,
+    legacy_from_pretrained_to_config,
+    load_generator_config_from_file,
+    normalize_generation_request,
+    normalize_generator_config,
+    request_to_pipeline_overrides,
+    request_to_sampling_param,
+)
+from fastvideo.api.results import GenerationResult
+from fastvideo.api.schema import GenerationRequest, GeneratorConfig
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ForwardBatch
+from fastvideo.utils import align_to, shallow_asdict
+from fastvideo.worker.executor import Executor
+
+logger = init_logger(__name__)
+
+_FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
+    "num_gpus",
+    "revision",
+    "trust_remote_code",
+    "distributed_executor_backend",
+    "tp_size",
+    "sp_size",
+    "hsdp_replicate_dim",
+    "hsdp_shard_dim",
+    "dist_timeout",
+    "use_fsdp_inference",
+    "disable_autocast",
+    "enable_stage_verification",
+    "dit_cpu_offload",
+    "dit_layerwise_offload",
+    "text_encoder_cpu_offload",
+    "image_encoder_cpu_offload",
+    "vae_cpu_offload",
+    "pin_cpu_memory",
+    "enable_torch_compile",
+    "torch_compile_kwargs",
+    "transformer_quant",
+})
+
+
+def _infer_latent_batch_size(batch: ForwardBatch) -> int:
+    if isinstance(batch.prompt, list):
+        latent_batch_size = len(batch.prompt)
+    elif batch.prompt is not None:
+        latent_batch_size = 1
+    elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
+        latent_batch_size = batch.prompt_embeds[0].shape[0]
+    else:
+        raise ValueError("Cannot infer batch size from batch; no prompt or prompt_embeds found")
+    latent_batch_size *= batch.num_videos_per_prompt
+    return latent_batch_size
+
+
+class VideoGenerator:
+    """
+    A unified class for generating videos using diffusion models.
+    
+    This class provides a simple interface for video generation with rich
+    customization options, similar to popular frameworks like HF Diffusers.
+    """
+
+    def __init__(
+        self,
+        fastvideo_args: FastVideoArgs,
+        executor_class: type[Executor],
+        log_stats: bool,
+        *,
+        log_queue=None,
+    ):
+        """
+        Initialize the video generator.
+
+        Args:
+            fastvideo_args: The inference arguments
+            executor_class: The executor class to use for inference
+            log_stats: Whether to log statistics
+            log_queue: Optional multiprocessing.Queue to forward worker logs to
+        """
+        self.config: GeneratorConfig | None = None
+        self.fastvideo_args = fastvideo_args
+        self.executor = executor_class(fastvideo_args, log_queue=log_queue)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str | GeneratorConfig | Mapping[str, Any] | None = None,
+        **kwargs,
+    ) -> "VideoGenerator":
+        """
+        Create a video generator from a pretrained model.
+        
+        Args:
+            model_path: Path or identifier for the pretrained model
+            pipeline_config: Pipeline config to use for inference
+            **kwargs: Additional arguments to customize model loading, set any FastVideoArgs or PipelineConfig attributes here.
+                
+        Returns:
+            The created video generator
+
+        Priority level: Default pipeline config < User's pipeline config < User's kwargs
+
+        Stable convenience kwargs remain supported here for common engine and
+        offload settings. Advanced model- or pipeline-specific options should
+        move to VideoGenerator.from_config(...).
+        """
+        log_queue = kwargs.pop("log_queue", None)
+        typed_config = kwargs.pop("config", None)
+        if typed_config is not None:
+            if model_path is not None:
+                raise TypeError("Pass either model_path or config to from_pretrained, not both")
+            if kwargs:
+                unexpected = ", ".join(sorted(kwargs))
+                raise TypeError(f"Unexpected keyword arguments with config: {unexpected}")
+            return cls.from_config(typed_config, log_queue=log_queue)
+
+        if isinstance(model_path, GeneratorConfig | Mapping):
+            if kwargs:
+                unexpected = ", ".join(sorted(kwargs))
+                raise TypeError(f"Unexpected keyword arguments with typed config: {unexpected}")
+            return cls.from_config(model_path, log_queue=log_queue)
+
+        if model_path is None:
+            raise TypeError("model_path or config is required")
+
+        legacy_only_kwargs = sorted(set(kwargs) - _FROM_PRETRAINED_CONVENIENCE_KWARGS)
+        if legacy_only_kwargs:
+            warnings.warn(
+                "VideoGenerator.from_pretrained(...) received legacy-only kwargs "
+                f"({', '.join(legacy_only_kwargs)}); prefer VideoGenerator.from_config(...) "
+                "for advanced configuration.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return cls.from_config(
+            legacy_from_pretrained_to_config(model_path, kwargs),
+            log_queue=log_queue,
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        config: GeneratorConfig | Mapping[str, Any],
+        *,
+        log_queue=None,
+    ) -> "VideoGenerator":
+        normalized = normalize_generator_config(config)
+        fastvideo_args = generator_config_to_fastvideo_args(normalized)
+        generator = cls.from_fastvideo_args(fastvideo_args, log_queue=log_queue)
+        generator.config = normalized
+        return generator
+
+    @classmethod
+    def from_file(
+        cls,
+        path: str,
+        overrides: list[str] | Mapping[str, Any] | None = None,
+        *,
+        log_queue=None,
+    ) -> "VideoGenerator":
+        return cls.from_config(
+            load_generator_config_from_file(path, overrides=overrides),
+            log_queue=log_queue,
+        )
+
+    @classmethod
+    def from_fastvideo_args(
+        cls,
+        fastvideo_args: FastVideoArgs,
+        *,
+        log_queue=None,
+    ) -> "VideoGenerator":
+        """
+        Create a video generator with the specified arguments.
+
+        Args:
+            fastvideo_args: The inference arguments
+            log_queue: Optional multiprocessing.Queue to forward worker logs to
+
+        Returns:
+            The created video generator
+        """
+        # Initialize distributed environment if needed
+        # initialize_distributed_and_parallelism(fastvideo_args)
+
+        executor_class = Executor.get_class(fastvideo_args)
+        return cls(
+            fastvideo_args=fastvideo_args,
+            executor_class=executor_class,
+            log_stats=False,  # TODO: implement
+            log_queue=log_queue,
+        )
+
+    def generate(
+        self,
+        request: GenerationRequest | Mapping[str, Any],
+        *,
+        log_queue=None,
+    ) -> GenerationResult | list[GenerationResult]:
+        """
+        Generate video or image outputs from a typed inference request.
+
+        Args:
+            request: A `GenerationRequest` instance or a mapping that can be
+                parsed into one. This is the primary public inference
+                entrypoint for the typed API.
+            log_queue: Optional multiprocessing.Queue to forward worker logs to
+                during this request.
+
+        Returns:
+            A `GenerationResult` for single-request generation, or a list of
+            `GenerationResult` objects when the request expands into multiple
+            prompts.
+        """
+        normalized_request = normalize_generation_request(request)
+        if log_queue:
+            self.executor.set_log_queue(log_queue)
+
+        try:
+            return self._generate_request_impl(normalized_request)
+        finally:
+            if log_queue:
+                self.executor.clear_log_queue()
+
+    def generate_video(
+        self,
+        prompt: str | None = None,
+        sampling_param: SamplingParam | None = None,
+        # Action control inputs (Matrix-Game)
+        mouse_cond: torch.Tensor | None = None,
+        keyboard_cond: torch.Tensor | None = None,
+        grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
+        | None = None,
+        **kwargs,
+    ) -> dict[str, Any] | list[dict[str, Any]]:
+        """
+        Generate a video based on the given prompt.
+        
+        Args:
+            prompt: The prompt to use for generation (optional if prompt_txt is provided)
+            negative_prompt: The negative prompt to use (overrides the one in fastvideo_args)
+            output_path: Path to save the video (overrides the one in fastvideo_args)
+            prompt_path: Path to prompt file
+            save_video: Whether to save the video to disk
+            return_frames: Whether to include raw frames in the result dict
+            num_inference_steps: Number of denoising steps (overrides fastvideo_args)
+            guidance_scale: Classifier-free guidance scale (overrides fastvideo_args)
+            num_frames: Number of frames to generate (overrides fastvideo_args)
+            height: Height of generated video (overrides fastvideo_args)
+            width: Width of generated video (overrides fastvideo_args)
+            fps: Frames per second for saved video (overrides fastvideo_args)
+            seed: Random seed for generation (overrides fastvideo_args)
+            callback: Callback function called after each step
+            callback_steps: Number of steps between each callback
+            
+        Returns:
+            A metadata dictionary for single-prompt generation, or a list of
+            metadata dictionaries for prompt-file batch generation.
+        """
+        log_queue = kwargs.pop("log_queue", None)
+        warnings.warn(
+            "VideoGenerator.generate_video(...) is deprecated; use "
+            "VideoGenerator.generate(request=...) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if log_queue:
+            self.executor.set_log_queue(log_queue)
+
+        try:
+            return self._generate_video_impl(
+                prompt=prompt,
+                sampling_param=sampling_param,
+                mouse_cond=mouse_cond,
+                keyboard_cond=keyboard_cond,
+                grid_sizes=grid_sizes,
+                **kwargs,
+            )
+        finally:
+            if log_queue:
+                self.executor.clear_log_queue()
+
+    def _generate_request_impl(
+        self,
+        request: GenerationRequest,
+    ) -> GenerationResult | list[GenerationResult]:
+        if isinstance(request.prompt, list):
+            if request.inputs.prompt_path is not None:
+                raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path")
+            results: list[GenerationResult] = []
+            for index, single_request in enumerate(expand_request_prompt_batch(request)):
+                prompt = single_request.prompt
+                wrapped = self._generate_single_request(single_request)
+                if isinstance(wrapped, list):
+                    results.extend(wrapped)
+                    continue
+                wrapped.prompt_index = index
+                if wrapped.prompt is None and isinstance(prompt, str):
+                    wrapped.prompt = prompt
+                results.append(wrapped)
+            return results
+
+        return self._generate_single_request(request)
+
+    def _generate_single_request(
+        self,
+        request: GenerationRequest,
+    ) -> GenerationResult | list[GenerationResult]:
+        fastvideo_args = self.fastvideo_args
+        pipeline_overrides = request_to_pipeline_overrides(request)
+        if pipeline_overrides:
+            fastvideo_args = deepcopy(self.fastvideo_args)
+            for key, value in pipeline_overrides.items():
+                if not hasattr(fastvideo_args.pipeline_config, key):
+                    raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
+                setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
+
+        sampling_param = request_to_sampling_param(
+            request,
+            model_path=self.fastvideo_args.model_path,
+        )
+        result = self._generate_video_impl(
+            prompt=request.prompt,
+            sampling_param=sampling_param,
+            fastvideo_args=fastvideo_args,
+        )
+        return self._wrap_legacy_result(result)
+
+    def _generate_video_impl(
+        self,
+        prompt: str | list[str] | None = None,
+        sampling_param: SamplingParam | None = None,
+        mouse_cond: torch.Tensor | None = None,
+        keyboard_cond: torch.Tensor | None = None,
+        grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
+        | None = None,
+        fastvideo_args: FastVideoArgs | None = None,
+        **kwargs,
+    ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]:
+        """Internal implementation of generate_video."""
+        if fastvideo_args is None:
+            fastvideo_args = self.fastvideo_args
+
+        # Handle batch processing from text file
+        if sampling_param is None:
+            sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
+
+        # Add action control inputs to kwargs if provided
+        if mouse_cond is not None:
+            kwargs['mouse_cond'] = mouse_cond
+        if keyboard_cond is not None:
+            kwargs['keyboard_cond'] = keyboard_cond
+        if grid_sizes is not None:
+            kwargs['grid_sizes'] = grid_sizes
+
+        sampling_param.update(kwargs)
+
+        if fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None:
+            prompt_txt_path = sampling_param.prompt_path or fastvideo_args.prompt_txt
+            if not prompt_txt_path or not os.path.exists(prompt_txt_path):
+                raise FileNotFoundError(f"Prompt text file not found: {prompt_txt_path}")
+
+            # Read prompts from file
+            with open(prompt_txt_path, encoding='utf-8') as f:
+                prompts = [line.strip() for line in f if line.strip()]
+
+            if not prompts:
+                raise ValueError(f"No prompts found in file: {prompt_txt_path}")
+
+            logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
+
+            results = []
+            for i, batch_prompt in enumerate(prompts):
+                logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
+                try:
+                    # Generate video for this prompt using the same logic below
+                    output_path = self._prepare_output_path(sampling_param.output_path, batch_prompt)
+                    kwargs["output_path"] = output_path
+                    result = self._generate_single_video(
+                        prompt=batch_prompt,
+                        sampling_param=sampling_param,
+                        fastvideo_args=fastvideo_args,
+                        **kwargs,
+                    )
+
+                    # Add prompt info to result
+                    result["prompt_index"] = i
+                    result["prompt"] = batch_prompt
+
+                    results.append(result)
+                    logger.info("Successfully generated video for prompt %d", i + 1)
+
+                except Exception as e:
+                    logger.error("Failed to generate video for prompt %d: %s", i + 1, e)
+                    continue
+
+            logger.info("Completed batch processing. Generated %d videos successfully.", len(results))
+            return results
+
+        # Single prompt generation (original behavior)
+        if prompt is None:
+            raise ValueError("Either prompt or prompt_txt must be provided")
+        if not isinstance(prompt, str):
+            raise ValueError("Single-prompt generation expects a string prompt")
+        output_path = self._prepare_output_path(sampling_param.output_path, prompt)
+        kwargs["output_path"] = output_path
+        return self._generate_single_video(
+            prompt=prompt,
+            sampling_param=sampling_param,
+            fastvideo_args=fastvideo_args,
+            **kwargs,
+        )
+
+    def _is_image_workload(self) -> bool:
+        """Return True when the workload produces a single image (t2i, i2i …)."""
+        args = getattr(self, "fastvideo_args", None)
+        if args is None:
+            return False
+        return args.workload_type.value.endswith("2i")
+
+    def _prepare_output_path(
+        self,
+        output_path: str,
+        prompt: str,
+    ) -> str:
+        """Build a unique, sanitized output file path.
+
+        The file extension is chosen automatically based on the workload type:
+        ``.png`` for image workloads (``t2i``, ``i2i``, …) and ``.mp4`` for
+        video workloads.
+
+        - If ``output_path`` already carries the correct extension, treat it
+          as a file path.
+        - Otherwise, treat ``output_path`` as a directory and derive the
+          filename from the prompt.
+        - Invalid filename characters are removed; if the name changes, a
+          warning is logged.
+        - If the target path already exists, a numeric suffix is appended.
+        """
+        target_ext = ".png" if self._is_image_workload() else ".mp4"
+
+        def _sanitize_filename_component(name: str) -> str:
+            # Remove characters invalid on common filesystems, strip spaces/dots
+            sanitized = re.sub(r'[\\/:*?"<>|]', '', name)
+            sanitized = sanitized.strip().strip('.')
+            sanitized = re.sub(r'\s+', ' ', sanitized)
+            return sanitized or "output"
+
+        base_path, extension = os.path.splitext(output_path)
+        extension_lower = extension.lower()
+
+        if extension_lower == target_ext:
+            output_dir = os.path.dirname(output_path)
+            base_name = os.path.basename(base_path)  # filename without extension
+            sanitized_base = _sanitize_filename_component(base_name)
+            if sanitized_base != base_name:
+                logger.warning(
+                    "The output name '%s' contained invalid characters. "
+                    "It has been renamed to '%s%s'",
+                    os.path.basename(output_path),
+                    sanitized_base,
+                    target_ext,
+                )
+            out_name = f"{sanitized_base}{target_ext}"
+        else:
+            # Treat as directory; inform if an unexpected extension was
+            # provided.
+            if extension:
+                logger.info(
+                    "Output path '%s' has extension '%s' which does not "
+                    "match the target '%s'; treating it as a directory",
+                    output_path,
+                    extension,
+                    target_ext,
+                )
+            output_dir = output_path
+            prompt_component = _sanitize_filename_component(prompt[:100])
+            out_name = f"{prompt_component}{target_ext}"
+
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+
+        new_output_path = os.path.join(output_dir, out_name)
+        counter = 1
+        while os.path.exists(new_output_path):
+            name_part, ext_part = os.path.splitext(out_name)
+            new_name = f"{name_part}_{counter}{ext_part}"
+            new_output_path = os.path.join(output_dir, new_name)
+            counter += 1
+        return new_output_path
+
+    def _generate_single_video(
+        self,
+        prompt: str,
+        sampling_param: SamplingParam | None = None,
+        fastvideo_args: FastVideoArgs | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Internal method for single video generation"""
+        if fastvideo_args is None:
+            fastvideo_args = self.fastvideo_args
+
+        # Validate inputs
+        if not isinstance(prompt, str):
+            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
+        prompt = prompt.strip()
+        sampling_param = deepcopy(sampling_param)
+        output_path = kwargs["output_path"]
+        sampling_param.prompt = prompt
+        # Process negative prompt
+        if sampling_param.negative_prompt is not None:
+            sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
+
+        # Validate dimensions
+        if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
+            raise ValueError(f"Height, width, and num_frames must be positive integers, got "
+                             f"height={sampling_param.height}, width={sampling_param.width}, "
+                             f"num_frames={sampling_param.num_frames}")
+
+        # Calculate sizes
+        target_height = align_to(sampling_param.height, 16)
+        target_width = align_to(sampling_param.width, 16)
+
+        # Calculate latent sizes
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+        # Log parameters
+        debug_str = f"""
+                      height: {target_height}
+                       width: {target_width}
+                video_length: {sampling_param.num_frames}
+                      prompt: {sampling_param.prompt}
+                      image_path: {sampling_param.image_path}
+                  neg_prompt: {sampling_param.negative_prompt}
+                        seed: {sampling_param.seed}
+                 infer_steps: {sampling_param.num_inference_steps}
+       num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
+              guidance_scale: {sampling_param.guidance_scale}
+                    n_tokens: {n_tokens}
+                  flow_shift: {fastvideo_args.pipeline_config.flow_shift}
+     embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale}
+                  save_video: {sampling_param.save_video}
+                  output_path: {output_path}
+        """ # type: ignore[attr-defined]
+        logger.info(debug_str)
+
+        # Prepare batch
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            eta=0.0,
+            n_tokens=n_tokens,
+            VSA_sparsity=fastvideo_args.VSA_sparsity,
+        )
+
+        # Run inference
+        start_time = time.perf_counter()
+
+        # Execute forward pass in a new thread for non-blocking tensor
+        # allocation. Capture thread exceptions so we can surface the true
+        # failure in the main thread instead of later hitting None outputs.
+        result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)}
+        thread_error: dict[str, BaseException | None] = {"error": None}
+        thread_error_traceback: dict[str, str] = {"traceback": ""}
+
+        def execute_forward_thread():
+            import traceback
+            try:
+                result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args)
+            except BaseException as error:  # noqa: BLE001
+                thread_error["error"] = error
+                thread_error_traceback["traceback"] = traceback.format_exc()
+
+        thread = threading.Thread(target=execute_forward_thread)
+        thread.start()
+        latent_batch_size = _infer_latent_batch_size(batch)
+        samples = torch.empty(
+            (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
+            device='cpu',
+            pin_memory=fastvideo_args.pin_cpu_memory)
+        thread.join()
+
+        if thread_error["error"] is not None:
+            raise RuntimeError("Forward execution thread failed.\n"
+                               f"{thread_error_traceback['traceback']}") from thread_error["error"]
+
+        output_batch = result_container["output_batch"]
+        if output_batch.output is None:
+            raise RuntimeError("Forward execution returned no output tensor. "
+                               "This usually means the executor/pipeline failed earlier.")
+
+        if output_batch.output.shape == samples.shape:
+            samples.copy_(output_batch.output)
+        else:
+            logger.warning("Output shape %s does not match expected shape %s; use slow path", output_batch.output.shape,
+                           samples.shape)
+            samples = output_batch.output.cpu()
+        logging_info = output_batch.logging_info
+
+        gen_time = time.perf_counter() - start_time
+        logger.info("Generated successfully in %.2f seconds", gen_time)
+
+        # Process outputs
+        videos = rearrange(samples, "b c t h w -> t b c h w")
+        frames = []
+        for x in videos:
+            x = torchvision.utils.make_grid(x, nrow=6)
+            x = x.permute(1, 2, 0).squeeze(-1)
+            x = (x * 255).to(torch.uint8)
+            frames.append(x.cpu().numpy())
+
+        # Save output if requested
+        if batch.save_video:
+            if self._is_image_workload():
+                # Image workloads (t2i, i2i, …): save the first frame as PNG.
+                imageio.imwrite(output_path, frames[0])
+                logger.info("Saved image to %s", output_path)
+            else:
+                imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
+                logger.info("Saved video to %s", output_path)
+                audio = output_batch.extra.get("audio")
+                audio_sample_rate = output_batch.extra.get("audio_sample_rate")
+                if (audio is not None and audio_sample_rate is not None
+                        and not self._mux_audio(output_path, audio, audio_sample_rate)):
+                    logger.warning("Audio mux failed; saved video without audio.")
+
+        result: dict[str, Any] = {
+            "prompts": prompt,
+            "samples": samples if batch.return_frames else None,
+            "frames": frames if batch.return_frames else None,
+            "audio": output_batch.extra.get("audio") if batch.return_frames else None,
+            "size": (target_height, target_width, batch.num_frames),
+            "generation_time": gen_time,
+            "logging_info": logging_info,
+            "trajectory": output_batch.trajectory_latents,
+            "trajectory_timesteps": output_batch.trajectory_timesteps,
+            "trajectory_decoded": output_batch.trajectory_decoded,
+            "video_path": output_path if batch.save_video else None,
+            "peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
+        }
+
+        return result
+
+    @staticmethod
+    def _wrap_legacy_result(
+        result: dict[str, Any] | list[dict[str, Any]], ) -> GenerationResult | list[GenerationResult]:
+        if isinstance(result, list):
+            return [GenerationResult.from_legacy_result(item) for item in result]
+        return GenerationResult.from_legacy_result(result)
+
+    @staticmethod
+    def _unwrap_typed_result(
+        result: GenerationResult | list[GenerationResult], ) -> dict[str, Any] | list[dict[str, Any]]:
+        if isinstance(result, list):
+            return [item.to_legacy_dict() for item in result]
+        return result.to_legacy_dict()
+
+    @staticmethod
+    def _mux_audio(
+        video_path: str,
+        audio: torch.Tensor | np.ndarray,
+        sample_rate: int,
+    ) -> bool:
+        """Mux audio into video using PyAV."""
+        try:
+            import av
+        except ImportError:
+            logger.warning("PyAV not installed; cannot mux audio. "
+                           "Install with: pip install av")
+            return False
+
+        if torch.is_tensor(audio):
+            audio_np = audio.detach().cpu().float().numpy()
+        else:
+            audio_np = np.asarray(audio, dtype=np.float32)
+
+        if audio_np.ndim == 1:
+            audio_np = audio_np[:, None]
+        elif audio_np.ndim == 2:
+            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
+                audio_np = audio_np.T
+        else:
+            logger.warning("Unexpected audio shape %s; skipping mux.", audio_np.shape)
+            return False
+
+        audio_np = np.clip(audio_np, -1.0, 1.0)
+        audio_int16 = (audio_np * 32767.0).astype(np.int16)
+        num_channels = audio_int16.shape[1]
+        layout = "stereo" if num_channels == 2 else "mono"
+
+        try:
+            import wave
+            with tempfile.TemporaryDirectory() as tmpdir:
+                out_path = os.path.join(tmpdir, "muxed.mp4")
+                wav_path = os.path.join(tmpdir, "audio.wav")
+
+                # Write audio to WAV file
+                with wave.open(wav_path, "wb") as wav_file:
+                    wav_file.setnchannels(num_channels)
+                    wav_file.setsampwidth(2)
+                    wav_file.setframerate(sample_rate)
+                    wav_file.writeframes(audio_int16.tobytes())
+
+                # Open input video and audio
+                input_video = av.open(video_path)
+                input_audio = av.open(wav_path)
+
+                # Create output with both streams
+                output = av.open(out_path, mode="w")
+
+                # Add video stream (copy codec from input)
+                in_video_stream = input_video.streams.video[0]
+                out_video_stream = output.add_stream(
+                    codec_name=in_video_stream.codec_context.name,
+                    rate=in_video_stream.average_rate,
+                )
+                out_video_stream.width = in_video_stream.width
+                out_video_stream.height = in_video_stream.height
+                out_video_stream.pix_fmt = in_video_stream.pix_fmt
+
+                # Add audio stream (AAC)
+                out_audio_stream = output.add_stream("aac", rate=sample_rate)
+                out_audio_stream.layout = layout
+
+                # Remux video (decode and re-encode to be safe)
+                for frame in input_video.decode(video=0):
+                    for packet in out_video_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_video_stream.encode():
+                    output.mux(packet)
+
+                # Encode audio
+                for frame in input_audio.decode(audio=0):
+                    frame.pts = None  # Let encoder assign PTS
+                    for packet in out_audio_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_audio_stream.encode():
+                    output.mux(packet)
+
+                input_video.close()
+                input_audio.close()
+                output.close()
+                shutil.move(out_path, video_path)
+            return True
+        except Exception as e:
+            logger.warning("Audio mux failed: %s", e)
+            return False
+
+    def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None) -> None:
+        self.executor.set_lora_adapter(lora_nickname, lora_path)
+
+    def unmerge_lora_weights(self) -> None:
+        """
+        Use unmerged weights for inference to produce videos that align with 
+        validation videos generated during training.
+        """
+        self.executor.unmerge_lora_weights()
+
+    def merge_lora_weights(self) -> None:
+        self.executor.merge_lora_weights()
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the video generator.
+        """
+        self.executor.shutdown()
+        del self.executor

backend_snapshot/fastvideo/fastvideo_args.py

@@ -0,0 +1,1188 @@
+# SPDX-License-Identifier: Apache-2.0
+# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
+"""The arguments of FastVideo Inference."""
+import argparse
+import dataclasses
+import json
+from contextlib import contextmanager
+from dataclasses import field
+from enum import Enum
+from typing import Any, TYPE_CHECKING
+
+from fastvideo.configs.configs import PreprocessConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.utils import clean_cli_args
+from fastvideo.layers.quantization import QUANTIZATION_METHODS, QuantizationMethods
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser, StoreBoolean
+
+if TYPE_CHECKING:
+    from ray.runtime_env import RuntimeEnv
+    from ray.util.placement_group import PlacementGroup
+else:
+    RuntimeEnv = Any
+    PlacementGroup = Any
+
+logger = init_logger(__name__)
+
+
+class ExecutionMode(str, Enum):
+    """
+    Enumeration for different pipeline modes.
+    
+    Inherits from str to allow string comparison for backward compatibility.
+    """
+    INFERENCE = "inference"
+    PREPROCESS = "preprocess"
+    FINETUNING = "finetuning"
+    DISTILLATION = "distillation"
+
+    @classmethod
+    def from_string(cls, value: str) -> "ExecutionMode":
+        """Convert string to ExecutionMode enum."""
+        try:
+            return cls(value.lower())
+        except ValueError:
+            raise ValueError(f"Invalid mode: {value}. Must be one of: {', '.join([m.value for m in cls])}") from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        """Get all available choices as strings for argparse."""
+        return [mode.value for mode in cls]
+
+
+class WorkloadType(str, Enum):
+    """
+    Enumeration for different workload types.
+    
+    Inherits from str to allow string comparison for backward compatibility.
+    """
+    I2V = "i2v"  # Image to Video
+    T2V = "t2v"  # Text to Video
+    T2I = "t2i"  # Text to Image
+    I2I = "i2i"  # Image to Image
+
+    @classmethod
+    def from_string(cls, value: str) -> "WorkloadType":
+        """Convert string to WorkloadType enum."""
+        try:
+            return cls(value.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid workload type: {value}. Must be one of: {', '.join([m.value for m in cls])}") from None
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        """Get all available choices as strings for argparse."""
+        return [workload.value for workload in cls]
+
+
+# args for fastvideo framework
+@dataclasses.dataclass
+class FastVideoArgs:
+    # Model and path configuration (for convenience)
+    model_path: str
+
+    # Running mode
+    mode: ExecutionMode = ExecutionMode.INFERENCE
+
+    # Workload type
+    workload_type: WorkloadType = WorkloadType.T2V
+
+    # Distributed executor backend
+    distributed_executor_backend: str = "mp"
+
+    # a few attributes for ray related
+    ray_placement_group: PlacementGroup | None = None
+    ray_runtime_env: RuntimeEnv | None = None
+
+    inference_mode: bool = True  # if False == training mode
+
+    # HuggingFace specific parameters
+    trust_remote_code: bool = False
+    revision: str | None = None
+
+    # Parallelism
+    num_gpus: int = 1
+    tp_size: int = -1
+    sp_size: int = -1
+    hsdp_replicate_dim: int = 1
+    hsdp_shard_dim: int = -1
+    dist_timeout: int | None = None  # timeout for torch.distributed
+
+    pipeline_config: PipelineConfig = field(default_factory=PipelineConfig)
+    preprocess_config: PreprocessConfig | None = None
+
+    # LoRA parameters
+    # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
+    lora_path: str | None = None
+    lora_nickname: str = "default"  # for swapping adapters in the pipeline
+    # can restrict layers to adapt, e.g. ["q_proj"]
+    # Will adapt only q, k, v, o by default.
+    lora_target_modules: list[str] | None = None
+
+    output_type: str = "pil"
+
+    # CPU offload parameters
+    dit_cpu_offload: bool = True
+    use_fsdp_inference: bool = False
+    dit_layerwise_offload: bool = True
+    text_encoder_cpu_offload: bool = True
+    image_encoder_cpu_offload: bool = True
+    vae_cpu_offload: bool = True
+    pin_cpu_memory: bool = True
+
+    # Compilation
+    enable_torch_compile: bool = False
+    torch_compile_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    disable_autocast: bool = False
+
+    # VSA parameters
+    VSA_sparsity: float = 0.0  # inference/validation sparsity
+
+    # V-MoBA parameters
+    moba_config_path: str | None = None
+    moba_config: dict[str, Any] = field(default_factory=dict)
+
+    # Master port for distributed training/inference
+    master_port: int | None = None
+
+    # Stage verification
+    enable_stage_verification: bool = True
+
+    # Prompt text file for batch processing
+    prompt_txt: str | None = None
+
+    # LTX-2 VAE tiling overrides
+    ltx2_vae_tiling: bool | None = None
+    ltx2_vae_spatial_tile_size_in_pixels: int | None = None
+    ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None
+    ltx2_vae_temporal_tile_size_in_frames: int | None = None
+    ltx2_vae_temporal_tile_overlap_in_frames: int | None = None
+    ltx2_initial_latent_path: str | None = None
+
+    # model paths for correct deallocation
+    model_paths: dict[str, str] = field(default_factory=dict)
+    model_loaded: dict[str, bool] = field(default_factory=lambda: {
+        "transformer": True,
+        "vae": True,
+        "upsampler": True,
+    })
+
+    override_text_encoder_safetensors: str | None = None  # path to safetensors file for text encoder override
+    override_text_encoder_quant: QuantizationMethods = None
+    transformer_quant: QuantizationMethods = None
+
+    override_transformer_cls_name: str | None = None
+    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
+    init_weights_from_safetensors_2: str = ""  # path to safetensors file for initial weight loading for transformer_2
+
+    override_pipeline_cls_name: str | None = None
+
+    # # DMD parameters
+    # dmd_denoising_steps: List[int] | None = field(default=None)
+
+    # MoE parameters used by Wan2.2
+    boundary_ratio: float = 0.875
+
+    @property
+    def training_mode(self) -> bool:
+        return not self.inference_mode
+
+    def __post_init__(self):
+        if self.moba_config_path:
+            try:
+                with open(self.moba_config_path) as f:
+                    self.moba_config = json.load(f)
+                logger.info("Loaded V-MoBA config from %s", self.moba_config_path)
+            except (FileNotFoundError, json.JSONDecodeError) as e:
+                logger.error("Failed to load V-MoBA config from %s: %s", self.moba_config_path, e)
+                raise
+        self._apply_ltx2_vae_overrides()
+        self.check_fastvideo_args()
+
+    def __getattr__(self, name: str) -> Any:
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def _apply_ltx2_vae_overrides(self) -> None:
+        if self.pipeline_config is None:
+            return
+        vae_config = self.pipeline_config.vae_config
+        has_any = any(value is not None for value in (
+            self.ltx2_vae_spatial_tile_size_in_pixels,
+            self.ltx2_vae_spatial_tile_overlap_in_pixels,
+            self.ltx2_vae_temporal_tile_size_in_frames,
+            self.ltx2_vae_temporal_tile_overlap_in_frames,
+        ))
+        if self.ltx2_vae_tiling is not None and hasattr(self.pipeline_config, "vae_tiling"):
+            self.pipeline_config.vae_tiling = self.ltx2_vae_tiling
+        elif has_any and hasattr(self.pipeline_config, "vae_tiling"):
+            self.pipeline_config.vae_tiling = True
+
+        if hasattr(vae_config,
+                   "ltx2_spatial_tile_size_in_pixels") and self.ltx2_vae_spatial_tile_size_in_pixels is not None:
+            vae_config.ltx2_spatial_tile_size_in_pixels = (self.ltx2_vae_spatial_tile_size_in_pixels)
+        if hasattr(vae_config,
+                   "ltx2_spatial_tile_overlap_in_pixels") and self.ltx2_vae_spatial_tile_overlap_in_pixels is not None:
+            vae_config.ltx2_spatial_tile_overlap_in_pixels = (self.ltx2_vae_spatial_tile_overlap_in_pixels)
+        if hasattr(vae_config,
+                   "ltx2_temporal_tile_size_in_frames") and self.ltx2_vae_temporal_tile_size_in_frames is not None:
+            vae_config.ltx2_temporal_tile_size_in_frames = (self.ltx2_vae_temporal_tile_size_in_frames)
+        if hasattr(
+                vae_config,
+                "ltx2_temporal_tile_overlap_in_frames") and self.ltx2_vae_temporal_tile_overlap_in_frames is not None:
+            vae_config.ltx2_temporal_tile_overlap_in_frames = (self.ltx2_vae_temporal_tile_overlap_in_frames)
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+        # Model and path configuration
+        parser.add_argument(
+            "--model-path",
+            type=str,
+            help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
+        )
+
+        # Running mode
+        parser.add_argument(
+            "--mode",
+            type=str,
+            choices=ExecutionMode.choices(),
+            default=FastVideoArgs.mode.value,
+            help="The mode to run FastVideo",
+        )
+
+        # Workload type
+        parser.add_argument(
+            "--workload-type",
+            type=str,
+            choices=WorkloadType.choices(),
+            default=FastVideoArgs.workload_type.value,
+            help="The workload type",
+        )
+
+        # distributed_executor_backend
+        parser.add_argument(
+            "--distributed-executor-backend",
+            type=str,
+            choices=["mp"],
+            default=FastVideoArgs.distributed_executor_backend,
+            help="The distributed executor backend to use",
+        )
+
+        parser.add_argument(
+            "--inference-mode",
+            action=StoreBoolean,
+            default=FastVideoArgs.inference_mode,
+            help="Whether to use inference mode",
+        )
+
+        # HuggingFace specific parameters
+        parser.add_argument(
+            "--trust-remote-code",
+            action=StoreBoolean,
+            default=FastVideoArgs.trust_remote_code,
+            help="Trust remote code when loading HuggingFace models",
+        )
+        parser.add_argument(
+            "--revision",
+            type=str,
+            default=FastVideoArgs.revision,
+            help="The specific model version to use (can be a branch name, tag name, or commit id)",
+        )
+
+        # Parallelism
+        parser.add_argument(
+            "--num-gpus",
+            type=int,
+            default=FastVideoArgs.num_gpus,
+            help="The number of GPUs to use.",
+        )
+        parser.add_argument(
+            "--tp-size",
+            type=int,
+            default=FastVideoArgs.tp_size,
+            help="The tensor parallelism size.",
+        )
+        parser.add_argument(
+            "--sp-size",
+            type=int,
+            default=FastVideoArgs.sp_size,
+            help="The sequence parallelism size.",
+        )
+        parser.add_argument(
+            "--hsdp-replicate-dim",
+            type=int,
+            default=FastVideoArgs.hsdp_replicate_dim,
+            help="The data parallelism size.",
+        )
+        parser.add_argument(
+            "--hsdp-shard-dim",
+            type=int,
+            default=FastVideoArgs.hsdp_shard_dim,
+            help="The data parallelism shards.",
+        )
+        parser.add_argument(
+            "--dist-timeout",
+            type=int,
+            default=FastVideoArgs.dist_timeout,
+            help="Set timeout for torch.distributed initialization.",
+        )
+
+        # Output type
+        parser.add_argument(
+            "--output-type",
+            type=str,
+            default=FastVideoArgs.output_type,
+            choices=["pil"],
+            help="Output type for the generated video",
+        )
+
+        # Prompt text file for batch processing
+        parser.add_argument(
+            "--prompt-txt",
+            type=str,
+            default=FastVideoArgs.prompt_txt,
+            help="Path to a text file containing prompts (one per line) for batch processing",
+        )
+
+        # LTX-2 VAE tiling overrides
+        parser.add_argument(
+            "--ltx2-vae-tiling",
+            action=StoreBoolean,
+            default=FastVideoArgs.ltx2_vae_tiling,
+            help="Enable LTX-2 VAE tiling overrides.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-spatial-tile-size-in-pixels",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_spatial_tile_size_in_pixels,
+            help="LTX-2 VAE spatial tile size in pixels.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-spatial-tile-overlap-in-pixels",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_spatial_tile_overlap_in_pixels,
+            help="LTX-2 VAE spatial tile overlap in pixels.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-temporal-tile-size-in-frames",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_temporal_tile_size_in_frames,
+            help="LTX-2 VAE temporal tile size in frames.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-temporal-tile-overlap-in-frames",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_temporal_tile_overlap_in_frames,
+            help="LTX-2 VAE temporal tile overlap in frames.",
+        )
+        parser.add_argument(
+            "--ltx2-initial-latent-path",
+            type=str,
+            default=FastVideoArgs.ltx2_initial_latent_path,
+            help="Path to load/save a precomputed LTX-2 initial latent.",
+        )
+
+        # LoRA parameters (inference-time adapter loading)
+        parser.add_argument(
+            "--lora-path",
+            type=str,
+            default=FastVideoArgs.lora_path,
+            help="Path to a LoRA adapter (directory or HF repo id). If set, LoRA will be applied at inference.",
+        )
+        parser.add_argument(
+            "--lora-nickname",
+            type=str,
+            default=FastVideoArgs.lora_nickname,
+            help="Nickname to refer to the loaded LoRA adapter (useful for swapping).",
+        )
+        parser.add_argument(
+            "--lora-target-modules",
+            nargs="+",
+            type=str,
+            default=FastVideoArgs.lora_target_modules,
+            help="Optional list of module name substrings to restrict LoRA injection (e.g. q_proj k_proj v_proj).",
+        )
+
+        # BSA runtime control (LongCat)
+        parser.add_argument(
+            "--enable-bsa",
+            action=StoreBoolean,
+            help="Enable Block Sparse Attention (BSA) at runtime (overrides config).",
+        )
+        parser.add_argument(
+            "--bsa-sparsity",
+            type=float,
+            help="BSA sparsity (e.g., 0.9375).",
+        )
+        parser.add_argument(
+            "--bsa-cdf-threshold",
+            type=float,
+            help="BSA CDF threshold (optional).",
+        )
+        parser.add_argument(
+            "--bsa-chunk-q",
+            nargs=3,
+            type=int,
+            metavar=("T", "H", "W"),
+            help="BSA chunk_3d_shape_q as three ints, e.g., 4 4 4.",
+        )
+        parser.add_argument(
+            "--bsa-chunk-k",
+            nargs=3,
+            type=int,
+            metavar=("T", "H", "W"),
+            help="BSA chunk_3d_shape_k as three ints, e.g., 4 4 4.",
+        )
+
+        parser.add_argument(
+            "--enable-torch-compile",
+            action=StoreBoolean,
+            default=FastVideoArgs.enable_torch_compile,
+            help="Use torch.compile to speed up DiT inference." +
+            "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
+        )
+        parser.add_argument(
+            "--torch-compile-kwargs",
+            type=str,
+            default=None,
+            help=
+            "JSON string of kwargs to pass to torch.compile. Example: '{\"backend\":\"inductor\",\"mode\":\"reduce-overhead\"}'",
+        )
+
+        parser.add_argument(
+            "--dit-cpu-offload",
+            action=StoreBoolean,
+            help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
+        )
+        parser.add_argument(
+            "--dit-layerwise-offload",
+            action=StoreBoolean,
+            help="Enable layerwise CPU offload with async H2D prefetch overlap.",
+        )
+        parser.add_argument(
+            "--use-fsdp-inference",
+            action=StoreBoolean,
+            help=
+            "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--text-encoder-cpu-offload",
+            action=StoreBoolean,
+            help="Use CPU offload for text encoder. Enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--image-encoder-cpu-offload",
+            action=StoreBoolean,
+            help="Use CPU offload for image encoder. Enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--vae-cpu-offload",
+            action=StoreBoolean,
+            help="Use CPU offload for VAE. Enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--pin-cpu-memory",
+            action=StoreBoolean,
+            help=
+            "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
+            "Should be enabled in almost all cases",
+        )
+        parser.add_argument(
+            "--disable-autocast",
+            action=StoreBoolean,
+            help="Disable autocast for denoising loop and vae decoding in pipeline sampling",
+        )
+
+        # VSA parameters
+        parser.add_argument(
+            "--VSA-sparsity",
+            type=float,
+            default=FastVideoArgs.VSA_sparsity,
+            help="Validation sparsity for VSA",
+        )
+
+        # Master port for distributed training/inference
+        parser.add_argument(
+            "--master-port",
+            type=int,
+            default=FastVideoArgs.master_port,
+            help="Master port for distributed training/inference",
+        )
+
+        # Stage verification
+        parser.add_argument(
+            "--enable-stage-verification",
+            action=StoreBoolean,
+            default=FastVideoArgs.enable_stage_verification,
+            help="Enable input/output verification for pipeline stages",
+        )
+        parser.add_argument(
+            "--override-text-encoder-safetensors",
+            type=str,
+            default=FastVideoArgs.override_text_encoder_safetensors,
+            help="Path to safetensors file for text encoder override",
+        )
+        parser.add_argument(
+            "--override-text-encoder-quant",
+            type=str,
+            choices=QUANTIZATION_METHODS,
+            default=FastVideoArgs.override_text_encoder_quant,
+            help="Quantization method for text encoder override",
+        )
+        parser.add_argument(
+            "--transformer-quant",
+            type=str,
+            choices=QUANTIZATION_METHODS,
+            default=FastVideoArgs.transformer_quant,
+            help="Quantization method for transformer loading",
+        )
+        parser.add_argument(
+            "--override-transformer-cls-name",
+            type=str,
+            default=FastVideoArgs.override_transformer_cls_name,
+            help="Override transformer cls name",
+        )
+        parser.add_argument(
+            "--override-pipeline-cls-name",
+            type=str,
+            default=FastVideoArgs.override_pipeline_cls_name,
+            help="Override pipeline cls name",
+        )
+        parser.add_argument("--init-weights-from-safetensors",
+                            type=str,
+                            help="Path to safetensors file for initial weight loading")
+        parser.add_argument("--init-weights-from-safetensors-2",
+                            type=str,
+                            help="Path to safetensors file for initial weight loading")
+
+        # Add pipeline configuration arguments
+        PipelineConfig.add_cli_args(parser)
+
+        # Add preprocessing configuration arguments
+        PreprocessConfig.add_cli_args(parser)
+
+        return parser
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
+        provided_args = clean_cli_args(args)
+        # Get all fields from the dataclass
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+
+        # Create a dictionary of attribute values, with defaults for missing attributes
+        kwargs: dict[str, Any] = {}
+        for attr in attrs:
+            if attr == 'pipeline_config':
+                pipeline_config = PipelineConfig.from_kwargs(provided_args)
+                kwargs['pipeline_config'] = pipeline_config
+            elif attr == 'preprocess_config':
+                preprocess_config = PreprocessConfig.from_kwargs(provided_args)
+                kwargs['preprocess_config'] = preprocess_config
+            elif attr == 'mode':
+                # Convert string to ExecutionMode enum
+                mode_value = getattr(args, attr, FastVideoArgs.mode.value)
+                kwargs['mode'] = ExecutionMode.from_string(mode_value) if isinstance(mode_value, str) else mode_value
+            elif attr == 'torch_compile_kwargs':
+                # Parse JSON string for torch.compile kwargs
+                torch_compile_kwargs_str = getattr(args, 'torch_compile_kwargs', None)
+                if torch_compile_kwargs_str:
+                    try:
+                        import json
+                        kwargs['torch_compile_kwargs'] = json.loads(torch_compile_kwargs_str)
+                    except json.JSONDecodeError as e:
+                        raise ValueError(f"Invalid JSON for torch_compile_kwargs: {e}") from e
+                else:
+                    kwargs['torch_compile_kwargs'] = {}
+            elif attr == 'workload_type':
+                # Convert string to WorkloadType enum
+                workload_type_value = getattr(args, 'workload_type', FastVideoArgs.workload_type.value)
+                kwargs['workload_type'] = WorkloadType.from_string(workload_type_value) if isinstance(
+                    workload_type_value, str) else workload_type_value
+            # Use getattr with default value from the dataclass for potentially missing attributes
+            else:
+                # Get the field to check if it has a default_factory
+                field = dataclasses.fields(cls)[next(i for i, f in enumerate(dataclasses.fields(cls))
+                                                     if f.name == attr)]
+                if field.default_factory is not dataclasses.MISSING:
+                    # Use the default_factory to create the default value
+                    default_value = field.default_factory()
+                else:
+                    default_value = getattr(cls, attr, None)
+                value = getattr(args, attr, default_value)
+                kwargs[attr] = value  # type: ignore
+
+        return cls(**kwargs)  # type: ignore
+
+    @classmethod
+    def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs":
+        # Convert mode string to enum if necessary
+        if 'mode' in kwargs and isinstance(kwargs['mode'], str):
+            kwargs['mode'] = ExecutionMode.from_string(kwargs['mode'])
+
+        # Convert workload_type string to enum if necessary
+        if 'workload_type' in kwargs and isinstance(kwargs['workload_type'], str):
+            kwargs['workload_type'] = WorkloadType.from_string(kwargs['workload_type'])
+
+        kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
+        kwargs['preprocess_config'] = PreprocessConfig.from_kwargs(kwargs)
+        # Filter to only FastVideoArgs dataclass fields — pipeline-specific CLI
+        # args (e.g. enable_bsa, bsa_sparsity) live in PipelineConfig and must
+        # not be forwarded to the FastVideoArgs constructor.
+        valid_fields = {f.name for f in dataclasses.fields(cls)}
+        return cls(**{k: v for k, v in kwargs.items() if k in valid_fields})
+
+    def check_fastvideo_args(self) -> None:
+        """Validate inference arguments for consistency"""
+        from fastvideo.platforms import current_platform
+
+        if current_platform.is_mps():
+            self.use_fsdp_inference = False
+            self.dit_layerwise_offload = False
+
+        if self.dit_layerwise_offload:
+            if self.use_fsdp_inference:
+                logger.warning("dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference.")
+                self.use_fsdp_inference = False
+            if self.dit_cpu_offload:
+                logger.warning("dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload.")
+                self.dit_cpu_offload = False
+
+        # Validate mode and inference_mode consistency
+        assert isinstance(self.mode, ExecutionMode), f"Mode must be an ExecutionMode enum, got {type(self.mode)}"
+        assert self.mode in ExecutionMode.choices(), f"Invalid execution mode: {self.mode}"
+
+        # Validate workload type
+        assert isinstance(self.workload_type,
+                          WorkloadType), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}"
+        assert self.workload_type in WorkloadType.choices(), f"Invalid workload type: {self.workload_type}"
+
+        if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING] and self.inference_mode:
+            logger.warning("Mode is 'training' but inference_mode is True. Setting inference_mode to False.")
+            self.inference_mode = False
+        elif self.mode in [ExecutionMode.INFERENCE, ExecutionMode.PREPROCESS] and not self.inference_mode:
+            logger.warning("Mode is '%s' but inference_mode is False. Setting inference_mode to True.", self.mode)
+            self.inference_mode = True
+
+        if not self.inference_mode:
+            assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
+            assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
+            assert self.sp_size != -1, "sp_size must be set for training"
+
+        if self.tp_size == -1:
+            self.tp_size = 1
+        if self.sp_size == -1:
+            self.sp_size = self.num_gpus
+        if self.hsdp_shard_dim == -1:
+            self.hsdp_shard_dim = self.num_gpus
+
+        assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size"
+        assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim"
+        assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim"
+
+        if self.num_gpus < max(self.tp_size, self.sp_size):
+            self.num_gpus = max(self.tp_size, self.sp_size)
+
+        if self.pipeline_config is None:
+            raise ValueError("pipeline_config is not set in FastVideoArgs")
+
+        self.pipeline_config.check_pipeline_config()
+
+        # Add preprocessing config validation if needed
+        if self.mode == ExecutionMode.PREPROCESS:
+            if self.preprocess_config is None:
+                raise ValueError("preprocess_config is not set in FastVideoArgs when mode is PREPROCESS")
+            if self.preprocess_config.model_path == "":
+                self.preprocess_config.model_path = self.model_path
+            if not self.pipeline_config.vae_config.load_encoder:
+                self.pipeline_config.vae_config.load_encoder = True
+            self.preprocess_config.check_preprocess_config()
+
+
+_current_fastvideo_args = None
+
+
+def prepare_fastvideo_args(argv: list[str]) -> FastVideoArgs:
+    """
+    Prepare the inference arguments from the command line arguments.
+
+    Args:
+        argv: The command line arguments. Typically, it should be `sys.argv[1:]`
+            to ensure compatibility with `parse_args` when no arguments are passed.
+
+    Returns:
+        The inference arguments.
+    """
+    parser = FlexibleArgumentParser()
+    FastVideoArgs.add_cli_args(parser)
+    raw_args = parser.parse_args(argv)
+    fastvideo_args = FastVideoArgs.from_cli_args(raw_args)
+    global _current_fastvideo_args
+    _current_fastvideo_args = fastvideo_args
+    return fastvideo_args
+
+
+@contextmanager
+def set_current_fastvideo_args(fastvideo_args: FastVideoArgs):
+    """
+    Temporarily set the current fastvideo config.
+    Used during model initialization.
+    We save the current fastvideo config in a global variable,
+    so that all modules can access it, e.g. custom ops
+    can access the fastvideo config to determine how to dispatch.
+    """
+    global _current_fastvideo_args
+    old_fastvideo_args = _current_fastvideo_args
+    try:
+        _current_fastvideo_args = fastvideo_args
+        yield
+    finally:
+        _current_fastvideo_args = old_fastvideo_args
+
+
+def get_current_fastvideo_args() -> FastVideoArgs:
+    if _current_fastvideo_args is None:
+        # in ci, usually when we test custom ops/modules directly,
+        # we don't set the fastvideo config. In that case, we set a default
+        # config.
+        # TODO(will): may need to handle this for CI.
+        raise ValueError("Current fastvideo args is not set.")
+    return _current_fastvideo_args
+
+
+@dataclasses.dataclass
+class TrainingArgs(FastVideoArgs):
+    """
+    Training arguments. Inherits from FastVideoArgs and adds training-specific
+    arguments. If there are any conflicts, the training arguments will take
+    precedence.
+    """
+    data_path: str = ""
+    dataloader_num_workers: int = 0
+    num_height: int = 0
+    num_width: int = 0
+    num_frames: int = 0
+
+    train_batch_size: int = 0
+    num_latent_t: int = 0
+    group_frame: bool = False
+    group_resolution: bool = False
+
+    # text encoder & vae & diffusion model
+    pretrained_model_name_or_path: str = ""
+
+    # DMD model paths - separate paths for each network
+    real_score_model_path: str = ""  # path for real score (teacher) model
+    fake_score_model_path: str = ""  # path for fake score (critic) model
+
+    # diffusion setting
+    ema_decay: float = 0.0
+    ema_start_step: int = 0
+    training_cfg_rate: float = 0.0
+    precondition_outputs: bool = False
+
+    # validation & logs
+    validation_dataset_file: str = ""
+    validation_preprocessed_path: str = ""
+    validation_sampling_steps: str = ""
+    validation_guidance_scale: str = ""
+    validation_steps: float = 0.0
+    log_validation: bool = False
+    trackers: list[str] = dataclasses.field(default_factory=list)
+    tracker_project_name: str = ""
+    wandb_run_name: str = ""
+    seed: int = 0
+    _loading_teacher_critic_model: bool = False
+
+    # output
+    output_dir: str = ""
+    checkpoints_total_limit: int = 0
+    resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
+
+    # optimizer & scheduler
+    num_train_epochs: int = 0
+    max_train_steps: int = 0
+    gradient_accumulation_steps: int = 0
+    learning_rate: float = 0.0
+    scale_lr: bool = False
+    lr_scheduler: str = "constant"
+    lr_warmup_steps: int = 0
+    max_grad_norm: float = 0.0
+    enable_gradient_checkpointing_type: str | None = None
+    selective_checkpointing: float = 0.0
+    mixed_precision: str = ""
+    train_sp_batch_size: int = 0
+    fsdp_sharding_startegy: str = ""
+
+    weighting_scheme: str = ""
+    logit_mean: float = 0.0
+    logit_std: float = 1.0
+    mode_scale: float = 0.0
+
+    num_euler_timesteps: int = 0
+    lr_num_cycles: int = 0
+    lr_power: float = 0.0
+    min_lr_ratio: float = 0.5  # minimum learning rate ratio for cosine_with_min_lr scheduler
+    not_apply_cfg_solver: bool = False
+    distill_cfg: float = 0.0
+    scheduler_type: str = ""
+    linear_quadratic_threshold: float = 0.0
+    linear_range: float = 0.0
+    weight_decay: float = 0.0
+    betas: str = "0.9,0.999"  # betas for optimizer, format: "beta1,beta2"
+    use_ema: bool = False
+    multi_phased_distill_schedule: str = ""
+    pred_decay_weight: float = 0.0
+    pred_decay_type: str = ""
+    hunyuan_teacher_disable_cfg: bool = False
+
+    # master_weight_type
+    master_weight_type: str = ""
+
+    # VSA training decay parameters
+    VSA_decay_rate: float = 0.01  # decay rate -> 0.02
+    VSA_decay_interval_steps: int = 1  # decay interval steps -> 50
+    VSA_init_sparsity: float = 0.0  # initial sparsity (default 0, ramp from 0)
+    VSA_warmup_steps: int = 0  # keep init_sparsity for this many steps before ramping
+
+    # LoRA training parameters
+    lora_rank: int | None = None
+    lora_alpha: int | None = None
+    lora_training: bool = False
+    ltx2_first_frame_conditioning_p: float = 0.1
+
+    # distillation args
+    generator_update_interval: int = 5
+    dfake_gen_update_ratio: int = 5  # self-forcing: how often to train generator vs critic
+    min_timestep_ratio: float = 0.2
+    max_timestep_ratio: float = 0.98
+    real_score_guidance_scale: float = 3.5
+    fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
+    fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
+    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
+    training_state_checkpointing_steps: int = 0  # for resuming training
+    weight_only_checkpointing_steps: int = 0  # for inference
+    log_visualization: bool = False
+    visualization_steps: int = 0
+    # simulate generator forward to match inference
+    simulate_generator_forward: bool = False
+    warp_denoising_step: bool = False
+    generator_4bit_attn: bool = False
+    generator_4bit_linear: bool = False
+
+    # Self-forcing specific arguments
+    num_frame_per_block: int = 3
+    independent_first_frame: bool = False
+    enable_gradient_masking: bool = True
+    gradient_mask_last_n_frames: int = 21
+    same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
+    last_step_only: bool = False  # Only use the last timestep for training
+    context_noise: int = 0  # Context noise level for cache updates
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
+        provided_args = clean_cli_args(args)
+        # Get all fields from the dataclass
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        logger.info(provided_args)
+        # Create a dictionary of attribute values, with defaults for missing attributes
+        kwargs: dict[str, Any] = {}
+        for attr in attrs:
+            if attr == 'pipeline_config':
+                pipeline_config = PipelineConfig.from_kwargs(provided_args)
+                kwargs[attr] = pipeline_config
+            elif attr == 'mode':
+                # Convert string to ExecutionMode enum
+                mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value)
+                kwargs[attr] = ExecutionMode.from_string(mode_value) if isinstance(mode_value, str) else mode_value
+            elif attr == 'workload_type':
+                # Convert string to WorkloadType enum
+                workload_type_value = getattr(args, 'workload_type', WorkloadType.T2V.value)
+                kwargs[attr] = WorkloadType.from_string(workload_type_value) if isinstance(workload_type_value,
+                                                                                           str) else workload_type_value
+            # Use getattr with default value from the dataclass for potentially missing attributes
+            else:
+                # Get the field to check its default value
+                field = dataclasses.fields(cls)[next(i for i, f in enumerate(dataclasses.fields(cls))
+                                                     if f.name == attr)]
+
+                # Check if the attribute is provided in args
+                if hasattr(args, attr):
+                    value = getattr(args, attr)
+                else:
+                    # Use the field's default value
+                    if field.default_factory is not dataclasses.MISSING:
+                        value = field.default_factory()
+                    elif field.default is not dataclasses.MISSING:
+                        value = field.default
+                    else:
+                        # No default value, use None
+                        value = None
+
+                kwargs[attr] = value
+
+        return cls(**kwargs)  # type: ignore
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+        parser.add_argument("--data-path", type=str, required=True, help="Path to parquet files")
+        parser.add_argument("--dataloader-num-workers",
+                            type=int,
+                            required=True,
+                            help="Number of workers for dataloader")
+        parser.add_argument("--num-height", type=int, required=True, help="Number of heights")
+        parser.add_argument("--num-width", type=int, required=True, help="Number of widths")
+        parser.add_argument("--num-frames", type=int, required=True, help="Number of frames")
+
+        # Training batch and model configuration
+        parser.add_argument("--train-batch-size", type=int, required=True, help="Training batch size")
+        parser.add_argument("--num-latent-t", type=int, required=True, help="Number of latent time steps")
+        parser.add_argument("--group-frame", action=StoreBoolean, help="Whether to group frames during training")
+        parser.add_argument("--group-resolution",
+                            action=StoreBoolean,
+                            help="Whether to group resolutions during training")
+
+        # Model paths
+        parser.add_argument("--pretrained-model-name-or-path",
+                            type=str,
+                            required=True,
+                            help="Path to pretrained model or model name")
+        parser.add_argument("--dit-model-name-or-path",
+                            type=str,
+                            required=False,
+                            help="Path to DiT model or model name")
+        parser.add_argument("--cache-dir", type=str, help="Directory to cache models")
+
+        # DMD model paths - separate paths for each network
+        parser.add_argument("--generator-model-path",
+                            type=str,
+                            help="Path to generator (student) model for DMD distillation")
+        parser.add_argument("--real-score-model-path",
+                            type=str,
+                            help="Path to real score (teacher) model for DMD distillation")
+        parser.add_argument("--fake-score-model-path",
+                            type=str,
+                            help="Path to fake score (critic) model for DMD distillation")
+
+        # Diffusion settings
+        parser.add_argument("--ema-decay", type=float, default=0.999, help="EMA decay rate")
+        parser.add_argument("--ema-start-step", type=int, default=0, help="Step to start EMA")
+        parser.add_argument("--training-cfg-rate", type=float, help="Classifier-free guidance scale")
+        parser.add_argument("--precondition-outputs",
+                            action=StoreBoolean,
+                            help="Whether to precondition the outputs of the model")
+
+        # Validation and logging
+        parser.add_argument("--validation-dataset-file", type=str, help="Path to unprocessed validation dataset")
+        parser.add_argument("--validation-preprocessed-path", type=str, help="Path to processed validation dataset")
+        parser.add_argument("--validation-sampling-steps", type=str, help="Validation sampling steps")
+        parser.add_argument("--validation-guidance-scale", type=str, help="Validation guidance scale")
+        parser.add_argument("--validation-steps", type=float, help="Number of validation steps")
+        parser.add_argument("--log-validation", action=StoreBoolean, help="Whether to log validation results")
+        parser.add_argument("--visualization-steps", type=int, help="Number of visualization steps")
+        parser.add_argument("--tracker-project-name", type=str, help="Project name for tracking")
+        parser.add_argument("--wandb-run-name", type=str, help="Run name for wandb")
+        parser.add_argument("--seed", type=int, default=42, help="Seed for deterministic training")
+
+        # Output configuration
+        parser.add_argument("--output-dir", type=str, required=True, help="Output directory for checkpoints and logs")
+        parser.add_argument("--checkpoints-total-limit", type=int, help="Maximum number of checkpoints to keep")
+        parser.add_argument("--training-state-checkpointing-steps",
+                            type=int,
+                            help="Steps between training state checkpoints (for resuming training)")
+        parser.add_argument("--weight-only-checkpointing-steps",
+                            type=int,
+                            help="Steps between weight-only checkpoints (for inference)")
+        parser.add_argument("--resume-from-checkpoint", type=str, help="Path to checkpoint to resume from")
+        parser.add_argument("--logging-dir", type=str, help="Directory for logging")
+
+        # Training configuration
+        parser.add_argument("--num-train-epochs", type=int, help="Number of training epochs")
+        parser.add_argument("--max-train-steps", type=int, help="Maximum number of training steps")
+        parser.add_argument("--gradient-accumulation-steps", type=int, help="Number of steps to accumulate gradients")
+        parser.add_argument("--learning-rate", type=float, required=True, help="Learning rate")
+        parser.add_argument("--scale-lr", action=StoreBoolean, help="Whether to scale learning rate")
+        parser.add_argument("--lr-scheduler", type=str, default="constant", help="Learning rate scheduler type")
+        parser.add_argument("--lr-warmup-steps", type=int, default=10, help="Number of warmup steps for learning rate")
+        parser.add_argument("--max-grad-norm", type=float, help="Maximum gradient norm")
+        parser.add_argument("--enable-gradient-checkpointing-type",
+                            type=str,
+                            choices=["full", "ops", "block_skip"],
+                            default=None,
+                            help="Gradient checkpointing type")
+        parser.add_argument("--selective-checkpointing", type=float, help="Selective checkpointing threshold")
+        parser.add_argument("--mixed-precision", type=str, help="Mixed precision training type")
+        parser.add_argument("--train-sp-batch-size", type=int, help="Training spatial parallelism batch size")
+
+        parser.add_argument("--fsdp-sharding-strategy", type=str, help="FSDP sharding strategy")
+
+        parser.add_argument(
+            "--weighting_scheme",
+            type=str,
+            default="uniform",
+            choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
+        )
+        parser.add_argument(
+            "--logit_mean",
+            type=float,
+            default=0.0,
+            help="mean to use when using the `'logit_normal'` weighting scheme.",
+        )
+        parser.add_argument(
+            "--logit_std",
+            type=float,
+            default=1.0,
+            help="std to use when using the `'logit_normal'` weighting scheme.",
+        )
+        parser.add_argument(
+            "--mode_scale",
+            type=float,
+            default=1.29,
+            help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+        )
+
+        # Additional training parameters
+        parser.add_argument("--num-euler-timesteps", type=int, help="Number of Euler timesteps")
+        parser.add_argument("--lr-num-cycles", type=int, help="Number of learning rate cycles")
+        parser.add_argument("--lr-power", type=float, help="Learning rate power")
+        parser.add_argument("--min-lr-ratio",
+                            type=float,
+                            default=TrainingArgs.min_lr_ratio,
+                            help="Minimum learning rate ratio for cosine_with_min_lr scheduler")
+        parser.add_argument("--not-apply-cfg-solver", action=StoreBoolean, help="Whether to not apply CFG solver")
+        parser.add_argument("--distill-cfg", type=float, help="Distillation CFG scale")
+        parser.add_argument("--scheduler-type", type=str, help="Scheduler type")
+        parser.add_argument("--linear-quadratic-threshold", type=float, help="Linear quadratic threshold")
+        parser.add_argument("--linear-range", type=float, help="Linear range")
+        parser.add_argument("--weight-decay", type=float, help="Weight decay")
+        parser.add_argument("--betas",
+                            type=str,
+                            default=TrainingArgs.betas,
+                            help="Betas for optimizer (format: 'beta1,beta2')")
+        parser.add_argument("--use-ema", action=StoreBoolean, help="Whether to use EMA")
+        parser.add_argument("--multi-phased-distill-schedule", type=str, help="Multi-phased distillation schedule")
+        parser.add_argument("--pred-decay-weight", type=float, help="Prediction decay weight")
+        parser.add_argument("--pred-decay-type", type=str, help="Prediction decay type")
+        parser.add_argument("--hunyuan-teacher-disable-cfg",
+                            action=StoreBoolean,
+                            help="Whether to disable CFG for Hunyuan teacher")
+        parser.add_argument("--master-weight-type", type=str, help="Master weight type")
+
+        # VSA parameters for training with dense to sparse adaption
+        parser.add_argument(
+            "--VSA-decay-rate",  # decay rate, how much sparsity you want to decay each step
+            type=float,
+            default=TrainingArgs.VSA_decay_rate,
+            help="VSA decay rate")
+        parser.add_argument(
+            "--VSA-decay-interval-steps",  # how many steps for training with current sparsity
+            type=int,
+            default=TrainingArgs.VSA_decay_interval_steps,
+            help="VSA decay interval steps")
+        parser.add_argument(
+            "--VSA-init-sparsity",
+            type=float,
+            default=TrainingArgs.VSA_init_sparsity,
+            help="Initial sparsity to start from (default 0)")
+        parser.add_argument(
+            "--VSA-warmup-steps",
+            type=int,
+            default=TrainingArgs.VSA_warmup_steps,
+            help="Keep init sparsity for N steps before ramping (default 0)")
+        parser.add_argument("--lora-training", action=StoreBoolean, help="Whether to use LoRA training")
+        parser.add_argument("--lora-rank", type=int, help="LoRA rank")
+        parser.add_argument("--lora-alpha", type=int, help="LoRA alpha")
+        parser.add_argument(
+            "--ltx2-first-frame-conditioning-p",
+            type=float,
+            default=TrainingArgs.ltx2_first_frame_conditioning_p,
+            help="Probability of conditioning on the first frame during LTX-2 training",
+        )
+
+        # V-MoBA parameters
+        parser.add_argument(
+            "--moba-config-path",
+            type=str,
+            default=None,
+            help="Path to a JSON file containing V-MoBA specific configurations.",
+        )
+
+        # Distillation arguments
+        parser.add_argument("--generator-update-interval",
+                            type=int,
+                            default=TrainingArgs.generator_update_interval,
+                            help="Ratio of student updates to critic updates.")
+        parser.add_argument(
+            "--dfake-gen-update-ratio",
+            type=int,
+            default=TrainingArgs.dfake_gen_update_ratio,
+            help="Self-forcing: How often to train generator vs critic (train generator every N steps).")
+        parser.add_argument("--min-timestep-ratio",
+                            type=float,
+                            default=TrainingArgs.min_timestep_ratio,
+                            help="Minimum step ratio")
+        parser.add_argument("--max-timestep-ratio",
+                            type=float,
+                            default=TrainingArgs.max_timestep_ratio,
+                            help="Maximum step ratio")
+        parser.add_argument("--real-score-guidance-scale",
+                            type=float,
+                            default=TrainingArgs.real_score_guidance_scale,
+                            help="Teacher guidance scale")
+        parser.add_argument("--fake-score-learning-rate",
+                            type=float,
+                            default=TrainingArgs.fake_score_learning_rate,
+                            help="Learning rate for fake score transformer")
+        parser.add_argument("--fake-score-betas",
+                            type=str,
+                            default=TrainingArgs.fake_score_betas,
+                            help="Betas for fake score optimizer (format: 'beta1,beta2')")
+        parser.add_argument("--fake-score-lr-scheduler",
+                            type=str,
+                            default=TrainingArgs.fake_score_lr_scheduler,
+                            help="Learning rate scheduler for fake score transformer")
+        parser.add_argument("--log-visualization", action=StoreBoolean, help="Whether to log visualization")
+        parser.add_argument("--simulate-generator-forward",
+                            action=StoreBoolean,
+                            help="Whether to simulate generator forward to match inference")
+        parser.add_argument("--warp-denoising-step",
+                            action=StoreBoolean,
+                            help="Whether to warp denoising step according to the scheduler time shift")
+
+        # Self-forcing specific arguments
+        parser.add_argument("--num-frame-per-block",
+                            type=int,
+                            default=TrainingArgs.num_frame_per_block,
+                            help="Number of frames per block for causal generation")
+        parser.add_argument("--independent-first-frame",
+                            action=StoreBoolean,
+                            help="Whether the first frame is independent in causal generation")
+        parser.add_argument("--enable-gradient-masking",
+                            action=StoreBoolean,
+                            help="Whether to enable frame-level gradient masking")
+        parser.add_argument("--gradient-mask-last-n-frames",
+                            type=int,
+                            default=TrainingArgs.gradient_mask_last_n_frames,
+                            help="Number of last frames to enable gradients for")
+        parser.add_argument("--validate-cache-structure",
+                            action=StoreBoolean,
+                            help="Whether to validate KV cache structure (debug flag)")
+        parser.add_argument("--same-step-across-blocks",
+                            action=StoreBoolean,
+                            help="Whether to use the same exit timestep for all blocks")
+        parser.add_argument("--last-step-only",
+                            action=StoreBoolean,
+                            help="Whether to only use the last timestep for training")
+        parser.add_argument("--context-noise",
+                            type=int,
+                            default=TrainingArgs.context_noise,
+                            help="Context noise level for cache updates")
+
+        return parser
+
+
+def parse_int_list(value: str) -> list[int]:
+    if not value:
+        return []
+    return [int(x.strip()) for x in value.split(",")]

backend_snapshot/fastvideo/pipelines/basic/wan/wan_pipeline.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan video diffusion pipeline implementation.
+
+This module contains an implementation of the Wan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler)
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage,
+                                        LatentPreparationStage, TextEncodingStage, TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class WanPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Wan video diffusion pipeline with LoRA support.
+    """
+
+    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers.
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
+                                                    transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(transformer=self.get_module("transformer"),
+                                            transformer_2=self.get_module("transformer_2", None),
+                                            scheduler=self.get_module("scheduler"),
+                                            vae=self.get_module("vae"),
+                                            pipeline=self))
+
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
+
+
+EntryClass = WanPipeline

backend_snapshot/fastvideo/pipelines/composed_pipeline_base.py

@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Base class for composed pipelines.
+
+This module defines the base class for pipelines that are composed of multiple stages.
+"""
+
+import argparse
+import os
+from abc import ABC, abstractmethod
+from typing import Any, cast
+
+import torch
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, get_world_group)
+from fastvideo.distributed.communication_op import (warmup_sequence_parallel_communication)
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.profiler import get_or_create_profiler
+from fastvideo.models.loader.component_loader import PipelineComponentLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages import PipelineStage
+import fastvideo.envs as envs
+from fastvideo.utils import (maybe_download_model, verify_model_config_and_directory)
+
+logger = init_logger(__name__)
+
+
+class ComposedPipelineBase(ABC):
+    """
+    Base class for pipelines composed of multiple stages.
+    
+    This class provides the framework for creating pipelines by composing multiple
+    stages together. Each stage is responsible for a specific part of the diffusion
+    process, and the pipeline orchestrates the execution of these stages.
+    """
+
+    is_video_pipeline: bool = False  # To be overridden by video pipelines
+    _required_config_modules: list[str] = []
+    _extra_config_module_map: dict[str, str] = {}
+    training_args: Any = None
+    fastvideo_args: Any = None
+    modules: dict[str, Any] = {}
+    # do not need to include moe related transformers
+    trainable_transformer_names: list[str] = ["transformer"]
+    trainable_transformer_modules: dict[str, torch.nn.Module] = {}
+    post_init_called: bool = False
+
+    # TODO(will): args should support both inference args and training args
+    def __init__(self,
+                 model_path: str,
+                 fastvideo_args: FastVideoArgs | TrainingArgs,
+                 required_config_modules: list[str] | None = None,
+                 loaded_modules: dict[str, torch.nn.Module] | None = None):
+        """
+        Initialize the pipeline. After __init__, the pipeline should be ready to
+        use. The pipeline should be stateless and not hold any batch state.
+        """
+        self.fastvideo_args = fastvideo_args
+
+        self.model_path: str = model_path
+        self._stages: list[PipelineStage] = []
+        self._stage_name_mapping: dict[str, PipelineStage] = {}
+
+        if required_config_modules is not None:
+            self._required_config_modules = required_config_modules
+
+        if self._required_config_modules is None:
+            raise NotImplementedError("Subclass must set _required_config_modules")
+
+        maybe_init_distributed_environment_and_model_parallel(fastvideo_args.tp_size, fastvideo_args.sp_size)
+
+        # Torch profiler. Enabled and configured through env vars:
+        # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
+        trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR
+        self.profiler_controller = get_or_create_profiler(trace_dir)
+        self.profiler = self.profiler_controller.profiler
+
+        self.local_rank = get_world_group().local_rank
+
+        # Load modules directly in initialization
+        logger.info("Loading pipeline modules...")
+        with self.profiler_controller.region("profiler_region_model_loading"):
+            self.modules = self.load_modules(fastvideo_args, loaded_modules)
+
+    def set_trainable(self) -> None:
+        # Only train DiT
+        if getattr(self.fastvideo_args, "training_mode", False):
+            for name, module in self.trainable_transformer_modules.items():
+                logger.info("Setting %s to requires_grad=True", name)
+                if not isinstance(module, torch.nn.Module):
+                    logger.info("Skipping %s because it is not a torch.nn.Module", name)
+                    continue
+                module.requires_grad_(True)
+                module.train()
+
+    @staticmethod
+    def _compile_with_conditions(
+        module: torch.nn.Module,
+        compile_kwargs: dict[str, Any],
+    ) -> int:
+        """Compile submodules that match module._compile_conditions."""
+        compile_conditions = getattr(module, "_compile_conditions", None)
+        if not compile_conditions:
+            return 0
+
+        compiled_count = 0
+        for name, submodule in module.named_modules():
+            if not name:
+                continue
+            if any(cond(name, submodule) for cond in compile_conditions):
+                submodule.forward = torch.compile(submodule.forward, **compile_kwargs)
+                compiled_count += 1
+        return compiled_count
+
+    def _maybe_compile_pipeline_module(
+        self,
+        module_name: str,
+        fsdp_module_cls: type | None,
+        compile_kwargs: dict[str, Any],
+    ) -> None:
+        if module_name not in self.modules:
+            return
+
+        module = self.modules[module_name]
+        if fsdp_module_cls is not None and isinstance(module, fsdp_module_cls):
+            logger.info(
+                "%s is already FSDP-wrapped; skipping torch.compile in pipeline",
+                module_name.capitalize(),
+            )
+            return
+
+        compiled_count = self._compile_with_conditions(module, compile_kwargs)
+        if compiled_count > 0:
+            logger.info(
+                "Enabled torch.compile for %d submodules in %s via _compile_conditions with kwargs=%s",
+                compiled_count,
+                module_name,
+                compile_kwargs,
+            )
+            return
+
+        # Backward-compatible fallback: compile full module if no condition matched.
+        logger.info("Enabling torch.compile for %s with kwargs=%s", module_name, compile_kwargs)
+        self.modules[module_name] = torch.compile(module, **compile_kwargs)
+
+    def post_init(self) -> None:
+        assert self.fastvideo_args is not None, "fastvideo_args must be set"
+        if self.post_init_called:
+            return
+        self.post_init_called = True
+        if self.fastvideo_args.training_mode:
+            assert isinstance(self.fastvideo_args, TrainingArgs)
+            self.training_args = self.fastvideo_args
+            assert self.training_args is not None
+            self.initialize_training_pipeline(self.training_args)
+            if self.training_args.log_validation:
+                self.initialize_validation_pipeline(self.training_args)
+
+        self.initialize_pipeline(self.fastvideo_args)
+        if self.fastvideo_args.enable_torch_compile:
+            if self.fastvideo_args.training_mode:
+                logger.info("Torch Compile enabled via FSDP loader for training; skipping additional pipeline compile")
+            else:
+                fsdp_module_cls = None
+                try:
+                    from torch.distributed.fsdp import FSDPModule  # type: ignore
+                    fsdp_module_cls = FSDPModule
+                except Exception:  # pragma: no cover - FSDP not always available
+                    fsdp_module_cls = None
+
+                compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {}
+                self._maybe_compile_pipeline_module(
+                    module_name="transformer",
+                    fsdp_module_cls=fsdp_module_cls,
+                    compile_kwargs=compile_kwargs,
+                )
+                self._maybe_compile_pipeline_module(
+                    module_name="transformer_2",
+                    fsdp_module_cls=fsdp_module_cls,
+                    compile_kwargs=compile_kwargs,
+                )
+                logger.info("Torch Compile enabled for DiT")
+
+        if not self.fastvideo_args.training_mode:
+            logger.info("Creating pipeline stages...")
+            self.create_pipeline_stages(self.fastvideo_args)
+
+            # Warmup NCCL communicators for sequence parallelism to avoid
+            # slow first forward pass due to lazy initialization
+            warmup_sequence_parallel_communication()
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError("if training_mode is True, the pipeline must implement this method")
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError("if log_validation is True, the pipeline must implement this method")
+
+    @classmethod
+    def from_pretrained(cls,
+                        model_path: str,
+                        device: str | None = None,
+                        torch_dtype: torch.dtype | None = None,
+                        pipeline_config: str | PipelineConfig | None = None,
+                        args: argparse.Namespace | FastVideoArgs | TrainingArgs | None = None,
+                        required_config_modules: list[str] | None = None,
+                        loaded_modules: dict[str, torch.nn.Module]
+                        | None = None,
+                        **kwargs) -> "ComposedPipelineBase":
+        """
+        Load a pipeline from a pretrained model.
+        loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
+        If provided, loaded_modules will be used instead of loading from config/pretrained weights.
+        """
+        if args is None or (isinstance(args, FastVideoArgs) and args.inference_mode):
+
+            kwargs['model_path'] = model_path
+            fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+        else:
+            if isinstance(args, TrainingArgs):
+                fastvideo_args = args
+            else:
+                assert isinstance(args, argparse.Namespace), "training mode expects argparse.Namespace args"
+                fastvideo_args = TrainingArgs.from_cli_args(args)
+            # TODO(will): fix this so that its not so ugly
+            fastvideo_args.model_path = model_path
+            for key, value in kwargs.items():
+                setattr(fastvideo_args, key, value)
+
+            fastvideo_args.dit_cpu_offload = False
+            # we hijack the precision to be the master weight type so that the
+            # model is loaded with the correct precision. Subsequently we will
+            # use FSDP2's MixedPrecisionPolicy to set the precision for the
+            # fwd, bwd, and other operations' precision.
+            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+
+        logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
+
+        pipe = cls(model_path,
+                   fastvideo_args,
+                   required_config_modules=required_config_modules,
+                   loaded_modules=loaded_modules)
+        pipe.post_init()
+        return pipe
+
+    def get_module(self, module_name: str, default_value: Any = None) -> Any:
+        if module_name not in self.modules:
+            return default_value
+        return self.modules[module_name]
+
+    def add_module(self, module_name: str, module: Any):
+        self.modules[module_name] = module
+
+    def __getattr__(self, name: str) -> Any:
+        if "_stage_name_mapping" in self.__dict__ and name in self._stage_name_mapping:
+            return self._stage_name_mapping[name]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def _load_config(self, model_path: str) -> dict[str, Any]:
+        model_path = maybe_download_model(self.model_path)
+        self.model_path = model_path
+        # fastvideo_args.downloaded_model_path = model_path
+        logger.info("Model path: %s", model_path)
+        config = verify_model_config_and_directory(model_path)
+        return cast(dict[str, Any], config)
+
+    @property
+    def required_config_modules(self) -> list[str]:
+        """
+        List of modules that are required by the pipeline. The names should match
+        the diffusers directory and model_index.json file. These modules will be
+        loaded using the PipelineComponentLoader and made available in the
+        modules dictionary. Access these modules using the get_module method.
+
+        class ConcretePipeline(ComposedPipelineBase):
+            _required_config_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"]
+            
+
+            @property
+            def required_config_modules(self):
+                return self._required_config_modules
+        """
+        return self._required_config_modules
+
+    @property
+    def stages(self) -> list[PipelineStage]:
+        """
+        List of stages in the pipeline.
+        """
+        return self._stages
+
+    @abstractmethod
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """
+        Create the inference pipeline stages.
+        """
+        raise NotImplementedError
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        Create the training pipeline stages.
+        """
+        raise NotImplementedError
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """
+        Initialize the pipeline.
+        """
+        return
+
+    def load_modules(self,
+                     fastvideo_args: FastVideoArgs,
+                     loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]:
+        """
+        Load the modules from the config.
+        loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, 
+        If provided, loaded_modules will be used instead of loading from config/pretrained weights.
+        """
+
+        model_index = self._load_config(self.model_path)
+        logger.info("Loading pipeline modules from config: %s", model_index)
+
+        # remove keys that are not pipeline modules
+        model_index.pop("_class_name")
+        model_index.pop("_diffusers_version")
+        model_index.pop("_name_or_path", None)
+        model_index.pop("workload_type", None)
+        if "boundary_ratio" in model_index and model_index["boundary_ratio"] is not None:
+            logger.info("MoE pipeline detected. Adding transformer_2 to self.required_config_modules...")
+            self.required_config_modules.append("transformer_2")
+            logger.info("MoE pipeline detected. Setting boundary ratio to %s", model_index["boundary_ratio"])
+            fastvideo_args.pipeline_config.dit_config.boundary_ratio = model_index["boundary_ratio"]
+
+        model_index.pop("boundary_ratio", None)
+        # used by Wan2.2 ti2v
+        model_index.pop("expand_timesteps", None)
+
+        # some sanity checks
+        assert len(model_index) > 1, "model_index.json must contain at least one pipeline module"
+
+        for module_name in self.required_config_modules:
+            if module_name not in model_index and module_name in self._extra_config_module_map:
+                extra_module_value = self._extra_config_module_map[module_name]
+                logger.warning(
+                    "model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.",
+                    module_name, module_name, extra_module_value)
+                if extra_module_value in model_index:
+                    logger.info("Using module %s for %s", extra_module_value, module_name)
+                    model_index[module_name] = model_index[extra_module_value]
+                    continue
+                else:
+                    raise ValueError(
+                        f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}"
+                    )
+
+        # all the component models used by the pipeline
+        required_modules = self.required_config_modules
+        logger.info("Loading required modules: %s", required_modules)
+
+        modules = {}
+        for module_name, module_spec in model_index.items():
+            if not isinstance(module_spec, list | tuple):
+                logger.info(
+                    "Skipping non-module config entry %s=%s",
+                    module_name,
+                    module_spec,
+                )
+                continue
+            if len(module_spec) < 1:
+                logger.warning(
+                    "Skipping module %s due to invalid empty spec in model_index.json",
+                    module_name,
+                )
+                continue
+            transformers_or_diffusers = module_spec[0]
+            if transformers_or_diffusers is None:
+                logger.warning("Module %s in model_index.json has null value, removing from required_config_modules",
+                               module_name)
+                if module_name in self.required_config_modules:
+                    self.required_config_modules.remove(module_name)
+                continue
+            if module_name not in required_modules:
+                logger.info("Skipping module %s", module_name)
+                continue
+            if loaded_modules is not None and module_name in loaded_modules:
+                logger.info("Using module %s already provided", module_name)
+                modules[module_name] = loaded_modules[module_name]
+                continue
+
+            # we load the module from the extra config module map if it exists
+            if module_name in self._extra_config_module_map:
+                load_module_name = self._extra_config_module_map[module_name]
+            else:
+                load_module_name = module_name
+
+            component_model_path = os.path.join(self.model_path, load_module_name)
+            module = PipelineComponentLoader.load_module(
+                module_name=load_module_name,
+                component_model_path=component_model_path,
+                transformers_or_diffusers=transformers_or_diffusers,
+                fastvideo_args=fastvideo_args,
+            )
+            logger.info("Loaded module %s from %s", module_name, component_model_path)
+
+            if module_name in modules:
+                logger.warning("Overwriting module %s", module_name)
+            modules[module_name] = module
+
+        # Check if all required modules were loaded
+        for module_name in required_modules:
+            if module_name not in modules or modules[module_name] is None:
+                raise ValueError(
+                    f"Required module key: {module_name} value: {modules.get(module_name)} was not found in loaded modules {modules.keys()}"
+                )
+
+        return modules
+
+    def add_stage(self, stage_name: str, stage: PipelineStage):
+        assert self.modules is not None, "No modules are registered"
+        self._stages.append(stage)
+        self._stage_name_mapping[stage_name] = stage
+        setattr(self, stage_name, stage)
+
+    def profile(self, is_start: bool = True):
+        if self.profiler is None:
+            raise RuntimeError("Profiler is not enabled.")
+        if is_start:
+            self.profiler.start()
+        else:
+            self.profiler.stop()
+            # only print profiler results on rank 0
+            if self.local_rank == 0:
+                print(self.profiler.key_averages().table(sort_by="self_cuda_time_total"))
+
+    # TODO(will): don't hardcode no_grad
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Generate a video or image using the pipeline.
+        
+        Args:
+            batch: The batch to generate from.
+            fastvideo_args: The inference arguments.
+        Returns:
+            ForwardBatch: The batch with the generated video or image.
+        """
+        if not self.post_init_called:
+            self.post_init()
+
+        # Execute each stage
+        logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys())
+        # logger.info("Batch: %s", batch)
+        for stage in self.stages:
+            batch = stage(batch, fastvideo_args)
+
+        # Return the output
+        return batch
+
+    def train(self) -> None:
+        raise NotImplementedError("if training_mode is True, the pipeline must implement this method")
+
+    def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        raise NotImplementedError(f"{type(self).__name__} does not support streaming_reset")
+
+    def streaming_step(self, *args: Any, **kwargs: Any) -> ForwardBatch:
+        raise NotImplementedError(f"{type(self).__name__} does not support streaming_step")
+
+    def streaming_clear(self) -> None:
+        raise NotImplementedError(f"{type(self).__name__} does not support streaming_clear")

backend_snapshot/manifest.sha256

@@ -1,17 +1,33 @@
-45ff4b677a84fad92bd2ff596bf432cb1b9386c5923b6c0824f896074e7cfbc6  ./README.md
+033b7cce6eb0ead450a35b910adc9ae1323b8d2323aea6731b72b940e222fb46  ./README.md
+9c7dec8f1b8160954d0566231b0952a5f6a5d81f546affd71d190b2b3fc79cb6  ./examples/inference/basic/basic.py
 9d1d8dc58aab529270fe31eb1735d6a1382c0c6d36fccca122a8dbffa1b714fd  ./fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ours_p.py
 211c7f0445fbe9488250f01fa83457c6620e83bd6f3877db791fd155de93c08b  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton_ours_p.py
 3f3a407a88612ea17ad65e1b6b9cf6b7b02df56956d8301c4b13bffa92095016  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py
 56f17c602dede53c7c3677058f81274681530f1b83c086d9d1d44c6b51feefbb  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py
+58f4ac013e6755336212a7a6c9948b19dab0dafc00f4a3298591598df270cb39  ./fastvideo/api/compat.py
 2b821b0e2e7bdb3581be6312ebbece42380a6ee28a7a982f0cf2dc71fab849c8  ./fastvideo/attention/backends/sparse_fp4_ours_p_attn.py
 a97adcc52d7558c49f418c09395fd1665e988ad290d2276b95f21dfca0f8eb7d  ./fastvideo/attention/backends/video_sparse_attn.py
 79ef6f38ec0f5bfe16b2b98327ad2ccd15f3c863dd87fd03affc5dbdaa0a8224  ./fastvideo/configs/models/dits/base.py
+4bda44746a3626551ea9a9380d890f036087092fb99fce2d302642cce14a97ed  ./fastvideo/configs/pipelines/wan.py
+5926e29a594db13b116922f131db50631bf8adbf90fe5cec00a5e2f446bfb4ca  ./fastvideo/configs/sample/base.py
+d99adcf607d982b38bbb5a70be60bf87f35d0e9f6f50752f3bceb68b34ce46c2  ./fastvideo/configs/sample/wan.py
+49775ce42fd9643c78d8fad4ab8248c1755c7f1524ad771cbd1863d76c513c38  ./fastvideo/configs/wan_1.3B_t2v_pipeline.json
+ae2d8309472b09927da3e450dea52d9715dcabe5d6722fc2917130ae8d85adb4  ./fastvideo/entrypoints/cli/generate.py
+d0466769626e7fd497376c544904d56ba62847745eb52527896d96b99d76ba03  ./fastvideo/entrypoints/video_generator.py
+73afe6b2ebe0f8cfe0a8ec762a7126161621ad97a64ebad628995f4a164b8b0e  ./fastvideo/fastvideo_args.py
 ddcab6f4fd33c9813840571b6bf83bbbcea164b564166951ed4301297db6cef0  ./fastvideo/forward_context.py
+e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855  ./fastvideo/pipelines/basic/wan/__init__.py
+deac1e22530a6a41c501629f5e8fce47a7af4e008f321cc8a4d734c5120ef4fe  ./fastvideo/pipelines/basic/wan/wan_pipeline.py
+8908223b3ff99cdb3206148a68a730c2a13d554a2fb1316db6f2f9672efac9e8  ./fastvideo/pipelines/composed_pipeline_base.py
 6cfd128e782b7787a27ddd28a5e2d50cb4b0e2e9425d51d9780f14c91e8206f0  ./fastvideo/pipelines/stages/denoising.py
 489388dbdd9e5e3ad24db3012bd9b108794509a9729891d7dd315a102abba828  ./fastvideo/platforms/cuda.py
 c046b1914041b59254bcdfe577aed20d6f007a72632ea1fe1ae92fa678eca760  ./fastvideo/platforms/interface.py
 2456d39ca28019e12bb7ab007774e86348f0582a017bf0e6c91e2a01d654a1a0  ./fastvideo/train/models/wan/wan.py
 bc46e84b732567de6c0325223405daecd1226c623e303be33c7be9b5b7fdec08  ./fastvideo/training/training_pipeline.py
 1d3898fa37e21029df6c37e05dc34ed7805a211c2f87de6642db890e5a8c6f2e  ./fastvideo/training/wan_training_pipeline.py
+57fd33c78a16c9b4a239734c76726d94df1daf86f15cf22451c6107fcc197834  ./scripts/inference/run_sfp4_ours_p_checkpoint_750.sh
+0162b26dddb2a249e4e2cc56a7a28eba4fb3ea77e938eaf9c6b16a241edcd3ec  ./scripts/inference/run_sfp4_single.sh
+159579109f9fb7d7dded977f5c7fc974583c95f3f895e418d6c466463c036304  ./scripts/inference/run_validate_and_gen.sh
 5c982b64653fae83ebfdeb43fda8f29b3e2cb581fb4daee38cd3cf56aa9d73f5  ./scripts/training/run_sparse_fp4_train_v4_1n_sparse09_hpo_on_ours_p_init2050_interactive.sh
 5c1d5ce9ecc8b90e59ddfc2ddb3e2dae500bcd3acb90429c901444b1630f05fb  ./scripts/training/run_sparse_fp4_train_v4_common.sh
+75455829ca55a80daaa7e3c7faa080b6eec3c7109bdbd1e198b722face62eed0  ./training_attention_settings.json

backend_snapshot/scripts/inference/run_sfp4_ours_p_checkpoint_750.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+set -euo pipefail
+
+REPO_ROOT="${REPO_ROOT:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/yitongl/code/FastVideo}"
+MODEL_PATH="${MODEL_PATH:-Wan-AI/Wan2.1-T2V-1.3B-Diffusers}"
+HF_REPO="${HF_REPO:-yitongl/sparse_quant_exp}"
+CHECKPOINT_DIR="${CHECKPOINT_DIR:-${REPO_ROOT}/checkpoints/hf_download/sparse_quant_exp}"
+WEIGHTS_PATH="${WEIGHTS_PATH:-${CHECKPOINT_DIR}/transformer/diffusion_pytorch_model.safetensors}"
+OUTPUT_DIR="${OUTPUT_DIR:-${REPO_ROOT}/outputs/sfp4_v4_sparse09_checkpoint_750}"
+PROMPT="${PROMPT:-A cinematic shot of a futuristic city street at dusk, reflective pavement, soft volumetric light, detailed motion, stable camera.}"
+NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards}"
+
+cd "${REPO_ROOT}"
+
+if [[ -f ".venv/bin/activate" ]]; then
+    source .venv/bin/activate
+fi
+
+export PYTHONPATH="${REPO_ROOT}/fastvideo-kernel/python:${REPO_ROOT}/fastvideo-kernel:${PYTHONPATH:-}"
+export FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_OURS_P_ATTN
+export FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O=1
+
+if [[ ! -f "${WEIGHTS_PATH}" ]]; then
+    echo "Missing ${WEIGHTS_PATH}"
+    echo "Download the uploaded transformer weights first:"
+    echo "  hf download ${HF_REPO} --repo-type model --local-dir ${CHECKPOINT_DIR} --include 'transformer/*'"
+    exit 1
+fi
+
+mkdir -p "${OUTPUT_DIR}"
+
+fastvideo generate \
+    --model-path "${MODEL_PATH}" \
+    --init-weights-from-safetensors "${WEIGHTS_PATH}" \
+    --sp-size 1 \
+    --tp-size 1 \
+    --num-gpus 1 \
+    --dit-cpu-offload False \
+    --vae-cpu-offload False \
+    --text-encoder-cpu-offload True \
+    --pin-cpu-memory False \
+    --height 448 \
+    --width 832 \
+    --num-frames 77 \
+    --num-inference-steps 50 \
+    --fps 16 \
+    --guidance-scale 5.0 \
+    --flow-shift 1.0 \
+    --prompt "${PROMPT}" \
+    --negative-prompt "${NEGATIVE_PROMPT}" \
+    --seed 1000 \
+    --VSA-sparsity 0.9 \
+    --output-path "${OUTPUT_DIR}/"

backend_snapshot/scripts/inference/run_sfp4_single.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+#SBATCH --job-name=sfp4-s0
+#SBATCH --account=nvr_elm_llm
+#SBATCH --partition=interactive
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=64G
+#SBATCH --time=00:30:00
+#SBATCH --output=slurm_logs/sfp4_s0_%j.out
+#SBATCH --error=slurm_logs/sfp4_s0_%j.err
+
+set -ex
+cd /lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/yitongl/code/FastVideo
+source .venv/bin/activate
+export PYTHONPATH=fastvideo-kernel/python:fastvideo-kernel:$PYTHONPATH
+export FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_ATTN
+
+mkdir -p outputs_sfp4_s0
+
+# Same prompt, seed, params as dense FP4 run
+fastvideo generate \
+    --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --sp-size 1 --tp-size 1 --num-gpus 1 \
+    --dit-cpu-offload False --vae-cpu-offload False \
+    --text-encoder-cpu-offload True --pin-cpu-memory False \
+    --height 480 --width 832 --num-frames 81 \
+    --num-inference-steps 50 --fps 16 \
+    --guidance-scale 6.0 --flow-shift 8.0 \
+    --prompt "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." \
+    --seed 1024 \
+    --VSA-sparsity 0.0 \
+    --output-path outputs_sfp4_s0/
+
+echo "=== Done ==="
+ls -lh outputs_sfp4_s0/*.mp4
+echo "--- Dense FP4 reference ---"
+ls -lh outputs_dense_fp4/*.mp4

backend_snapshot/scripts/inference/run_validate_and_gen.sh

@@ -0,0 +1,91 @@
+#!/bin/bash
+#SBATCH --job-name=sfp4-val-gen
+#SBATCH --account=nvr_elm_llm
+#SBATCH --partition=interactive
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --time=02:00:00
+#SBATCH --output=slurm_logs/sfp4_val_gen_%j.out
+#SBATCH --error=slurm_logs/sfp4_val_gen_%j.err
+
+set -ex
+
+REPO_ROOT="/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/yitongl/code/FastVideo"
+KERNEL_ROOT="${REPO_ROOT}/fastvideo-kernel"
+
+mkdir -p "${REPO_ROOT}/slurm_logs"
+cd "${REPO_ROOT}"
+source .venv/bin/activate
+export PYTHONPATH="${KERNEL_ROOT}/python:${KERNEL_ROOT}:${PYTHONPATH}"
+
+echo "=== Environment ==="
+nvidia-smi -L | head -1
+python -c "import torch; print(f'torch={torch.__version__}, cuda={torch.cuda.is_available()}, gpus={torch.cuda.device_count()}')"
+python -c "import triton; print(f'triton={triton.__version__}')"
+
+echo ""
+echo "######################################################################"
+echo "# Generate 8 videos with sparse FP4 attention                       #"
+echo "######################################################################"
+cd "${REPO_ROOT}"
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+PROMPT="Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
+NEGATIVE_PROMPT="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+SEED=1024
+
+SPARSITY_LIST=(0.0 0.1 0.2 0.4 0.5 0.7 0.8 0.9)
+OUTPUT_BASE="${REPO_ROOT}/outputs_sparse_fp4_sweep"
+mkdir -p "${OUTPUT_BASE}"
+
+echo "Sparsity levels: ${SPARSITY_LIST[*]}"
+
+PIDS=()
+for i in $(seq 0 7); do
+    SPARSITY=${SPARSITY_LIST[$i]}
+    OUT_DIR="${OUTPUT_BASE}/sparsity_${SPARSITY}"
+    mkdir -p "${OUT_DIR}"
+    echo "[GPU ${i}] sparsity=${SPARSITY}"
+
+    (
+        export CUDA_VISIBLE_DEVICES=${i}
+        export FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_ATTN
+
+        fastvideo generate \
+            --model-path "${MODEL_PATH}" \
+            --sp-size 1 --tp-size 1 --num-gpus 1 \
+            --dit-cpu-offload False \
+            --vae-cpu-offload False \
+            --text-encoder-cpu-offload True \
+            --pin-cpu-memory False \
+            --height 480 --width 832 --num-frames 81 \
+            --num-inference-steps 50 --fps 16 \
+            --guidance-scale 6.0 --flow-shift 8.0 \
+            --prompt "${PROMPT}" \
+            --negative-prompt "${NEGATIVE_PROMPT}" \
+            --seed ${SEED} \
+            --VSA-sparsity ${SPARSITY} \
+            --output-path "${OUT_DIR}/" \
+            2>&1 | tee "${OUT_DIR}/log.txt"
+
+        echo "[GPU ${i}] sparsity=${SPARSITY} DONE"
+    ) &
+    PIDS+=($!)
+done
+
+echo "=== Waiting for all 8 jobs ==="
+FAIL=0
+for i in $(seq 0 7); do
+    wait ${PIDS[$i]} || { echo "[GPU ${i}] FAILED"; FAIL=1; }
+done
+
+echo ""
+if [ $FAIL -eq 0 ]; then
+    echo "=== All 8 videos generated ==="
+else
+    echo "=== Some failed ==="
+fi
+find "${OUTPUT_BASE}" -name "*.mp4" | sort
+echo "=== Done ==="

backend_snapshot/training_attention_settings.json

@@ -0,0 +1,62 @@
+{
+  "run_name": "sfp4_v4_sparse09_hpo_on_ours_p_init2050_1n_interactive",
+  "checkpoint": "checkpoint-750",
+  "training_method": "legacy_sft_wan_training_pipeline",
+  "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+  "init_weights_from_safetensors": "checkpoints/init/sfp4_v4_sparse06_hpo_on_ours_p_1n_interactive_v2_ckpt2050/transformer/diffusion_pytorch_model.safetensors",
+  "environment": {
+    "FASTVIDEO_ATTENTION_BACKEND": "SPARSE_FP4_OURS_P_ATTN",
+    "FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O": "1",
+    "FASTVIDEO_VALIDATION_ONE_PROMPT_PER_RANK": "1",
+    "WANDB_MODE": "online",
+    "WANDB_RESUME": "allow"
+  },
+  "vsa_schedule": {
+    "VSA_SPARSITY": 0.9,
+    "VSA_INIT_SPARSITY": 0.9,
+    "VSA_WARMUP_STEPS": 0,
+    "VSA_DECAY_RATE": 0.03,
+    "VSA_DECAY_INTERVAL_STEPS": 50,
+    "effective_sparsity_from_step_0": 0.9
+  },
+  "attention_semantics": {
+    "selected_backend": "SPARSE_FP4_OURS_P_ATTN",
+    "self_attention": {
+      "backend_path": "fastvideo/attention/backends/sparse_fp4_ours_p_attn.py",
+      "kernel_path": "fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton_ours_p.py",
+      "tile_size_video": [4, 4, 4],
+      "tile_tokens": 64,
+      "qkv_quantization": "FP4 fake quantization with STE, no q/k mean subtraction in quantization",
+      "block_selection": "top-k blocks from q_c @ k_c tile-mean scores",
+      "p_quantization": "group-local exp2(qk - group_max) FP4 fake quantization; compensation multiplies exp2(group_max - running_row_m)",
+      "dropped_tile_handling": "tile-level q_mean/k_mean score and mean_v compensation"
+    },
+    "cross_attention": {
+      "backend": "dense_sdpa",
+      "reason": "sparse_fp4_ours_p_attn.py treats query_length != key_length as cross attention and returns _dense_sdpa_blhd",
+      "quantized": false,
+      "sparse": false
+    },
+    "force_dense": {
+      "backend": "dense_sdpa",
+      "used_for": "teacher or explicitly forced dense paths, not the normal SFT student self-attention path"
+    }
+  },
+  "validation_and_checkpointing": {
+    "save_steps": 50,
+    "eval_steps": 50,
+    "validation_sampling_steps": 50,
+    "validation_guidance_scale": 5.0,
+    "checkpoints_total_limit": 5,
+    "flow_shift": 1.0
+  },
+  "training_shape": {
+    "num_latent_t": 20,
+    "num_frames": 77,
+    "height": 448,
+    "width": 832,
+    "batch_size_per_gpu": 1,
+    "sp_size": 1,
+    "tp_size": 1
+  }
+}