Pipelines, runners, and applications

FlashDreams model integrations use these public layers:

• Pipelines (StreamInferencePipelineConfig) that define model behavior.

• Runners (RunnerConfig + Runner) that define CLI-facing I/O.

• V2 applications (IApplication) that bind reusable demo infrastructure directly to pipeline configs.

Most actively developed model implementations now live under integrations/* as plugin-style standalone packages. This page keeps documenting the in-tree pipeline modules that are still exposed from flashdreams.recipes.

Note

Pipeline modules import the heavy GPU stack (transformer-engine, CUDA ops) at import time, so this page shows them by automodule with :no-undoc-members: to keep the rendered API focused on the names that these in-tree modules actually expose. The unified flashdreams-run CLI shows end-to-end usage; see Models for model launch examples.

Integration structure (current)

V2 demo ports follow integrations_v2/<name>/:

• apps/<demo>/adapter.py: create_app() -> IApplication binding.

• config.py: the model’s unique pipeline config literal or model-specific StreamInferencePipelineConfig wrapper.

• impl/: all model-specific implementation.

• tests/: model-specific tests, when needed.

• apps/<demo>/README.md: launch instructions for that demo.

• pyproject.toml: packaging plus flashdreams.applications_v2 entry points.

Apart from that unique config, the integration root contains no implementation modules. These packages do not add runner.py or flashdreams.runner_configs just to launch a v2 demo. Older plugin-style integrations under integrations/ may still expose runners through flashdreams-run.

The default application entry point uses <demo-slug>-<model-slug> and create_app. Additional compatible configurations may use <demo-slug>-<model-slug>-<suffix> and the matching create_app_<suffix> factory, such as a -fast entry point backed by create_app_fast.

Reference integration folders

• omnidreams

• self_forcing

• causal_forcing

• lingbot

• hy_worldplay

• wan21

• fastvideo_causal_wan22

• flashvsr

• cosmos_predict2

NVIDIA OmniDreams

OmniDreams ships under integrations_v2/omnidreams. Its model code lives under impl/ and its application adapter binds the reusable interactive-drive package under apps/ through the flashdreams.applications_v2 entry-point group. See integrations_v2/omnidreams/README.md for launch details.

Wan

Public Wan integration surface for integration plugins.

class Wan21Transformer(*args: Any, **kwargs: Any)

Bases: Transformer[Wan21TransformerCache]

Wan 2.1 DiT adapted to the infra Transformer interface.

property latent_shape: tuple[int, ...]

Per-rank post-patchify latent shape [*batch_shape, L/cp, D].

Wan flattens THW into one token axis and shards across the THW CP group. D = network.out_dim * prod(patch_size) is the noise channel count; the mask / image-latent channels added by concat_image_mask_to_latent come from input in predict_flow, not from the noise tensor.

Per-rollout (height, width) is populated by initialize_autoregressive_cache(); reading earlier asserts.

initialize_autoregressive_cache(*, height: int, width: int, text_embeddings: Tensor, image_embeddings: Tensor | None = None, negative_text_embeddings: Tensor | None = None, **_unused: Any) → Wan21TransformerCache

Build a seeded transformer cache for a new rollout.

I2V state is not baked into the cache; the latent + injection mask are passed per AR step as the input argument to predict_flow / postprocess_clean_latent.

Parameters:

• height – Pre-patchify latent height (post-VAE).

• width – Pre-patchify latent width (post-VAE).

• text_embeddings – Conditional UMT5 embeddings [..., text_len, text_dim].

• image_embeddings – Conditional CLIP image embeddings (only used by networks with cross_attn_enable_img=True). Shared with the uncond branch.

• negative_text_embeddings – Negative-prompt embeddings. Required iff config.guidance_scale > 1.0; must be None otherwise.

Returns:

Populated cache. network_cache_uncond is None iff CFG is disabled.

predict_flow(noisy_latent: torch.Tensor, timestep: torch.Tensor, cache: Wan21TransformerCache, input: I2VCtrl | None = None, network_extra_kwargs: dict[str, Any] | None = None) → torch.Tensor

Predict the flow for one denoising step.

timestep may be a scalar / per-batch tensor (standard Wan 2.1 / 14B path) or a per-token tensor with the same trailing token axis as noisy_latent (Wan 2.2 TI2V 5B first-frame seeding at AR step 0). The per-token layout flows through WanDiTNetwork.forward(), which dispatches the sinusoidal embedding + AdaLN modulation on the native shape.

CUDA-graph capture is shape-sensitive: the captured replay region only sees AR step >= self._cuda_graph_capture_ar_idx (steady state). TI2V 5B is configured with len_t == window_size_t so the threshold lands at AR 1, putting the per-token AR-0 step inside the eager .drain branch where shape changes are safe. After AR 0 the pipeline switches back to scalar timesteps, so the captured branch sees a single stable shape across all AR steps it owns.

postprocess_clean_latent(clean_latent: torch.Tensor, cache: Wan21TransformerCache, input: I2VCtrl | None = None) → torch.Tensor

Re-stamp x0 masked positions with the image latent (mask-inject I2V only).

T2V and the channel-concat I2V mode fall through unchanged.

patchify_and_maybe_split_cp(x: torch.Tensor) → torch.Tensor

patchify_and_maybe_split_cp(x: I2VCtrl) → I2VCtrl

Patchify and CP-split a noisy latent or an I2V control payload.

Tensors delegate to the network helper; I2V payloads patchify the latent and mask fields independently so the per-field channel layouts are preserved for the mask-inject blend downstream.

unpatchify_and_maybe_gather_cp(x: torch.Tensor) → torch.Tensor

Inverse of patchify_and_maybe_split_cp for the network output.

class Wan21TransformerConfig(*, _target: type[~flashdreams.recipes.wan.transformer.wan21.Wan21Transformer] = <factory>, network: ~flashdreams.recipes.wan.transformer.impl.network.WanDiTNetworkConfig = <factory>, dtype: torch.dtype = torch.bfloat16, checkpoint_path: str | None = None, checkpoint_min_free_gb: float | None = None, state_dict_transform: ~collections.abc.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None, stream_checkpoint: bool = False, init_device: str | None = None, batch_shape: tuple[int, ...] = (1,), len_t: int = 21, guidance_scale: float = 1.0, window_size_t: int = 21, sink_size_t: int = 0, h_extrapolation_ratio: float = 1.0, w_extrapolation_ratio: float = 1.0, compile_network: bool = True, use_cuda_graph: bool = True, cuda_graph_warmup_iters: int = 2, stamp_image_latent: bool = False, concat_image_mask_to_latent: bool = False, ti2v_first_frame_per_token_timestep: bool = False, first_frame_timestep_value: float = 0.0)

Bases: TransformerConfig

Config for the Wan 2.1 transformer.

Bakes in the temporal layout (len_t, window_size_t, optional sink_size_t) and the CFG / compile knobs. Per-rollout spatial layout (height, width) is supplied to Wan21Transformer.initialize_autoregressive_cache() so one instance can serve multiple resolutions. Wan flattens T*H*W into one token axis and shards it across the THW CP group; the CP size is auto-detected from torch.distributed.get_world_size() at construction time, so the launcher (torchrun --nproc_per_node=N) is the single source of truth.

The two I2V flags are independent and composable:

• stamp_image_latent: overwrite the noisy latent with the clean image latent at masked positions every denoising step, and re-stamp the predicted x0 the same way. network.in_dim unchanged. (flashdreams mask-inject integration; used by the out-of-tree causal_forcing plugin.)

• concat_image_mask_to_latent: append the 4-channel mask and 16-channel image latent along the channel dim. Builders that set this flag must also set network.in_dim = 16 + 4 + 16 to match the official Wan 2.1 14B I2V layout.

With both enabled, the stamp runs first and the result is then concatenated with the mask + image latent.

checkpoint_min_free_gb: float | None = None

Optional first-run disk preflight for checkpoint downloads.

None uses the generic cache reserve. Larger model integrations can set this to their documented storage requirement; users can still override it with FLASHDREAMS_MIN_CACHE_FREE_GB.

state_dict_transform: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None

Pre-load state-dict remap (e.g. Self-Forcing’s generator_ema.model.… layout).

stream_checkpoint: bool = False

Load cached safetensors directly into the model with bounded host residency.

init_device: str | None = None

Optional device used for initial network parameter allocation.

Large streaming-checkpoint models can set this to the final runtime device so the module is not first materialized as fp32 CPU tensors.

batch_shape: tuple[int, ...] = (1,)

Batch dims of the latent (excluding the L, D dims).

len_t: int = 21

Latent frames per AR chunk (post-VAE).

guidance_scale: float = 1.0

flow = uncond + s * (cond - uncond). 1.0 disables CFG; > 1.0 requires negative-text embeddings at cache build time.

Type:

CFG scale s

window_size_t: int = 21

Self-attention sliding-window size (pre-patchify T frames).

sink_size_t: int = 0

Prefix sink size (pre-patchify T frames) for self-attention KV cache.

compile_network: bool = True

torch.compile the network on init.

use_cuda_graph: bool = True

Wrap the network in CUDAGraphWrapper for steady-state replay. Caller must keep non-staged inputs at stable storage addresses across calls. predict_flow dispatches to wrapper.drain while the KV cache is still filling and to wrapper once it reaches steady state.

cuda_graph_warmup_iters: int = 2

Eager calls before capture (>= 2 to drain Inductor autotune).

stamp_image_latent: bool = False

See class docstring (mask-inject I2V integration).

concat_image_mask_to_latent: bool = False

See class docstring (channel-concat I2V layout).

ti2v_first_frame_per_token_timestep: bool = False

Wan 2.2 TI2V 5B first-frame conditioning. When True and an I2VCtrl input is provided at AR step 0, predict_flow rewrites the scheduler’s scalar timestep into a per-token tensor: t = first_frame_timestep_value at positions marked by the I2V mask (i.e. the first-frame latent), and the scheduler’s t elsewhere. AR steps >= 1 continue to use the scalar timestep, which keeps the CUDA-graph-captured replay branch on a single stable input shape.

Composes with stamp_image_latent: together they implement the upstream Wan 2.2 5B “VAE-seeded first-frame + per-token t=0” TI2V recipe – the latent is stamped clean every denoising step while the network sees t=0 for those tokens. The standard mask-inject I2V recipe leaves this flag off and relies on the classifier-free stamp alone.

first_frame_timestep_value: float = 0.0

Per-token timestep assigned to first-frame conditioning tokens when ti2v_first_frame_per_token_timestep is True.

Defaults to 0.0 (Wan 2.2 TI2V 5B’s base recipe — treats the first frame as fully clean by AdaLN). HY-WorldPlay’s distilled WAN-5B raises it to 14.0 (vendor’s stabilization_level - 1) so the AdaLN table sees a small nonzero sigma at the first frame.

Unused when ti2v_first_frame_per_token_timestep is False.

class Wan22TI2V5BVAEDecoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B/resolve/main/Wan2.2_VAE.pth', dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = False, base_dim: int = 160, decoder_base_dim: int | None = 256, z_dim: int = 48, patch_size: int = 2, is_residual: bool = True, latent_mean: tuple[float, ...] = (-0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.157, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.123, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.052, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667), latent_std: tuple[float, ...] = (0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.499, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.06, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744), state_dict_transform: ~typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None)

Bases: WanVAEDecoderConfig

Pre-rolled config for the Wan 2.2 TI2V 5B decoder.

Mirrors Wan22TI2V5BVAEEncoderConfig but with the asymmetric decoder_base_dim=256.

decoder_base_dim: int | None = 256

Decoder base channel count. None mirrors base_dim (Wan 2.1). Wan 2.2 TI2V 5B uses an asymmetric 256.

state_dict_transform: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None

Optional pre-load_state_dict key remap. See wan22_ti2v_5b_vae_state_dict_transform() for the Wan 2.2 TI2V 5B diffusers -> flashdreams remap.

class Wan22TI2V5BVAEEncoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B/resolve/main/Wan2.2_VAE.pth', dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = False, base_dim: int = 160, z_dim: int = 48, patch_size: int = 2, is_residual: bool = True, latent_mean: tuple[float, ...] = (-0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.157, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.123, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.052, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667), latent_std: tuple[float, ...] = (0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.499, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.06, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744), state_dict_transform: ~typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None)

Bases: WanVAEEncoderConfig

Pre-rolled config for the Wan 2.2 TI2V 5B encoder.

Pins upstream’s native Wan2.2_VAE.pth (loaded with no key remap), the 16x-spatial / 48ch / residual / patchify architecture knobs. Equivalent to the Wan 2.1 encoder config plus the 5B-specific knobs flipped on. To load the diffusers shard instead, override checkpoint_path=WAN22_TI2V_5B_VAE_DIFFUSERS_PATH and state_dict_transform=wan22_ti2v_5b_vae_state_dict_transform.

base_dim: int = 160

Encoder base channel count (WanVAE dim). 96 for Wan 2.1, 160 for Wan 2.2 TI2V 5B.

z_dim: int = 48

Latent channels. 16 for Wan 2.1, 48 for Wan 2.2 TI2V 5B.

patch_size: int = 2

Outer spatial pixel-shuffle factor (1 = no patchify; 2 for Wan 2.2 TI2V 5B).

is_residual: bool = True

Use ResidualDownBlock (Wan 2.2) instead of the legacy ResidualBlock + AttentionBlock down-stage (Wan 2.1).

latent_mean: tuple[float, ...] = (-0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.157, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.123, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.052, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667)

Per-channel latent mean used for normalisation; must match z_dim entries.

latent_std: tuple[float, ...] = (0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.499, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.06, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744)

Per-channel latent std used for normalisation.

state_dict_transform: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None

Optional pre-load_state_dict key remap (e.g. diffusers -> flashdreams layout). See wan22_ti2v_5b_vae_state_dict_transform() for the Wan 2.2 TI2V 5B remap.

class Wan22Transformer(*args: Any, **kwargs: Any)

Bases: Transformer[Wan22TransformerCache]

Wan 2.2 dual-network DiT.

predict_flow dispatches to the branch selected by the timestep; finalize_kv_cache re-runs both branches once at the context noise so neither KV cache lags behind.

property latent_shape: tuple[int, ...]

Per-rank latent shape (both branches share this, asserted by config).

initialize_autoregressive_cache(*, height: int, width: int, text_embeddings: Tensor, image_embeddings: Tensor | None = None, **_unused: Any) → Wan22TransformerCache

Build a seeded transformer cache for a new rollout.

Both branches see the same text/image conditioning and the same per-rollout spatial layout, matching upstream Wan 2.2. CFG is not supported here.

predict_flow(noisy_latent: torch.Tensor, timestep: torch.Tensor, cache: Wan22TransformerCache, input: Any = None) → torch.Tensor

Predict the flow at timestep.

Parameters:

• noisy_latent – Patchified noisy latent for this denoising step.

• timestep – Scalar timestep tensor.

• cache – Per-rollout AR cache.

• input – Patchified encoder output for this AR step, or None when the pipeline has no encoder. Subclasses should narrow the type to their encoder’s output type.

Returns:

Predicted flow tensor with the same shape as noisy_latent.

finalize_kv_cache(noisy_latent: torch.Tensor, timestep: torch.Tensor, cache: Wan22TransformerCache, input: Any = None) → None

Refresh both networks’ KV caches at the context-noise step.

Each Wan 2.2 denoising step only touches one branch, so the other lags by the end of the loop. Re-running both at context_noise keeps them in lock-step; flow outputs are discarded.

patchify_and_maybe_split_cp(x: torch.Tensor) → torch.Tensor

Patchify and (optionally) CP-split a noisy latent or encoder payload.

Tensors patchify and split. Structured payloads (e.g. an image-control struct with latent + mask) patchify each tensor field and return the same struct type. Implement as identity when neither token packing nor CP sharding applies. Output preserves the input Python type — only shapes change.

unpatchify_and_maybe_gather_cp(x: torch.Tensor) → torch.Tensor

Inverse of patchify_and_maybe_split_cp for the network output.

class Wan22TransformerConfig(*, _target: type[Wan22Transformer] = <factory>, transformer_high_noise: Wan21TransformerConfig = <factory>, transformer_low_noise: Wan21TransformerConfig = <factory>, boundary_ratio: float = 0.875, num_train_timesteps: int = 1000)

Bases: TransformerConfig

Config for the Wan 2.2 dual-network transformer.

Wan 2.2 dispatches to one of two Wan 2.1 networks based on whether the timestep is above or below boundary_ratio * num_train_timesteps. Both branches must agree on patch_size / in_dim / dim / batch_shape / len_t / guidance_scale (asserted in __post_init__). The per-rollout spatial layout (height, width) is supplied to Wan22Transformer.initialize_autoregressive_cache() and forwarded to both branches. The CP size is auto-detected from torch.distributed.get_world_size(), same as Wan 2.1. Wan 2.2 has no CFG or I2V support here.

transformer_high_noise: Wan21TransformerConfig

Sub-config for the high-noise branch (timestep > boundary).

transformer_low_noise: Wan21TransformerConfig

Sub-config for the low-noise branch (timestep <= boundary).

boundary_ratio: float = 0.875

Fraction of num_train_timesteps separating the two branches. Default 0.875 matches upstream Wan 2.2.

property boundary_timestep: float

Absolute timestep boundary.

class WanDiTNetwork(*args: Any, **kwargs: Any)

Bases: Module

WAN diffusion backbone for text-to-video and image-to-video.

set_context_parallel_group(cp_group: ProcessGroup | None = None) → None

Set context-parallel process group for all blocks.

This must be called before initialize_cache when CP is used.

patchify_and_maybe_split_cp(x: Tensor, process_groups: list[ProcessGroup | None] | None = None, cp_dims: list[int | None] | None = None) → Tensor

Patchify and optionally CP-split the input video tensor.

The patchify pattern is ... (t kt) c (h kh) (w kw) -> ... (t h w) (c kt kh kw).

Returns:

Patched tensor with shape [..., L, D] where L = T * H * W / (kt * kh * kw).

unpatchify_and_maybe_gather_cp(pH: int, pW: int, x: Tensor, process_groups: list[ProcessGroup | None] | None = None, cp_dims: list[int | None] | None = None) → Tensor

Unpatchify and optionally CP-gather the tensor back to video shape.

The unpatchify pattern is ... (t h w) (c kt kh kw) -> ... (t kt) c (h kh) (w kw).

Returns:

Unpatched tensor with shape [..., T, C, H, W].

initialize_cache(chunk_size: int, window_size: int, sink_size: int, text_embeddings: Tensor, img_embeddings: Tensor | None = None) → WanDiTNetworkCache

Initialize block caches from text/image context embeddings.

Parameters:

• chunk_size – Number of tokens appended per self-attention update.

• window_size – Rolling-window size in tokens for self-attention cache.

• sink_size – Sink-token capacity preserved across updates.

• text_embeddings – Text embeddings. UMT5 has shape […, 512, 4096].

• img_embeddings – Optional image embeddings for I2V. CLIP has shape […, 256, 1280].

Returns:

WanDiTNetworkCache containing per-block caches.

update_parameters_after_loading_checkpoint() → None

Fuse load-time-known ops into weights; call once after loading the checkpoint.

forward(x: torch.Tensor, timesteps: torch.Tensor, cache: WanDiTNetworkCache, rope_freqs: torch.Tensor, current_chunk_idx: int = 0, eager_mode: bool = True, block_extra_kwargs: dict[str, Any] = {}) → torch.Tensor

Run one denoising forward pass.

Parameters:

• x – Input tokens after patchify + CP, shape [..., L, D_in]; layout "... (t h w) (c kt kh kw)".

• timesteps –

  Diffusion timesteps. Two layouts are supported:

  • Scalar (per-batch). Shape broadcastable to [...] (i.e., to x.shape[:-2]). The same timestep is shared across every token, matching the standard Wan 2.1 / Wan 2.2 14B chunked-denoise path.

  • Per-token. Shape [..., L] matching x’s post- patchify token axis. Used by Wan 2.2 TI2V 5B at AR step 0 to stamp t=0 at the first-frame conditioning tokens while the rest of the chunk denoises at the current scheduler step. See Wan21Transformer.predict_flow for the higher-level entry point.

• cache – Network KV caches.

• rope_freqs – Full-width RoPE frequencies after CP. Standard mode uses current-chunk frequencies with shape [L, 1, 1, d]; KV-cache-relative mode uses frequencies relative to the KV cache.

• current_chunk_idx – Current chunk index for streaming cache update.

• eager_mode – If True, run cache before/after update hooks.

• block_extra_kwargs – Extra kwargs forwarded to each block.

Returns:

Network output, shape [..., L, prod(patch_size) * out_dim].

class WanDiTNetwork1pt3BConfig(_target: type[WanDiTNetwork] = <factory>, patch_size: tuple[int, int, int]=(1, 2, 2), text_len: int = 512, in_dim: int = 16, dim: int = 1536, ffn_dim: int = 8960, freq_dim: int = 256, text_dim: int = 4096, out_dim: int = 16, num_heads: int = 12, num_layers: int = 30, cross_attn_norm: bool = True, cross_attn_enable_img: bool = False, eps: float = 1e-06, concat_padding_mask: bool = False, patch_embedding_type: Literal['linear', 'conv3d']='conv3d', apply_rope_before_kvcache: bool = True, cp_method: Literal['ring', 'ulysses']='ring')

Bases: WanDiTNetworkConfig

Configuration for the 1.3B Wan DiT network.

dim: int = 1536

Transformer hidden size (width).

ffn_dim: int = 8960

Feed-forward hidden dimension.

num_heads: int = 12

Number of attention heads.

num_layers: int = 30

Number of transformer blocks.

class WanDiTNetwork14BConfig(_target: type[WanDiTNetwork] = <factory>, patch_size: tuple[int, int, int]=(1, 2, 2), text_len: int = 512, in_dim: int = 16, dim: int = 5120, ffn_dim: int = 13824, freq_dim: int = 256, text_dim: int = 4096, out_dim: int = 16, num_heads: int = 40, num_layers: int = 40, cross_attn_norm: bool = True, cross_attn_enable_img: bool = False, eps: float = 1e-06, concat_padding_mask: bool = False, patch_embedding_type: Literal['linear', 'conv3d']='conv3d', apply_rope_before_kvcache: bool = True, cp_method: Literal['ring', 'ulysses']='ring')

Bases: WanDiTNetworkConfig

Configuration for the 14B Wan DiT network.

dim: int = 5120

Transformer hidden size (width).

ffn_dim: int = 13824

Feed-forward hidden dimension.

num_heads: int = 40

Number of attention heads.

num_layers: int = 40

Number of transformer blocks.

class WanDiTNetworkConfig(_target: type[WanDiTNetwork] = <factory>, patch_size: tuple[int, int, int]=(1, 2, 2), text_len: int = 512, in_dim: int = 16, dim: int = 1536, ffn_dim: int = 8960, freq_dim: int = 256, text_dim: int = 4096, out_dim: int = 16, num_heads: int = 12, num_layers: int = 30, cross_attn_norm: bool = True, cross_attn_enable_img: bool = False, eps: float = 1e-06, concat_padding_mask: bool = False, patch_embedding_type: Literal['linear', 'conv3d']='conv3d', apply_rope_before_kvcache: bool = True, cp_method: Literal['ring', 'ulysses']='ring')

Bases: InstantiateConfig

Configuration for the Wan DiT network.

patch_size: tuple[int, int, int] = (1, 2, 2)

Patch size for the input tensor.

text_len: int = 512

Maximum text token length.

in_dim: int = 16

Number of input latent channels before patch embedding.

dim: int = 1536

Transformer hidden size (width).

ffn_dim: int = 8960

Feed-forward hidden dimension.

freq_dim: int = 256

Sinusoidal timestep embedding dimension.

text_dim: int = 4096

Text encoder output dimension.

out_dim: int = 16

Output latent channels after the head.

num_heads: int = 12

Number of attention heads.

num_layers: int = 30

Number of transformer blocks.

cross_attn_norm: bool = True

If True, apply LayerNorm before cross-attention.

cross_attn_enable_img: bool = False

If True, build image cross-attention and CLIP image projection (I2V).

eps: float = 1e-06

Epsilon for normalization layers.

concat_padding_mask: bool = False

If True, concatenate one mask channel into the input channels.

patch_embedding_type: Literal['linear', 'conv3d'] = 'conv3d'

"linear" (flattened patch MLP) or "conv3d" (strided conv).

Type:

Type of patch embedding

apply_rope_before_kvcache: bool = True

If True, apply RoPE to keys before storing them in the KV cache.

cp_method: Literal['ring', 'ulysses'] = 'ring'

Context-parallel attention method for transformer attention ops.

class WanDiTNetworkTI2V5BConfig(_target: type[WanDiTNetwork] = <factory>, patch_size: tuple[int, int, int]=(1, 2, 2), text_len: int = 512, in_dim: int = 48, dim: int = 3072, ffn_dim: int = 14336, freq_dim: int = 256, text_dim: int = 4096, out_dim: int = 48, num_heads: int = 24, num_layers: int = 30, cross_attn_norm: bool = True, cross_attn_enable_img: bool = False, eps: float = 1e-06, concat_padding_mask: bool = False, patch_embedding_type: Literal['linear', 'conv3d']='conv3d', apply_rope_before_kvcache: bool = True, cp_method: Literal['ring', 'ulysses']='ring')

Bases: WanDiTNetworkConfig

Configuration for the Wan 2.2 TI2V 5B DiT network.

Mirrors the official Wan-AI/Wan2.2-TI2V-5B-Diffusers/transformer config: 24 heads * 128 head_dim = 3072 inner dim, 30 layers, ffn_dim 14336, and 48-channel latent in/out (the matching 16x VAE in vae.py outputs 48 channels). Unlike Wan 2.1 14B I2V, TI2V 5B has no CLIP cross-attention branch (cross_attn_enable_img=False): the first frame is conditioned via a clean VAE-latent seed plus a per-token t=0 timestep on the AR-step-0 first-frame tokens, not via CLIP image features.

in_dim: int = 48

Number of input latent channels before patch embedding.

out_dim: int = 48

Output latent channels after the head.

dim: int = 3072

Transformer hidden size (width).

ffn_dim: int = 14336

Feed-forward hidden dimension.

num_heads: int = 24

Number of attention heads.

num_layers: int = 30

Number of transformer blocks.

cross_attn_enable_img: bool = False

If True, build image cross-attention and CLIP image projection (I2V).

class WanI2VCtrlEncoderConfig(*, _target: type[I2VCtrlEncoder] = <factory>, encoder: WanVAEEncoderConfig = <factory>)

Bases: EncoderConfig

Config for the I2V control encoder.

encoder: WanVAEEncoderConfig

Streaming Wan VAE encoder. Pin its checkpoint to the decoder’s so the encoded latent matches the network’s input distribution.

class WanInferencePipeline(*args: Any, **kwargs: Any)

Bases: StreamInferencePipeline[WanVAECache, Wan21TransformerCache | Wan22TransformerCache, WanVAECache]

Wan 2.1 / 2.2 inference pipeline, T2V and I2V.

T2V and I2V share the same rollout loop; the difference is whether you pass an image to initialize_cache. The pipeline config’s encoder slot must agree (None for T2V, an I2V config for I2V).

Example:

pipeline: WanInferencePipeline = ...

# T2V: pass latent ``height`` and ``width``.
cache = pipeline.initialize_cache(
    text=["A cat surfing."], height=60, width=104
)
_chunk = pipeline.generate(0, cache)
pipeline.finalize(0, cache)

# I2V: pass the first frame and let sizes derive from it.
_i2v_cache = pipeline.initialize_cache(
    text=["A cat surfing."], image=first_frame
)

initialize_cache(text: list[str], image: Tensor | None = None, *, height: int | None = None, width: int | None = None, release_oneshot_encoders: bool = True) → WanInferencePipelineCache

Initialize the per-rollout cache for a batch of prompts.

Parameters:

• text – One prompt per batch element. Length must match the transformer’s batch_shape.

• image – First-frame pixels of shape [*batch_shape, 1, 3, H, W] in [-1, 1]. Required for I2V (self.encoder is set), forbidden for T2V. H and W must equal height * decoder.spatial_compression_ratio and width * decoder.spatial_compression_ratio, respectively.

• height – Pre-patchify latent height (post-VAE). Optional for I2V — derived from image when omitted; required for T2V.

• width – Pre-patchify latent width (post-VAE). Same rules as height.

• release_oneshot_encoders – Free the text and image encoders after the cache is initialized. Later calls reload them from self.config before encoding new prompts/images.

Returns:

Cache to thread through generate / finalize.

release_oneshot_encoders() → None

Free the per-rollout text and first-frame image encoders.

Idempotent. A later initialize_cache() call can reload the encoders from self.config before encoding raw prompts/images.

generate(autoregressive_index: int, cache: WanInferencePipelineCache) → torch.Tensor

Generate one decoded video chunk.

Parameters:

• autoregressive_index – AR step index, starting at 0.

• cache – Per-rollout cache from initialize_cache.

Returns:

Decoded video of shape [*batch_shape, T, C, H, W] in [-1, 1].

get_num_input_frames(autoregressive_index: int) → int

Number of input video frames the model expects at this AR step.

get_num_output_frames(autoregressive_index: int) → int

Number of decoded video frames produced at this AR step.

class WanInferencePipelineCache(*, transformer_cache: TransformerCacheT, encoder_cache: StreamingEncoderCacheT | None = None, decoder_cache: StreamingDecoderCacheT | None = None, final_state: DiffusionModel.FinalState[TransformerCacheT] | None = None, autoregressive_index: int | None = None, event_profiler: EventProfiler | None = None, image: Tensor | None = None)

Bases: StreamInferencePipelineCache[WanVAECache, Wan21TransformerCache | Wan22TransformerCache, WanVAECache]

Per-rollout state for the Wan pipeline.

Adds the I2V first-frame pixels on top of the inherited caches. Pixel-to- latent encoding happens per AR step inside the encoder, not here.

image: Tensor | None = None

First-frame pixels [*batch_shape, 1, 3, H, W] in [-1, 1]; None for T2V.

class WanInferencePipelineConfig(*, _target: type['WanInferencePipeline'] = <factory>, name: str, diffusion_model: DiffusionModelConfig, decoder: DecoderConfig | None = None, encoder: EncoderConfig | None = None, enable_sync_and_profile: bool = False, text_encoder: UMT5TextEncoderConfig | None = <factory>, image_encoder: CLIPImageEncoderConfig | None = None)

Bases: StreamInferencePipelineConfig

Config for the Wan inference pipeline.

T2V vs I2V is selected by the inherited encoder slot: None for T2V, an I2V control-encoder config for I2V.

text_encoder: UMT5TextEncoderConfig | None

UMT5 text encoder run once per rollout.

image_encoder: CLIPImageEncoderConfig | None = None

CLIP image encoder for I2V variants trained with cross_attn_enable_img=True (Wan 2.1 14B I2V). None skips CLIP cross-attention conditioning.

class WanVAEDecoder(*args: Any, **kwargs: Any)

Bases: StreamingVideoDecoder[WanVAECache]

Wan VAE decoder.

Forward input is a latent [..., Tl, Cl, Hl, Wl]; output is a video tensor [..., T, C, H, W] in [-1, 1]. The cache is advanced in-place across AR decode steps; passing cache=None allocates a fresh single-shot cache.

initialize_autoregressive_cache() → WanVAECache

Build a fresh per-rollout cache.

Override to return the decoder’s concrete cache type.

property temporal_compression_ratio: int

Pixel frames ÷ latent frames in steady state (AR ≥ 1).

AR 0 typically yields fewer pixel frames per latent frame because of causal first-frame padding; that asymmetry lives inside get_output_temporal_size() / get_input_temporal_size().

property spatial_compression_ratio: int

Pixel side ÷ latent side. Constant across AR steps.

get_output_temporal_size(autoregressive_index: int, input_temporal_size: int) → int

Causal: AR 0 first latent frame decodes to a single pixel frame.

get_input_temporal_size(autoregressive_index: int, output_temporal_size: int) → int

Latent frame count needed to produce output_temporal_size pixels.

Inverse of get_output_temporal_size(). Implementations should assert output_temporal_size is achievable at this AR step (i.e. divisible by the right ratio after subtracting any causal padding).

Parameters:

• autoregressive_index – AR step index (0-based).

• output_temporal_size – Desired number of pixel frames.

Returns:

Number of latent frames needed at this step.

class WanVAEDecoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth', dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = False, base_dim: int = 96, decoder_base_dim: int | None = None, z_dim: int = 16, patch_size: int = 1, is_residual: bool = False, latent_mean: tuple[float, ...] = (-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921), latent_std: tuple[float, ...] = (2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916), state_dict_transform: ~typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None)

Bases: DecoderConfig

Config for the Wan VAE decoder.

Defaults reproduce the Wan 2.1 / 14B 8x-spatial 16-channel streaming VAE. Override the architecture knobs (or use Wan22TI2V5BVAEDecoderConfig) to load the Wan 2.2 TI2V 5B decoder, which has asymmetric base_dim=160 / decoder_base_dim =256 and the residual up-stage with DupUp3D shortcut.

use_cuda_graph: bool = True

Wrap the decoder forward in a CUDA graph for replay.

use_compile: bool = False

torch.compile(mode="max-autotune-no-cudagraphs"). See WanVAEEncoderConfig.use_compile for the VRAM caveat.

decoder_base_dim: int | None = None

Decoder base channel count. None mirrors base_dim (Wan 2.1). Wan 2.2 TI2V 5B uses an asymmetric 256.

state_dict_transform: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None

Optional pre-load_state_dict key remap. See wan22_ti2v_5b_vae_state_dict_transform() for the Wan 2.2 TI2V 5B diffusers -> flashdreams remap.

class WanVAEEncoder(*args: Any, **kwargs: Any)

Bases: StreamingVideoEncoder[WanVAECache]

Wan VAE encoder.

Forward input is a video tensor [..., T, C, H, W] in [-1, 1]; output is the latent [..., Tl, Cl, Hl, Wl]. The cache is advanced in-place across AR encode steps; passing cache=None allocates a fresh single-shot cache.

initialize_autoregressive_cache() → WanVAECache

Build a fresh per-rollout cache.

Override to return the encoder’s concrete cache type.

property temporal_compression_ratio: int

Pixel frames ÷ latent frames in steady state (AR ≥ 1).

AR 0 typically takes one extra (un-grouped) pixel frame to produce its first latent frame because of causal first-frame padding; that asymmetry lives inside get_output_temporal_size() / get_input_temporal_size().

property spatial_compression_ratio: int

Pixel side ÷ latent side. Constant across AR steps.

get_output_temporal_size(autoregressive_index: int, input_temporal_size: int) → int

Causal: AR 0 needs an extra (un-grouped) pixel frame for the first latent.

get_input_temporal_size(autoregressive_index: int, output_temporal_size: int) → int

Pixel frame count needed to produce output_temporal_size latents.

Inverse of get_output_temporal_size(). Implementations should assert output_temporal_size is achievable at this AR step (i.e. the corresponding pixel count comes out as a positive integer).

Parameters:

• autoregressive_index – AR step index (0-based).

• output_temporal_size – Desired number of latent frames.

Returns:

Number of pixel frames needed at this step.

class WanVAEEncoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth', dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = False, base_dim: int = 96, z_dim: int = 16, patch_size: int = 1, is_residual: bool = False, latent_mean: tuple[float, ...] = (-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921), latent_std: tuple[float, ...] = (2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916), state_dict_transform: ~typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None)

Bases: EncoderConfig

Config for the Wan VAE encoder.

Defaults reproduce the Wan 2.1 / 14B 8x-spatial 16-channel streaming VAE. Override base_dim / z_dim / patch_size / is_residual (and the latent mean/std) to load the Wan 2.2 TI2V 5B 16x-spatial 48-channel residual VAE; see Wan22TI2V5BVAEEncoderConfig for the pre-rolled set.

use_cuda_graph: bool = True

Wrap the encoder forward in a CUDA graph for replay.

use_compile: bool = False

torch.compile(mode="max-autotune-no-cudagraphs"). Off by default: Inductor autotune workspaces can add several GiB of transient VRAM per unique input shape, surfacing as ‘illegal memory access’ on smaller GPUs with the full-channel vae checkpoint.

base_dim: int = 96

Encoder base channel count (WanVAE dim). 96 for Wan 2.1, 160 for Wan 2.2 TI2V 5B.

z_dim: int = 16

Latent channels. 16 for Wan 2.1, 48 for Wan 2.2 TI2V 5B.

patch_size: int = 1

Outer spatial pixel-shuffle factor (1 = no patchify; 2 for Wan 2.2 TI2V 5B).

is_residual: bool = False

Use ResidualDownBlock (Wan 2.2) instead of the legacy ResidualBlock + AttentionBlock down-stage (Wan 2.1).

latent_mean: tuple[float, ...] = (-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921)

Per-channel latent mean used for normalisation; must match z_dim entries.

latent_std: tuple[float, ...] = (2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916)

Per-channel latent std used for normalisation.

state_dict_transform: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None

Optional pre-load_state_dict key remap (e.g. diffusers -> flashdreams layout). See wan22_ti2v_5b_vae_state_dict_transform() for the Wan 2.2 TI2V 5B remap.

wan22_ti2v_5b_vae_state_dict_transform(state_dict: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor]

Remap a diffusers AutoencoderKLWan state-dict to WanVAE keys.

Opt-in for callers who point Wan22TI2V5BVAEEncoderConfig / Wan22TI2V5BVAEDecoderConfig at the diffusers safetensors shard (WAN22_TI2V_5B_VAE_DIFFUSERS_PATH) instead of the native .pth default. Renames keys only – no tensors are copied or reshaped. Unmatched keys pass through and surface as unexpected_keys on load.

Unified Wan inference pipeline (Wan 2.1 / Wan 2.2, T2V and I2V).

class WanInferencePipelineCache(*, transformer_cache: TransformerCacheT, encoder_cache: StreamingEncoderCacheT | None = None, decoder_cache: StreamingDecoderCacheT | None = None, final_state: DiffusionModel.FinalState[TransformerCacheT] | None = None, autoregressive_index: int | None = None, event_profiler: EventProfiler | None = None, image: Tensor | None = None)

Bases: StreamInferencePipelineCache[WanVAECache, Wan21TransformerCache | Wan22TransformerCache, WanVAECache]

Per-rollout state for the Wan pipeline.

Adds the I2V first-frame pixels on top of the inherited caches. Pixel-to- latent encoding happens per AR step inside the encoder, not here.

image: Tensor | None = None

First-frame pixels [*batch_shape, 1, 3, H, W] in [-1, 1]; None for T2V.

class WanInferencePipelineConfig(*, _target: type['WanInferencePipeline'] = <factory>, name: str, diffusion_model: DiffusionModelConfig, decoder: DecoderConfig | None = None, encoder: EncoderConfig | None = None, enable_sync_and_profile: bool = False, text_encoder: UMT5TextEncoderConfig | None = <factory>, image_encoder: CLIPImageEncoderConfig | None = None)

Bases: StreamInferencePipelineConfig

Config for the Wan inference pipeline.

T2V vs I2V is selected by the inherited encoder slot: None for T2V, an I2V control-encoder config for I2V.

text_encoder: UMT5TextEncoderConfig | None

UMT5 text encoder run once per rollout.

image_encoder: CLIPImageEncoderConfig | None = None

CLIP image encoder for I2V variants trained with cross_attn_enable_img=True (Wan 2.1 14B I2V). None skips CLIP cross-attention conditioning.

class WanInferencePipeline(*args: Any, **kwargs: Any)

Bases: StreamInferencePipeline[WanVAECache, Wan21TransformerCache | Wan22TransformerCache, WanVAECache]

Wan 2.1 / 2.2 inference pipeline, T2V and I2V.

T2V and I2V share the same rollout loop; the difference is whether you pass an image to initialize_cache. The pipeline config’s encoder slot must agree (None for T2V, an I2V config for I2V).

Example:

pipeline: WanInferencePipeline = ...

# T2V: pass latent ``height`` and ``width``.
cache = pipeline.initialize_cache(
    text=["A cat surfing."], height=60, width=104
)
_chunk = pipeline.generate(0, cache)
pipeline.finalize(0, cache)

# I2V: pass the first frame and let sizes derive from it.
_i2v_cache = pipeline.initialize_cache(
    text=["A cat surfing."], image=first_frame
)

initialize_cache(text: list[str], image: Tensor | None = None, *, height: int | None = None, width: int | None = None, release_oneshot_encoders: bool = True) → WanInferencePipelineCache

Initialize the per-rollout cache for a batch of prompts.

Parameters:

• text – One prompt per batch element. Length must match the transformer’s batch_shape.

• image – First-frame pixels of shape [*batch_shape, 1, 3, H, W] in [-1, 1]. Required for I2V (self.encoder is set), forbidden for T2V. H and W must equal height * decoder.spatial_compression_ratio and width * decoder.spatial_compression_ratio, respectively.

• height – Pre-patchify latent height (post-VAE). Optional for I2V — derived from image when omitted; required for T2V.

• width – Pre-patchify latent width (post-VAE). Same rules as height.

• release_oneshot_encoders – Free the text and image encoders after the cache is initialized. Later calls reload them from self.config before encoding new prompts/images.

Returns:

Cache to thread through generate / finalize.

release_oneshot_encoders() → None

Free the per-rollout text and first-frame image encoders.

Idempotent. A later initialize_cache() call can reload the encoders from self.config before encoding raw prompts/images.

generate(autoregressive_index: int, cache: WanInferencePipelineCache) → torch.Tensor

Generate one decoded video chunk.

Parameters:

• autoregressive_index – AR step index, starting at 0.

• cache – Per-rollout cache from initialize_cache.

Returns:

Decoded video of shape [*batch_shape, T, C, H, W] in [-1, 1].

get_num_input_frames(autoregressive_index: int) → int

Number of input video frames the model expects at this AR step.

get_num_output_frames(autoregressive_index: int) → int

Number of decoded video frames produced at this AR step.

TAEHV

TAEHV video decoder and Hunyuan Video 1.5 codec configs.

AVAILABLE_TAEHV_CHECKPOINT_PATHS = {'hy1_5': 'https://huggingface.co/Overworld-Models/taehv1_5/resolve/main/taehv1_5.pth', 'lighttae': 'https://huggingface.co/lightx2v/Autoencoders/resolve/main/lighttaew2_1.pth'}

Checkpoint paths for the TAEHV decoder.

lighttae_state_dict_transform(sd: Mapping[str, torch.Tensor]) → dict[str, torch.Tensor]

Per-checkpoint remap for the lighttae weights. Rewrites the flat decoder.<i>.* keys to the current decoder.blocks.<i>.* layout and clips the stride=2 TGrow weights at idx 7 down to the stride=1 slice the live model expects.

class TeahvVAEDecoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/lightx2v/Autoencoders/resolve/main/lighttaew2_1.pth', state_dict_transform: ~collections.abc.Callable[[~collections.abc.Mapping[str, torch.Tensor]], dict[str, torch.Tensor]] | None = <function compose.<locals>.composed>, dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = True)

Bases: DecoderConfig

Config for the TAEHV decoder.

checkpoint_path: str = 'https://huggingface.co/lightx2v/Autoencoders/resolve/main/lighttaew2_1.pth'

Path to a pretrained TAEHV checkpoint. Defaults to the lighttae weights.

state_dict_transform() → dict[str, torch.Tensor]

Pre-load state-dict remap. Defaults to lighttae_state_dict_transform; None falls through to the bare TAEHV default (see load_from_checkpoint()).

dtype: torch.dtype

Network parameter / activation dtype.

use_cuda_graph: bool = True

Wrap the decoder forward in a CUDA graph for replay.

use_compile: bool = True

torch.compile(mode="max-autotune-no-cudagraphs").

class TeahvVAEDecoder(*args: Any, **kwargs: Any)

Bases: StreamingVideoDecoder[TAEHVCache]

TAEHV (Tiny AutoEncoder for Hunyuan Video) decoder.

Forward input is a latent [..., Tl, Cl, Hl, Wl]; output is a video tensor [..., T, C, H, W] in [-1, 1].

Set torch.backends.cudnn.benchmark = True at process start for ~5% extra on the eager seed/tail chunks.

mean: Tensor

Per-channel latent mean buffer; registered only when need_scaled.

std: Tensor

Per-channel latent standard deviation buffer; registered only when need_scaled.

initialize_autoregressive_cache() → TAEHVCache

Return an empty streaming decoder cache.

forward(input: torch.Tensor, autoregressive_index: int = 0, cache: TAEHVCache | None = None) → torch.Tensor

Decode a latent chunk to a video tensor in [-1, 1].

Parameters:

• input – Latent of shape [..., Tl, Cl, Hl, Wl].

• autoregressive_index – Unused by TAEHV; kept for the StreamingDecoder interface.

• cache – Streaming decoder cache; created on the fly when None.

Returns:

Video tensor of shape [..., T, C, H, W] in [-1, 1].

property temporal_compression_ratio: int

Pixel frames / latent frames in steady state (AR >= 1).

property spatial_compression_ratio: int

Pixel side / latent side.

get_output_temporal_size(autoregressive_index: int, input_temporal_size: int) → int

Return pixel frame count from input_temporal_size latent frames.

AR 0 applies causal padding: the first latent frame yields one pixel frame, remaining frames yield temporal_compression_ratio each. AR >= 1 is a plain multiply.

get_input_temporal_size(autoregressive_index: int, output_temporal_size: int) → int

Return latent frame count needed to produce output_temporal_size pixel frames.

Inverse of get_output_temporal_size().

class Hy15TAEHVDecoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/Overworld-Models/taehv1_5/resolve/main/taehv1_5.pth', state_dict_transform: ~collections.abc.Callable[[~collections.abc.Mapping[str, torch.Tensor]], dict[str, torch.Tensor]] | None = <function legacy_to_blocks_keys>, dtype: torch.dtype = torch.bfloat16, use_cuda_graph: bool = True, use_compile: bool = True)

Bases: DecoderConfig

Config for the Hunyuan Video 1.5 TAEHV decoder.

checkpoint_path: str = 'https://huggingface.co/Overworld-Models/taehv1_5/resolve/main/taehv1_5.pth'

Path or URL for the Hunyuan Video 1.5 TAEHV checkpoint.

state_dict_transform() → dict[str, torch.Tensor]

Pre-load state-dict remap from the published flat key layout.

dtype: torch.dtype

Network parameter / activation dtype.

use_cuda_graph: bool = True

Wrap the decoder forward in a CUDA graph for replay.

use_compile: bool = True

torch.compile(mode="max-autotune-no-cudagraphs").

class Hy15TAEHVDecoder(*args: Any, **kwargs: Any)

Bases: TeahvVAEDecoder

Hunyuan Video 1.5 TAEHV decoder with raw 32-channel latents.

class Hy15TAEHVEncoderConfig(*, _target: ~typing.Annotated[type, tyro.conf.Suppress] = <factory>, checkpoint_path: str = 'https://huggingface.co/Overworld-Models/taehv1_5/resolve/main/taehv1_5.pth', state_dict_transform: ~collections.abc.Callable[[~collections.abc.Mapping[str, torch.Tensor]], dict[str, torch.Tensor]] | None = <function legacy_to_blocks_keys>, dtype: torch.dtype = torch.bfloat16)

Bases: EncoderConfig

Config for the Hunyuan Video 1.5 TAEHV encoder.

checkpoint_path: str = 'https://huggingface.co/Overworld-Models/taehv1_5/resolve/main/taehv1_5.pth'

Path or URL for the Hunyuan Video 1.5 TAEHV checkpoint.

state_dict_transform() → dict[str, torch.Tensor]

Pre-load state-dict remap from the published flat key layout.

dtype: torch.dtype

Network parameter / activation dtype.

class Hy15TAEHVEncoder(*args: Any, **kwargs: Any)

Bases: StreamingVideoEncoder[TAEHVEncoderCache]

Hunyuan Video 1.5 TAEHV pixel-video encoder.

initialize_autoregressive_cache() → TAEHVEncoderCache

Return an empty causal encoder cache.

forward(input: torch.Tensor, autoregressive_index: int = 0, cache: TAEHVEncoderCache | None = None) → torch.Tensor

Encode frames in [-1, 1] to raw Hunyuan Video 1.5 latents.

property temporal_compression_ratio: int

Pixel frames / latent frames for complete causal groups.

property spatial_compression_ratio: int

Pixel side / latent side.

get_output_temporal_size(autoregressive_index: int, input_temporal_size: int) → int

Return latent count emitted from complete input frame groups.

get_input_temporal_size(autoregressive_index: int, output_temporal_size: int) → int

Return pixel frame count needed for output_temporal_size latents.