hunyuanvideo ¶

Classes¶

fastvideo.models.dits.hunyuanvideo.FinalLayer ¶

FinalLayer(hidden_size, patch_size, out_channels, dtype=None, prefix: str = '')

Bases: Module

The final layer of DiT that projects features to pixel space.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(self,
             hidden_size,
             patch_size,
             out_channels,
             dtype=None,
             prefix: str = "") -> None:
    super().__init__()

    # Normalization
    self.norm_final = nn.LayerNorm(hidden_size,
                                   eps=1e-6,
                                   elementwise_affine=False,
                                   dtype=dtype)

    output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels

    self.linear = ReplicatedLinear(hidden_size,
                                   output_dim,
                                   bias=True,
                                   params_dtype=dtype,
                                   prefix=f"{prefix}.linear")

    # Modulation
    self.adaLN_modulation = ModulateProjection(
        hidden_size,
        factor=2,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.adaLN_modulation")

fastvideo.models.dits.hunyuanvideo.HunyuanRMSNorm ¶

HunyuanRMSNorm(dim: int, elementwise_affine=True, eps: float = 1e-06, device=None, dtype=None)

Bases: Module

Initialize the RMSNorm normalization layer.

Parameters:

Name	Type	Description	Default
`dim`	`int`	The dimension of the input tensor.	required
`eps`	`float`	A small value added to the denominator for numerical stability. Default is 1e-6.	`1e-06`

Attributes:

Name	Type	Description
`eps`	`float`	A small value added to the denominator for numerical stability.
`weight`	`Parameter`	Learnable scaling parameter.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(
    self,
    dim: int,
    elementwise_affine=True,
    eps: float = 1e-6,
    device=None,
    dtype=None,
):
    """
    Initialize the RMSNorm normalization layer.

    Args:
        dim (int): The dimension of the input tensor.
        eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.

    Attributes:
        eps (float): A small value added to the denominator for numerical stability.
        weight (nn.Parameter): Learnable scaling parameter.

    """
    factory_kwargs = {"device": device, "dtype": dtype}
    super().__init__()
    self.eps = eps
    if elementwise_affine:
        self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))

Methods:¶

fastvideo.models.dits.hunyuanvideo.HunyuanRMSNorm.forward ¶

forward(x)

Forward pass through the RMSNorm layer.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	The input tensor.	required

Returns:

Type	Description
	torch.Tensor: The output tensor after applying RMSNorm.

Source code in fastvideo/models/dits/hunyuanvideo.py

def forward(self, x):
    """
    Forward pass through the RMSNorm layer.

    Args:
        x (torch.Tensor): The input tensor.

    Returns:
        torch.Tensor: The output tensor after applying RMSNorm.

    """
    output = self._norm(x.float()).type_as(x)
    if hasattr(self, "weight"):
        output = output * self.weight
    return output

fastvideo.models.dits.hunyuanvideo.HunyuanVideoTransformer3DModel ¶

HunyuanVideoTransformer3DModel(config: HunyuanVideoConfig, hf_config: dict[str, Any])

Bases: BaseDiT

HunyuanVideo Transformer backbone adapted for distributed training.

This implementation uses distributed attention and linear layers for efficient parallel processing across multiple GPUs.

Based on the architecture from: - Flux.1: https://github.com/black-forest-labs/flux - MMDiT: http://arxiv.org/abs/2403.03206

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):
    super().__init__(config=config, hf_config=hf_config)

    self.patch_size = [
        config.patch_size_t, config.patch_size, config.patch_size
    ]
    self.in_channels = config.in_channels
    self.num_channels_latents = config.num_channels_latents
    self.out_channels = config.in_channels if config.out_channels is None else config.out_channels
    self.unpatchify_channels = self.out_channels
    self.guidance_embeds = config.guidance_embeds
    self.rope_dim_list = list(config.rope_axes_dim)
    self.rope_theta = config.rope_theta
    self.text_states_dim = config.text_embed_dim
    self.text_states_dim_2 = config.pooled_projection_dim
    # TODO(will): hack?
    self.dtype = config.dtype

    pe_dim = config.hidden_size // config.num_attention_heads
    if sum(config.rope_axes_dim) != pe_dim:
        raise ValueError(
            f"Got {config.rope_axes_dim} but expected positional dim {pe_dim}"
        )

    self.hidden_size = config.hidden_size
    self.num_attention_heads = config.num_attention_heads
    self.num_channels_latents = config.num_channels_latents

    # Image projection
    self.img_in = PatchEmbed(self.patch_size,
                             self.in_channels,
                             self.hidden_size,
                             dtype=config.dtype,
                             prefix=f"{config.prefix}.img_in")

    self.txt_in = SingleTokenRefiner(self.text_states_dim,
                                     config.hidden_size,
                                     config.num_attention_heads,
                                     depth=config.num_refiner_layers,
                                     dtype=config.dtype,
                                     prefix=f"{config.prefix}.txt_in")

    # Time modulation
    self.time_in = TimestepEmbedder(self.hidden_size,
                                    act_layer="silu",
                                    dtype=config.dtype,
                                    prefix=f"{config.prefix}.time_in")

    # Text modulation
    self.vector_in = MLP(self.text_states_dim_2,
                         self.hidden_size,
                         self.hidden_size,
                         act_type="silu",
                         dtype=config.dtype,
                         prefix=f"{config.prefix}.vector_in")

    # Guidance modulation
    self.guidance_in = (TimestepEmbedder(
        self.hidden_size,
        act_layer="silu",
        dtype=config.dtype,
        prefix=f"{config.prefix}.guidance_in")
                        if self.guidance_embeds else None)

    # Double blocks
    self.double_blocks = nn.ModuleList([
        MMDoubleStreamBlock(
            config.hidden_size,
            config.num_attention_heads,
            mlp_ratio=config.mlp_ratio,
            dtype=config.dtype,
            supported_attention_backends=self._supported_attention_backends,
            prefix=f"{config.prefix}.double_blocks.{i}")
        for i in range(config.num_layers)
    ])

    # Single blocks
    self.single_blocks = nn.ModuleList([
        MMSingleStreamBlock(
            config.hidden_size,
            config.num_attention_heads,
            mlp_ratio=config.mlp_ratio,
            dtype=config.dtype,
            supported_attention_backends=self._supported_attention_backends,
            prefix=f"{config.prefix}.single_blocks.{i+config.num_layers}")
        for i in range(config.num_single_layers)
    ])

    self.final_layer = FinalLayer(config.hidden_size,
                                  self.patch_size,
                                  self.out_channels,
                                  dtype=config.dtype,
                                  prefix=f"{config.prefix}.final_layer")

    self.__post_init__()

Methods:¶

fastvideo.models.dits.hunyuanvideo.HunyuanVideoTransformer3DModel.forward ¶

forward(hidden_states: Tensor, encoder_hidden_states: Tensor | list[Tensor], timestep: LongTensor, encoder_hidden_states_image: Tensor | list[Tensor] | None = None, guidance=None, **kwargs)

Forward pass of the HunyuanDiT model.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input image/video latents [B, C, T, H, W]	required
`encoder_hidden_states`	`Tensor \| list[Tensor]`	Text embeddings [B, L, D]	required
`timestep`	`LongTensor`	Diffusion timestep	required
`guidance`		Guidance scale for CFG	`None`

Returns:

Type	Description
	Tuple of (output)

Source code in fastvideo/models/dits/hunyuanvideo.py

def forward(self,
            hidden_states: torch.Tensor,
            encoder_hidden_states: torch.Tensor | list[torch.Tensor],
            timestep: torch.LongTensor,
            encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
            | None = None,
            guidance=None,
            **kwargs):
    """
    Forward pass of the HunyuanDiT model.

    Args:
        hidden_states: Input image/video latents [B, C, T, H, W]
        encoder_hidden_states: Text embeddings [B, L, D]
        timestep: Diffusion timestep
        guidance: Guidance scale for CFG

    Returns:
        Tuple of (output)
    """
    if guidance is None:
        guidance = torch.tensor([6016.0],
                                device=hidden_states.device,
                                dtype=hidden_states.dtype)

    img = x = hidden_states
    t = timestep

    # Split text embeddings - first token is global, rest are per-token
    if isinstance(encoder_hidden_states, torch.Tensor):
        txt = encoder_hidden_states[:, 1:]
        text_states_2 = encoder_hidden_states[:, 0, :self.text_states_dim_2]
    else:
        txt = encoder_hidden_states[0]
        text_states_2 = encoder_hidden_states[1]

    # Get spatial dimensions
    _, _, ot, oh, ow = x.shape  # codespell:ignore
    tt, th, tw = (
        ot // self.patch_size[0],  # codespell:ignore
        oh // self.patch_size[1],
        ow // self.patch_size[2],
    )

    # Get rotary embeddings
    freqs_cos, freqs_sin = get_rotary_pos_embed(
        (tt, th, tw), self.hidden_size,
        self.num_attention_heads, self.rope_dim_list, self.rope_theta)
    freqs_cos = freqs_cos.to(x.device)
    freqs_sin = freqs_sin.to(x.device)
    # Prepare modulation vectors
    vec = self.time_in(t)

    # Add text modulation
    vec = vec + self.vector_in(text_states_2)

    # Add guidance modulation if needed
    if self.guidance_in and guidance is not None:
        vec = vec + self.guidance_in(guidance)
    # Embed image and text
    img = self.img_in(img)
    img, original_seq_len = sequence_model_parallel_shard(img, dim=1)
    txt = self.txt_in(txt, t)
    txt_seq_len = txt.shape[1]
    img_seq_len = img.shape[1]

    freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

    # Process through double stream blocks
    for index, block in enumerate(self.double_blocks):
        double_block_args = [img, txt, vec, freqs_cis, original_seq_len]
        img, txt = block(*double_block_args)
    # Merge txt and img to pass through single stream blocks
    x = torch.cat((img, txt), 1)

    # Process through single stream blocks
    if len(self.single_blocks) > 0:
        for index, block in enumerate(self.single_blocks):
            single_block_args = [
                x,
                vec,
                txt_seq_len,
                freqs_cis,
                original_seq_len,
            ]
            x = block(*single_block_args)

    # Extract image features
    img = x[:, :img_seq_len, ...]

    # Final layer processing
    img = sequence_model_parallel_all_gather_with_unpad(img, original_seq_len, dim=1)
    img = self.final_layer(img, vec)
    # Unpatchify to get original shape
    img = unpatchify(img, tt, th, tw, self.patch_size, self.out_channels)

    return img

fastvideo.models.dits.hunyuanvideo.IndividualTokenRefinerBlock ¶

IndividualTokenRefinerBlock(hidden_size, num_attention_heads, mlp_ratio=4.0, qkv_bias=True, dtype=None, prefix: str = '')

Bases: Module

A transformer block for refining individual tokens with self-attention.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(
    self,
    hidden_size,
    num_attention_heads,
    mlp_ratio=4.0,
    qkv_bias=True,
    dtype=None,
    prefix: str = "",
) -> None:
    super().__init__()
    self.num_attention_heads = num_attention_heads
    mlp_hidden_dim = int(hidden_size * mlp_ratio)

    # Normalization and attention
    self.norm1 = nn.LayerNorm(hidden_size,
                              eps=1e-6,
                              elementwise_affine=True,
                              dtype=dtype)

    self.self_attn_qkv = ReplicatedLinear(hidden_size,
                                          hidden_size * 3,
                                          bias=qkv_bias,
                                          params_dtype=dtype,
                                          prefix=f"{prefix}.self_attn_qkv")

    self.self_attn_proj = ReplicatedLinear(
        hidden_size,
        hidden_size,
        bias=qkv_bias,
        params_dtype=dtype,
        prefix=f"{prefix}.self_attn_proj")

    # MLP
    self.norm2 = nn.LayerNorm(hidden_size,
                              eps=1e-6,
                              elementwise_affine=True,
                              dtype=dtype)
    self.mlp = MLP(hidden_size,
                   mlp_hidden_dim,
                   bias=True,
                   act_type="silu",
                   dtype=dtype,
                   prefix=f"{prefix}.mlp")

    # Modulation
    self.adaLN_modulation = ModulateProjection(
        hidden_size,
        factor=2,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.adaLN_modulation")

    # Scaled dot product attention
    self.attn = LocalAttention(
        num_heads=num_attention_heads,
        head_size=hidden_size // num_attention_heads,
        # TODO: remove hardcode
        supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                      AttentionBackendEnum.TORCH_SDPA),
    )

fastvideo.models.dits.hunyuanvideo.MMDoubleStreamBlock ¶

MMDoubleStreamBlock(hidden_size: int, num_attention_heads: int, mlp_ratio: float, dtype: dtype | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '')

Bases: Module

A multimodal DiT block with separate modulation for text and image/video, using distributed attention and linear layers.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    mlp_ratio: float,
    dtype: torch.dtype | None = None,
    supported_attention_backends: tuple[AttentionBackendEnum, ...]
    | None = None,
    prefix: str = "",
):
    super().__init__()

    self.deterministic = False
    self.num_attention_heads = num_attention_heads
    head_dim = hidden_size // num_attention_heads
    mlp_hidden_dim = int(hidden_size * mlp_ratio)

    # Image modulation components
    self.img_mod = ModulateProjection(
        hidden_size,
        factor=6,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.img_mod",
    )

    # Fused operations for image stream
    self.img_attn_norm = LayerNormScaleShift(hidden_size,
                                             norm_type="layer",
                                             elementwise_affine=False,
                                             dtype=dtype)
    self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
        hidden_size,
        norm_type="layer",
        elementwise_affine=False,
        dtype=dtype)
    self.img_mlp_residual = ScaleResidual()

    # Image attention components
    self.img_attn_qkv = ReplicatedLinear(hidden_size,
                                         hidden_size * 3,
                                         bias=True,
                                         params_dtype=dtype,
                                         prefix=f"{prefix}.img_attn_qkv")

    self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
    self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

    self.img_attn_proj = ReplicatedLinear(hidden_size,
                                          hidden_size,
                                          bias=True,
                                          params_dtype=dtype,
                                          prefix=f"{prefix}.img_attn_proj")

    self.img_mlp = MLP(hidden_size,
                       mlp_hidden_dim,
                       bias=True,
                       dtype=dtype,
                       prefix=f"{prefix}.img_mlp")

    # Text modulation components
    self.txt_mod = ModulateProjection(
        hidden_size,
        factor=6,
        act_layer="silu",
        dtype=dtype,
        prefix=f"{prefix}.txt_mod",
    )

    # Fused operations for text stream
    self.txt_attn_norm = LayerNormScaleShift(hidden_size,
                                             norm_type="layer",
                                             elementwise_affine=False,
                                             dtype=dtype)
    self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
        hidden_size,
        norm_type="layer",
        elementwise_affine=False,
        dtype=dtype)
    self.txt_mlp_residual = ScaleResidual()

    # Text attention components
    self.txt_attn_qkv = ReplicatedLinear(hidden_size,
                                         hidden_size * 3,
                                         bias=True,
                                         params_dtype=dtype)

    # QK norm layers for text
    self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
    self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

    self.txt_attn_proj = ReplicatedLinear(hidden_size,
                                          hidden_size,
                                          bias=True,
                                          params_dtype=dtype)

    self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)

    # Distributed attention
    self.attn = DistributedAttention(
        num_heads=num_attention_heads,
        head_size=head_dim,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=f"{prefix}.attn")

fastvideo.models.dits.hunyuanvideo.MMSingleStreamBlock ¶

MMSingleStreamBlock(hidden_size: int, num_attention_heads: int, mlp_ratio: float = 4.0, dtype: dtype | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '')

Bases: Module

A DiT block with parallel linear layers using distributed attention and tensor parallelism.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    mlp_ratio: float = 4.0,
    dtype: torch.dtype | None = None,
    supported_attention_backends: tuple[AttentionBackendEnum, ...]
    | None = None,
    prefix: str = "",
):
    super().__init__()

    self.deterministic = False
    self.hidden_size = hidden_size
    self.num_attention_heads = num_attention_heads
    head_dim = hidden_size // num_attention_heads
    mlp_hidden_dim = int(hidden_size * mlp_ratio)
    self.mlp_hidden_dim = mlp_hidden_dim

    # Combined QKV and MLP input projection
    self.linear1 = ReplicatedLinear(hidden_size,
                                    hidden_size * 3 + mlp_hidden_dim,
                                    bias=True,
                                    params_dtype=dtype,
                                    prefix=f"{prefix}.linear1")

    # Combined projection and MLP output
    self.linear2 = ReplicatedLinear(hidden_size + mlp_hidden_dim,
                                    hidden_size,
                                    bias=True,
                                    params_dtype=dtype,
                                    prefix=f"{prefix}.linear2")

    # QK norm layers
    self.q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
    self.k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

    # Fused operations with better naming
    self.input_norm_scale_shift = LayerNormScaleShift(
        hidden_size,
        norm_type="layer",
        eps=1e-6,
        elementwise_affine=False,
        dtype=dtype)
    self.output_residual = ScaleResidual()

    # Activation function
    self.mlp_act = nn.GELU(approximate="tanh")

    # Modulation
    self.modulation = ModulateProjection(hidden_size,
                                         factor=3,
                                         act_layer="silu",
                                         dtype=dtype,
                                         prefix=f"{prefix}.modulation")

    # Distributed attention
    self.attn = DistributedAttention(
        num_heads=num_attention_heads,
        head_size=head_dim,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=f"{prefix}.attn")

fastvideo.models.dits.hunyuanvideo.SingleTokenRefiner ¶

SingleTokenRefiner(in_channels, hidden_size, num_attention_heads, depth=2, qkv_bias=True, dtype=None, prefix: str = '')

Bases: Module

A token refiner that processes text embeddings with attention to improve their representation for cross-attention with image features.

Source code in fastvideo/models/dits/hunyuanvideo.py

def __init__(
    self,
    in_channels,
    hidden_size,
    num_attention_heads,
    depth=2,
    qkv_bias=True,
    dtype=None,
    prefix: str = "",
) -> None:
    super().__init__()

    # Input projection
    self.input_embedder = ReplicatedLinear(
        in_channels,
        hidden_size,
        bias=True,
        params_dtype=dtype,
        prefix=f"{prefix}.input_embedder")

    # Timestep embedding
    self.t_embedder = TimestepEmbedder(hidden_size,
                                       act_layer="silu",
                                       dtype=dtype,
                                       prefix=f"{prefix}.t_embedder")

    # Context embedding
    self.c_embedder = MLP(in_channels,
                          hidden_size,
                          hidden_size,
                          act_type="silu",
                          dtype=dtype,
                          prefix=f"{prefix}.c_embedder")

    # Refiner blocks
    self.refiner_blocks = nn.ModuleList([
        IndividualTokenRefinerBlock(
            hidden_size,
            num_attention_heads,
            qkv_bias=qkv_bias,
            dtype=dtype,
            prefix=f"{prefix}.refiner_blocks.{i}",
        ) for i in range(depth)
    ])

hunyuanvideo ¶

Classes¶

fastvideo.models.dits.hunyuanvideo.FinalLayer ¶

fastvideo.models.dits.hunyuanvideo.HunyuanRMSNorm ¶

Methods:¶

fastvideo.models.dits.hunyuanvideo.HunyuanRMSNorm.forward ¶

fastvideo.models.dits.hunyuanvideo.HunyuanVideoTransformer3DModel ¶

Methods:¶

fastvideo.models.dits.hunyuanvideo.HunyuanVideoTransformer3DModel.forward ¶

fastvideo.models.dits.hunyuanvideo.IndividualTokenRefinerBlock ¶

fastvideo.models.dits.hunyuanvideo.MMDoubleStreamBlock ¶

fastvideo.models.dits.hunyuanvideo.MMSingleStreamBlock ¶

fastvideo.models.dits.hunyuanvideo.SingleTokenRefiner ¶

Functions:¶