causal_wanvideo

Compatibility imports for the family-local causal Wan transformer.

fastvideo.models.dits.causal_wanvideo.CausalWanSelfAttention

CausalWanSelfAttention(dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm = True, eps = 1e-06, parallel_attention = False, rope_cache_policy: str = 'absolute')

fastvideo.models.dits.causal_wanvideo.CausalWanSelfAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, frame_seqlen: int = 1560)

Parameters

x: Tensor

Shape [B, L, num_heads, C / num_heads]

seq_lens: Tensor

Shape [B]

grid_sizes: Tensor

Shape [B, 3], the second dimension contains (F, H, W)

freqs: Tensor

Rope freqs, shape [1024, C / num_heads / 2]

frame_seqlen: int

Number of tokens per latent frame, e.g. 1560 for 480x832 resolution.

fastvideo.models.dits.causal_wanvideo.CausalWanTransformer3DModel

CausalWanTransformer3DModel(config: WanVideoConfig, hf_config: dict[str, Any])

fastvideo.models.dits.causal_wanvideo.CausalWanTransformer3DModel.forward

forward(*args, **kwargs)

fastvideo.models.dits.causal_wanvideo.CausalWanTransformer3DModel.unpatchify

unpatchify(x, grid_sizes)

Parameters

x: List[Tensor]

List of patchified features, each with shape [L, C_out * prod(patch_size)]

grid_sizes: Tensor

Original spatial-temporal grid dimensions before patching,

Returns

Tensor

Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8]

fastvideo.models.dits.causal_wanvideo.CausalWanTransformerBlock

CausalWanTransformerBlock(dim: int, ffn_dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, added_kv_proj_dim: int | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', rope_cache_policy: str = 'absolute')

fastvideo.models.dits.causal_wanvideo.CausalWanTransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, kv_cache: dict | None = None, crossattn_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, frame_seqlen: int | None = None) -> torch.Tensor