# Inference Architecture

This section documents the FastVideo inference pipeline: how models are
discovered, configs resolved, components loaded, and stages composed to
generate video. Training-specific code paths (FSDP, gradient checkpointing)
are out of scope.

## Registries

FastVideo uses three registries that work together to resolve a
user-provided `model_path` into a runnable pipeline.

### Model Registry (`fastvideo/models/registry.py`)

Maps HuggingFace architecture class names (e.g. `"WanTransformer3DModel"`)
to FastVideo model classes. Two discovery mechanisms:

1. **Hardcoded dicts** — `_TEXT_TO_VIDEO_DIT_MODELS`, `_VAE_MODELS`,
   `_SCHEDULERS`, `_TEXT_ENCODER_MODELS`, `_IMAGE_ENCODER_MODELS`,
   `_UPSAMPLERS`, `_AUDIO_MODELS`. Each entry is
   `{hf_class_name: (component_name, module_name, class_name)}`.
2. **AST-based discovery** — `_discover_and_register_models()` walks
   `fastvideo/models/` and parses each `.py` file's AST looking for an
   `EntryClass` variable assignment. Discovered models take priority over
   hardcoded entries. For example,
   `fastvideo/models/wan/transformer.py` exports
   `EntryClass = WanTransformer3DModel`.

Both feed into a unified `_FAST_VIDEO_MODELS` dict, which populates the
singleton `ModelRegistry` — an instance of `_ModelRegistry`. Components are
wrapped in `_LazyRegisteredModel` for deferred import.

**Key API:** `ModelRegistry.resolve_model_cls(architectures)` iterates
candidate architecture strings and returns the first matching
`(model_cls, arch)` tuple. Called by `TransformerLoader`, `VAELoader`, etc.

### Config Registry (`fastvideo/registry.py`)

Maps model paths and names to `(PipelineConfig, SamplingParam)` class
pairs. Registration happens at module load via `_register_configs()`, which
calls `register_configs()` for each model family:

```python
register_configs(
    sampling_param_cls=WanT2V_1_3B_SamplingParam,
    pipeline_config_cls=WanT2V480PConfig,
    hf_model_paths=["Wan-AI/Wan2.1-T2V-1.3B-Diffusers"],
    model_detectors=[lambda path: "wanpipeline" in path.lower()],
)
```

Each call populates three data structures:
- `_CONFIG_REGISTRY: dict[str, ConfigInfo]` — auto-incrementing ID to
  `ConfigInfo(sampling_param_cls, pipeline_config_cls)`.
- `_MODEL_HF_PATH_TO_NAME: dict[str, str]` — HF path to registry ID.
- `_MODEL_NAME_DETECTORS: list[tuple[str, Callable]]` — lambda detectors.

**Resolution priority** (`_get_config_info()`):
1. Exact HF path match in `_MODEL_HF_PATH_TO_NAME`.
2. Partial match on short model name (last path segment, case-insensitive).
3. Detector-based match — runs each detector against the lowercased path
   and the `_class_name` from `model_index.json`.
4. `RuntimeError` if no match.

**Top-level resolver:** `get_model_info(model_path, pipeline_type,
workload_type)` combines config resolution with pipeline resolution to
return a `ModelInfo(pipeline_cls, sampling_param_cls, pipeline_config_cls)`.

### Pipeline Registry (`fastvideo/pipelines/pipeline_registry.py`)

Discovers pipeline classes by scanning Python packages under
`fastvideo/pipelines/{basic,preprocess,training}/`.

`import_pipeline_classes()` iterates architecture subdirectories
(e.g. `wan/`, `hunyuan/`), imports each module, and collects those
exporting an `EntryClass` attribute. Supports single class or list.
Returns `{pipeline_type_str: {pipeline_class_name: pipeline_cls}}`.

`_PipelineRegistry.resolve_pipeline_cls(pipeline_name, pipeline_type,
workload_type)` looks up the pipeline class by the `_class_name` field
from `model_index.json`.

## Config Mechanism

### Config Hierarchy

```
PipelineConfig                    (fastvideo/configs/pipelines/base.py)
├── WanT2V480PConfig              (fastvideo/models/wan/pipeline_config.py)
│   ├── WanT2V720PConfig
│   └── WanI2V480PConfig
├── HunyuanConfig                 (fastvideo/configs/pipelines/hunyuan.py)
├── LTX2T2VConfig                 (fastvideo/configs/pipelines/ltx2.py)
├── CosmosConfig                  (fastvideo/configs/pipelines/cosmos.py)
└── ... (15+ model families)
```

`PipelineConfig` holds:
- Video generation params: `embedded_cfg_scale`, `flow_shift`,
  `disable_autocast`, `is_causal`.
- Nested model configs: `dit_config: DiTConfig`, `vae_config: VAEConfig`,
  `text_encoder_configs: tuple[EncoderConfig, ...]`.
- Precision settings: `dit_precision`, `vae_precision`,
  `text_encoder_precisions`.

Model-specific subclasses override defaults. For example,
`WanT2V480PConfig` sets `flow_shift=3.0` and uses `WanVideoConfig` as
its DiT config.

Wan's `models/wan/definition.py` links each registered variant to its pipeline
config and sampling preset. The shared registry consumes these definitions
without changing detector precedence or checkpoint/override-based pipeline
selection. `configs/pipelines/wan.py` remains a compatibility import.

### ModelConfig / ArchConfig (`fastvideo/configs/models/base.py`)

`ModelConfig` wraps an `ArchConfig` using `__getattr__` proxy — attribute
access falls through to `arch_config` transparently. `ArchConfig` holds
architecture fields from `config.json` (hidden_size, num_attention_heads,
etc.) and is immutable after initialization. `update_model_arch()` writes
to `ArchConfig`; `update_model_config()` writes to `ModelConfig` fields.

Concrete hierarchy: `DiTConfig` → `DiTArchConfig`, `VAEConfig` →
`VAEArchConfig`, `EncoderConfig` → `EncoderArchConfig`.

### Config Construction

- `PipelineConfig.from_pretrained(model_path)` — resolves config class
  via `get_pipeline_config_cls_from_name()`, instantiates with defaults.
- `PipelineConfig.from_kwargs(kwargs)` — resolves class, optionally loads
  JSON via `load_from_json()`, then applies CLI overrides via
  `update_config_from_dict()`.
- `dump_to_json()` / `load_from_json()` — JSON persistence. Callable
  fields and `arch_config` are excluded from dumps.

### SamplingParam (`fastvideo/api/sampling_param.py`)

Generation parameters separate from pipeline config. Each model family
provides defaults via a profile (see `fastvideo/pipelines/basic/<family>/profiles.py`):

```python
sp = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
# sp.height == 480, sp.width == 832, sp.num_frames == 81, etc.
```

## Component Loading

### ComponentLoader (`fastvideo/models/loader/component_loader.py`)

Abstract base with a `load(model_path, fastvideo_args)` method.
`ComponentLoader.for_module_type(module_type, library)` is a factory
that dispatches to specialized loaders via a `module_loaders` dict:

| Module type | Loader class | Library |
|---|---|---|
| `scheduler` | `SchedulerLoader` | diffusers |
| `transformer`, `transformer_2`, `transformer_3` | `TransformerLoader` | diffusers |
| `vae` | `VAELoader` | diffusers |
| `text_encoder`, `text_encoder_2`, `text_encoder_3` | `TextEncoderLoader` | transformers |
| `tokenizer`, `tokenizer_2`, `tokenizer_3` | `TokenizerLoader` | transformers |
| `image_encoder` | `ImageEncoderLoader` | transformers |
| `image_processor`, `feature_extractor` | `ImageProcessorLoader` | transformers |
| `audio_vae`, `audio_decoder` | `AudioDecoderLoader` | diffusers |
| `vocoder` | `VocoderLoader` | diffusers |
| `upsampler`, `upsampler_2` | `UpsamplerLoader` | diffusers |

`TransformerLoader` reads `config.json` from the component directory,
resolves the class via `ModelRegistry.resolve_model_cls()`, instantiates
the model, and loads safetensors weights. CPU offload and layerwise
offload are applied based on `FastVideoArgs`.

Unknown module types fall back to `GenericComponentLoader`.

### model_index.json

Diffusers-format JSON at the model root. Keys are module names; values
are `[library, class_name]` tuples:

```json
{
    "_class_name": "WanPipeline",
    "_diffusers_version": "0.24.0",
    "transformer": ["diffusers", "WanTransformer3DModel"],
    "vae": ["diffusers", "AutoencoderKLWan"],
    "text_encoder": ["transformers", "UMT5EncoderModel"],
    "tokenizer": ["transformers", "AutoTokenizer"],
    "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"]
}
```

`ComposedPipelineBase.load_modules()` reads this file via
`_load_config()`, strips metadata keys (`_class_name`,
`_diffusers_version`, `_name_or_path`), detects MoE pipelines via
`boundary_ratio`, then loads each module listed in the pipeline's
`required_config_modules`. Modules not in `required_config_modules` are
skipped. `_extra_config_module_map` allows aliasing (e.g. mapping
`"transformer_2"` to an alternate directory name).

`PipelineComponentLoader.load_module()` orchestrates per-component
loading by calling `ComponentLoader.for_module_type()` then `.load()`.

## Stage Design

### PipelineStage (`fastvideo/pipelines/stages/base.py`)

Abstract base class using the Template Method pattern:

- `__call__(batch, fastvideo_args)` — orchestrates verification, timing,
  and error handling. Not overridden by subclasses.
- `forward(batch, fastvideo_args) -> ForwardBatch` — abstract, contains
  the stage logic.
- `verify_input()` / `verify_output()` — optional hooks returning
  `VerificationResult`. Default: no checks.

When `fastvideo_args.enable_stage_verification` is `True`, `__call__`
runs input verification before `forward()` and output verification after.
When `envs.FASTVIDEO_STAGE_LOGGING` is set, execution time is measured
with `torch.cuda.synchronize()` and logged.

### ForwardBatch (`fastvideo/pipelines/pipeline_batch_info.py`)

Dataclass carrying all pipeline state between stages. Key field groups:

- **Inputs**: `prompt`, `negative_prompt`, `image_path`, `pil_image`,
  `video_path`.
- **Embeddings**: `prompt_embeds: list[Tensor]`,
  `negative_prompt_embeds`, `prompt_attention_mask`, `image_embeds`.
- **Latents**: `latents`, `image_latent`, `noise_pred`,
  `lq_latents`.
- **Dimensions**: `height`, `width`, `num_frames`, `height_latents`,
  `width_latents`.
- **Scheduler**: `timesteps`, `num_inference_steps`, `guidance_scale`,
  `sigmas`.
- **Task-specific**: `mouse_cond`/`keyboard_cond` (Matrix-Game 2.0), `pose`
  (HYWorld), `c2ws_plucker_emb` (LingBotWorld).
- **Output**: `output: Tensor | None`.
- **Logging**: `logging_info: PipelineLoggingInfo`.

`__post_init__` enables CFG when `guidance_scale > 1.0` or LTX2 text
CFG scales differ from 1.0.

### Stage Catalog

Standard stages (typical execution order):

| Stage | File | Purpose |
|---|---|---|
| `InputValidationStage` | `stages/input_validation.py` | Validates input dimensions and types |
| `TextEncodingStage` | `stages/text_encoding.py` | Encodes prompts via text encoders |
| `ImageEncodingStage` | `stages/image_encoding.py` | Encodes input images (I2V pipelines) |
| `ConditioningStage` | `stages/conditioning.py` | Prepares conditioning embeddings |
| `TimestepPreparationStage` | `stages/timestep_preparation.py` | Sets up scheduler timesteps |
| `LatentPreparationStage` | `stages/latent_preparation.py` | Initializes noise latents |
| `DenoisingStage` | `stages/denoising.py` | Main diffusion denoising loop |
| `DecodingStage` | `stages/decoding.py` | Decodes latents to video via VAE |

Specialized variants: `CausalDenoisingStage`, `LTX2DenoisingStage`,
`HYWorldDenoisingStage`, `MatrixGame2CausalDenoisingStage`,
`SRDenoisingStage`, `LTX2AudioDecodingStage`, `SD35ConditioningStage`,
`LTX2TextEncodingStage`, `LTX2LatentPreparationStage`.

Wan owns its sampling recipes under `basic/wan/stages/`. `WanDenoisingStage`
specializes input packing, expert selection, timesteps, and first-frame
restoration around the shared dense loop. `WanFirstFrameEncodingStage`
produces normalized `ForwardBatch.first_frame_latent` before sampling; the
sampler no longer executes a VAE. Dense DMD and the two causal samplers have
family-local implementations and explicit scheduler ownership. Standard and
DMD causal sampling share cache allocation, not their sampling algorithm.
Legacy imports from `stages/` remain compatibility aliases. Sampling invariants
also live beside the code in `fastvideo/pipelines/basic/wan/AGENTS.md`.

### Verification System (`fastvideo/pipelines/stages/validators.py`)

`StageValidators` (aliased as `V`) provides static validators:
`not_none`, `positive_int`, `is_tensor`, `tensor_with_dims`,
`positive_int_divisible(divisor)`, etc.

`VerificationResult` collects check results:

```python
result = VerificationResult()
result.add_check("height", batch.height, V.positive_int_divisible(8))
result.add_check("width", batch.width, V.positive_int_divisible(8))
```

`is_valid()` returns whether all checks passed. `get_failure_summary()`
provides detailed error messages. Failed verification raises
`StageVerificationError`.

## Pipeline Architecture

### ComposedPipelineBase (`fastvideo/pipelines/composed_pipeline_base.py`)

Abstract base for all inference pipelines. Lifecycle:

1. **`__init__(model_path, fastvideo_args)`** — initializes distributed
   environment via `maybe_init_distributed_environment_and_model_parallel
   (tp_size, sp_size)`, then calls `load_modules()` to populate
   `self.modules`.
2. **`post_init()`** — calls `initialize_pipeline()` (model-specific
   setup), `create_pipeline_stages()` (abstract — subclasses wire stages),
   optionally applies `torch.compile` to transformers, and calls
   `warmup_sequence_parallel_communication()`.
3. **`forward(batch, fastvideo_args)`** — iterates `self.stages` calling
   each stage in order. Decorated with `@torch.no_grad()`.

Key class attributes:
- `_required_config_modules: list[str]` — module names to load from
  `model_index.json`.
- `_extra_config_module_map: dict[str, str]` — aliases for module dirs.
- `is_video_pipeline: bool` — whether this produces video output.

Key methods:
- `add_stage(name, stage)` — appends to `_stages` list and
  `_stage_name_mapping` dict, also sets attribute on `self`.
- `get_module(name, default)` — retrieves a loaded module.
- `from_pretrained(model_path, **kwargs)` — class method constructing
  `FastVideoArgs` and calling `cls(...)` then `post_init()`.

### LoRAPipeline (`fastvideo/pipelines/lora_pipeline.py`)

Extends `ComposedPipelineBase` with LoRA adapter support. Sits in the
MRO between the concrete pipeline and `ComposedPipelineBase`:

```python
class WanPipeline(LoRAPipeline, ComposedPipelineBase):
    ...
```

Key functionality:
- `convert_to_lora_layers()` — scans transformer blocks, replaces target
  linear layers (default: q/k/v/o projections) with LoRA equivalents via
  `get_lora_layer()`.
- `set_lora_adapter(path)` — loads safetensors containing `lora_A`,
  `lora_B`, `lora_alpha` and maps weights to internal layers.
- `merge_lora_weights()` / `unmerge_lora_weights()` — activates or
  deactivates LoRA in the forward pass.
- `LoRAModelLayers` — groups LoRA layers by transformer block for
  efficient layerwise offload.

### Distributed Inference

`maybe_init_distributed_environment_and_model_parallel(tp_size, sp_size)`
in `fastvideo/distributed/` initializes `torch.distributed` and creates
tensor-parallel (TP) and sequence-parallel (SP) process groups.

Key APIs: `get_tp_rank()`, `get_tp_world_size()`, `get_sp_rank()`,
`get_sp_world_size()`, `get_world_rank()`, `get_world_size()`.

`warmup_sequence_parallel_communication()` pre-warms NCCL communicators
to avoid slow first forward passes.

The default `mp` executor starts all workers locally. For synchronized offline
SPMD generation (including multi-node inference), select the
`external_launcher` executor in the typed run config and make
`generator.engine.num_gpus` equal the launcher's total world size. The
[MiniMax-H3 external-launcher example](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/minimax_h3_external_launcher.yaml)
contains a complete request:

```bash
torchrun --standalone --nproc-per-node=4 \
  -m fastvideo.entrypoints.cli.main generate \
  --config examples/inference/basic/minimax_h3_external_launcher.yaml
```

For two nodes with four GPUs each, launch the same command on both nodes and
set `NODE_RANK` to `0` or `1`:

```bash
torchrun --nnodes=2 --nproc-per-node=4 \
  --node-rank="$NODE_RANK" \
  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
  -m fastvideo.entrypoints.cli.main generate \
  --config examples/inference/basic/minimax_h3_external_launcher.yaml \
  --generator.engine.num_gpus 8 \
  --generator.engine.parallelism.sp_size 8
```

Native Slurm variables are also accepted: `SLURM_PROCID`, `SLURM_NTASKS`,
and `SLURM_LOCALID` map to global rank, world size, and local rank. Slurm does
not choose the env rendezvous endpoint, so export a reachable rank-zero host
and a job-unique port. Preserve a scheduler/CI-provided `MASTER_PORT` when one
already exists. This uniform `2 x 4` example exposes all four allocated GPUs
to every task so `SLURM_LOCALID` is a valid device index:

```bash
export MASTER_ADDR="${MASTER_ADDR:-$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)}"
export MASTER_PORT="${MASTER_PORT:-$((20000 + SLURM_JOB_ID % 20000))}"
srun --nodes=2 --ntasks-per-node=4 --gpus-per-node=4 --gpu-bind=none \
  --kill-on-bad-exit=1 \
  python -m fastvideo.entrypoints.cli.main generate \
  --config examples/inference/basic/minimax_h3_external_launcher.yaml \
  --generator.engine.num_gpus 8 \
  --generator.engine.parallelism.sp_size 8
```

Every rank must enter the same requests in the same order. The CLI config and
model/media inputs therefore need to resolve identically on every node. Rank
zero alone reads prompt-list files, creates output directories, saves media,
and returns generated frames/audio/latents/trajectories; only that rank needs
a writable output path. For MiniMax-H3 parallel VAE decode, every rank in the
SP group containing global rank zero participates in the decoder collectives,
but only rank zero assembles the CPU output. All ranks in every other SP group
skip the decode uniformly, so multiple SP groups do not materialize duplicate
outputs.

Set `generator.engine.parallelism.dist_timeout` to an optional positive
integer number of **seconds** to apply one timeout to the default process
group (when FastVideo initializes it) and every FastVideo-created device and
Gloo control group. Omitting the key preserves PyTorch's backend defaults. A
local-rank error, an out-of-range rendezvous port, or a device index outside
the process-visible accelerator set fails before model loading. The process-
group timeout bounds a missing collective participant;
`--kill-on-bad-exit=1` (or torchrun's worker-group supervision) remains
necessary to terminate the rest of the worker group after a rank failure.

`FASTVIDEO_EXTERNAL_LAUNCHER=1` is an equivalent offline-only opt-in for
callers that leave `execution_backend: mp`. External-launcher mode is rejected
by REST and streaming servers because those request loops do not provide a
rank-zero SPMD control plane.

### torch.compile Integration

When `fastvideo_args.enable_torch_compile` is `True`,
`_maybe_compile_pipeline_module()` checks for a `_compile_conditions`
attribute on the module. If present, only matching submodules are
compiled. Otherwise, the entire module is compiled. FSDP-wrapped
modules are skipped.

### Entry Points

**Python API** (`fastvideo/entrypoints/video_generator.py`):

```python
generator = VideoGenerator.from_pretrained(
    model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers",
    num_gpus=1, tp_size=1, sp_size=1,
)
result = generator.generate_video(
    prompt="A cat dancing",
    height=720, width=1280, num_frames=81,
)
```

**CLI** (`fastvideo/entrypoints/cli/`):

```bash
fastvideo generate \
    --model-path "Wan-AI/Wan2.1-T2V-14B-Diffusers" \
    --prompt "A cat dancing" \
    --num-gpus 1
```

**FastVideoArgs** (`fastvideo/fastvideo_args.py`): Central args dataclass.
Key fields: `model_path`, `mode` (`ExecutionMode`), `workload_type`
(`WorkloadType`), `pipeline_config` (`PipelineConfig`), `num_gpus`,
`tp_size`, `sp_size`, `lora_path`, `dit_cpu_offload`,
`dit_layerwise_offload`, `enable_torch_compile`,
`enable_stage_verification`.

Constructed via `FastVideoArgs.from_kwargs(**kwargs)` which resolves the
`PipelineConfig` from the registry, applies JSON config if provided, and
merges CLI overrides.

## End-to-End Inference Flow

```
User: VideoGenerator.from_pretrained(model_path, **kwargs)
  │
  ├─ FastVideoArgs.from_kwargs() → PipelineConfig resolved via registry
  ├─ get_model_info() → ModelInfo(pipeline_cls, sampling_param_cls, ...)
  │   ├─ model_index.json read → _class_name extracted
  │   ├─ pipeline_registry resolves pipeline_cls from _class_name
  │   └─ config_registry resolves config classes from model_path
  │
  ├─ pipeline_cls.__init__(model_path, fastvideo_args)
  │   ├─ maybe_init_distributed(tp_size, sp_size)
  │   └─ load_modules() → reads model_index.json, loads each component
  │       ├─ ComponentLoader.for_module_type() → specialized loader
  │       └─ loader.load() → model class resolved, weights loaded
  │
  └─ pipeline.post_init()
      ├─ initialize_pipeline() → model-specific setup
      ├─ create_pipeline_stages() → stages wired with modules
      ├─ torch.compile (if enabled)
      └─ warmup_sequence_parallel_communication()

User: generator.generate_video(prompt, ...)
  │
  ├─ ForwardBatch constructed from SamplingParam + user args
  └─ pipeline.forward(batch, fastvideo_args)
      ├─ InputValidationStage → validates dims
      ├─ TextEncodingStage → prompt → embeddings
      ├─ ConditioningStage → prepares conditioning
      ├─ TimestepPreparationStage → scheduler timesteps
      ├─ LatentPreparationStage → random noise
      ├─ DenoisingStage → iterative denoising loop
      └─ DecodingStage → latents → video frames
```

## Adding a New Model Family — Checklist

1. **Pipeline config** — Create a `PipelineConfig` subclass in
   `fastvideo/configs/pipelines/<model>.py`. Set DiT/VAE/encoder configs,
   flow_shift, precision defaults.

2. **Sampling param profile** — Create a profile in
   `fastvideo/pipelines/basic/<model>/profiles.py` with default height,
   width, num_frames, guidance_scale, num_inference_steps.

3. **Register configs** — In `fastvideo/registry.py`, add a
   `register_configs()` call inside `_register_configs()` with
   `hf_model_paths` and/or `model_detectors`.

4. **Pipeline class** — Create a subclass of `ComposedPipelineBase` (or
   `LoRAPipeline` + `ComposedPipelineBase`) in
   `fastvideo/pipelines/basic/<model>/<model>_pipeline.py`.
   - Set `_required_config_modules` listing needed components.
   - Implement `create_pipeline_stages()` wiring stages via `add_stage()`.
   - Optionally override `initialize_pipeline()` for custom setup.
   - Export `EntryClass = YourPipeline` at module level.

5. **Model classes** (if custom) — Add DiT/VAE implementations in
   `fastvideo/models/dits/` or `fastvideo/models/vaes/` with
   `EntryClass = YourModel`. The AST discovery will register them
   automatically.

6. **Custom stages** (if needed) — Subclass `PipelineStage` in
   `fastvideo/pipelines/stages/`, implement `forward()`, optionally
   implement `verify_input()`/`verify_output()`.

7. **Verify** — Run `fastvideo generate --config <config.yaml>` with a
   minimal nested config to confirm the pipeline loads and generates
   output.