hunyuanvae ¶

Classes¶

fastvideo.models.vaes.hunyuanvae.AutoencoderKLHunyuanVideo ¶

AutoencoderKLHunyuanVideo(config: HunyuanVAEConfig)

Bases: Module, ParallelTiledVAE

A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in HunyuanVideo.

This model inherits from [ModelMixin]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving).

Source code in fastvideo/models/vaes/hunyuanvae.py

def __init__(
    self,
    config: HunyuanVAEConfig,
) -> None:
    nn.Module.__init__(self)
    ParallelTiledVAE.__init__(self, config)

    # TODO(will): only pass in config. We do this by manually defining a
    # config for hunyuan vae
    self.block_out_channels = config.block_out_channels

    if config.load_encoder:
        self.encoder = HunyuanVideoEncoder3D(
            in_channels=config.in_channels,
            out_channels=config.latent_channels,
            down_block_types=config.down_block_types,
            block_out_channels=config.block_out_channels,
            layers_per_block=config.layers_per_block,
            norm_num_groups=config.norm_num_groups,
            act_fn=config.act_fn,
            double_z=True,
            mid_block_add_attention=config.mid_block_add_attention,
            temporal_compression_ratio=config.temporal_compression_ratio,
            spatial_compression_ratio=config.spatial_compression_ratio,
        )
        self.quant_conv = nn.Conv3d(2 * config.latent_channels,
                                    2 * config.latent_channels,
                                    kernel_size=1)

    if config.load_decoder:
        self.decoder = HunyuanVideoDecoder3D(
            in_channels=config.latent_channels,
            out_channels=config.out_channels,
            up_block_types=config.up_block_types,
            block_out_channels=config.block_out_channels,
            layers_per_block=config.layers_per_block,
            norm_num_groups=config.norm_num_groups,
            act_fn=config.act_fn,
            time_compression_ratio=config.temporal_compression_ratio,
            spatial_compression_ratio=config.spatial_compression_ratio,
            mid_block_add_attention=config.mid_block_add_attention,
        )
        self.post_quant_conv = nn.Conv3d(config.latent_channels,
                                         config.latent_channels,
                                         kernel_size=1)

Methods:¶

fastvideo.models.vaes.hunyuanvae.AutoencoderKLHunyuanVideo.forward ¶

forward(sample: Tensor, sample_posterior: bool = False, generator: Generator | None = None) -> Tensor

Parameters:

Name	Type	Description	Default
`sample`	`torch.Tensor`	Input sample.	required
`sample_posterior`	`bool`, optional, defaults to `False`	Whether to sample from the posterior.	`False`

Source code in fastvideo/models/vaes/hunyuanvae.py

def forward(
    self,
    sample: torch.Tensor,
    sample_posterior: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""
    Args:
        sample (`torch.Tensor`): Input sample.
        sample_posterior (`bool`, *optional*, defaults to `False`):
            Whether to sample from the posterior.
    """
    x = sample
    posterior = self.encode(x).latent_dist
    z = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
    dec = self.decode(z)
    return dec

fastvideo.models.vaes.hunyuanvae.HunyuanVideoDecoder3D ¶

HunyuanVideoDecoder3D(in_channels: int = 3, out_channels: int = 3, up_block_types: tuple[str, ...] = ('HunyuanVideoUpBlock3D', 'HunyuanVideoUpBlock3D', 'HunyuanVideoUpBlock3D', 'HunyuanVideoUpBlock3D'), block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = 'silu', mid_block_add_attention=True, time_compression_ratio: int = 4, spatial_compression_ratio: int = 8)

Bases: Module

Causal decoder for 3D video-like data introduced in Hunyuan Video.

Source code in fastvideo/models/vaes/hunyuanvae.py

def __init__(
    self,
    in_channels: int = 3,
    out_channels: int = 3,
    up_block_types: tuple[str, ...] = (
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
    ),
    block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
    layers_per_block: int = 2,
    norm_num_groups: int = 32,
    act_fn: str = "silu",
    mid_block_add_attention=True,
    time_compression_ratio: int = 4,
    spatial_compression_ratio: int = 8,
):
    super().__init__()
    self.layers_per_block = layers_per_block

    self.conv_in = HunyuanVideoCausalConv3d(in_channels,
                                            block_out_channels[-1],
                                            kernel_size=3,
                                            stride=1)
    self.up_blocks = nn.ModuleList([])

    # mid
    self.mid_block = HunyuanVideoMidBlock3D(
        in_channels=block_out_channels[-1],
        resnet_eps=1e-6,
        resnet_act_fn=act_fn,
        attention_head_dim=block_out_channels[-1],
        resnet_groups=norm_num_groups,
        add_attention=mid_block_add_attention,
    )

    # up
    reversed_block_out_channels = list(reversed(block_out_channels))
    output_channel = reversed_block_out_channels[0]
    for i, up_block_type in enumerate(up_block_types):
        if up_block_type != "HunyuanVideoUpBlock3D":
            raise ValueError(f"Unsupported up_block_type: {up_block_type}")

        prev_output_channel = output_channel
        output_channel = reversed_block_out_channels[i]
        is_final_block = i == len(block_out_channels) - 1
        num_spatial_upsample_layers = int(
            np.log2(spatial_compression_ratio))
        num_time_upsample_layers = int(np.log2(time_compression_ratio))

        if time_compression_ratio == 4:
            add_spatial_upsample = bool(i < num_spatial_upsample_layers)
            add_time_upsample = bool(
                i >= len(block_out_channels) - 1 - num_time_upsample_layers
                and not is_final_block)
        else:
            raise ValueError(
                f"Unsupported time_compression_ratio: {time_compression_ratio}"
            )

        upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1,
                                                                        1)
        upsample_scale_factor_T = (2, ) if add_time_upsample else (1, )
        upsample_scale_factor = tuple(upsample_scale_factor_T +
                                      upsample_scale_factor_HW)

        up_block = HunyuanVideoUpBlock3D(
            num_layers=self.layers_per_block + 1,
            in_channels=prev_output_channel,
            out_channels=output_channel,
            add_upsample=bool(add_spatial_upsample or add_time_upsample),
            upsample_scale_factor=upsample_scale_factor,
            resnet_eps=1e-6,
            resnet_act_fn=act_fn,
            resnet_groups=norm_num_groups,
        )

        self.up_blocks.append(up_block)
        prev_output_channel = output_channel

    # out
    self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0],
                                      num_groups=norm_num_groups,
                                      eps=1e-6)
    self.conv_act = nn.SiLU()
    self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0],
                                             out_channels,
                                             kernel_size=3)

    self.gradient_checkpointing = False

fastvideo.models.vaes.hunyuanvae.HunyuanVideoEncoder3D ¶

HunyuanVideoEncoder3D(in_channels: int = 3, out_channels: int = 3, down_block_types: tuple[str, ...] = ('HunyuanVideoDownBlock3D', 'HunyuanVideoDownBlock3D', 'HunyuanVideoDownBlock3D', 'HunyuanVideoDownBlock3D'), block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = 'silu', double_z: bool = True, mid_block_add_attention=True, temporal_compression_ratio: int = 4, spatial_compression_ratio: int = 8)

Bases: Module

Causal encoder for 3D video-like data introduced in Hunyuan Video.

Source code in fastvideo/models/vaes/hunyuanvae.py

def __init__(
    self,
    in_channels: int = 3,
    out_channels: int = 3,
    down_block_types: tuple[str, ...] = (
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
    ),
    block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
    layers_per_block: int = 2,
    norm_num_groups: int = 32,
    act_fn: str = "silu",
    double_z: bool = True,
    mid_block_add_attention=True,
    temporal_compression_ratio: int = 4,
    spatial_compression_ratio: int = 8,
) -> None:
    super().__init__()

    self.conv_in = HunyuanVideoCausalConv3d(in_channels,
                                            block_out_channels[0],
                                            kernel_size=3,
                                            stride=1)
    self.mid_block: HunyuanVideoMidBlock3D | None = None
    self.down_blocks = nn.ModuleList([])

    output_channel = block_out_channels[0]
    for i, down_block_type in enumerate(down_block_types):
        if down_block_type != "HunyuanVideoDownBlock3D":
            raise ValueError(
                f"Unsupported down_block_type: {down_block_type}")

        input_channel = output_channel
        output_channel = block_out_channels[i]
        is_final_block = i == len(block_out_channels) - 1
        num_spatial_downsample_layers = int(
            np.log2(spatial_compression_ratio))
        num_time_downsample_layers = int(
            np.log2(temporal_compression_ratio))

        if temporal_compression_ratio == 4:
            add_spatial_downsample = bool(i < num_spatial_downsample_layers)
            add_time_downsample = bool(i >= (len(block_out_channels) - 1 -
                                             num_time_downsample_layers)
                                       and not is_final_block)
        elif temporal_compression_ratio == 8:
            add_spatial_downsample = bool(i < num_spatial_downsample_layers)
            add_time_downsample = bool(i < num_time_downsample_layers)
        else:
            raise ValueError(
                f"Unsupported time_compression_ratio: {temporal_compression_ratio}"
            )

        downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
        downsample_stride_T = (2, ) if add_time_downsample else (1, )
        downsample_stride = tuple(downsample_stride_T +
                                  downsample_stride_HW)

        down_block = HunyuanVideoDownBlock3D(
            num_layers=layers_per_block,
            in_channels=input_channel,
            out_channels=output_channel,
            add_downsample=bool(add_spatial_downsample
                                or add_time_downsample),
            resnet_eps=1e-6,
            resnet_act_fn=act_fn,
            resnet_groups=norm_num_groups,
            downsample_stride=downsample_stride,
            downsample_padding=0,
        )

        self.down_blocks.append(down_block)

    self.mid_block = HunyuanVideoMidBlock3D(
        in_channels=block_out_channels[-1],
        resnet_eps=1e-6,
        resnet_act_fn=act_fn,
        attention_head_dim=block_out_channels[-1],
        resnet_groups=norm_num_groups,
        add_attention=mid_block_add_attention,
    )

    self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1],
                                      num_groups=norm_num_groups,
                                      eps=1e-6)
    self.conv_act = nn.SiLU()

    conv_out_channels = 2 * out_channels if double_z else out_channels
    self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1],
                                             conv_out_channels,
                                             kernel_size=3)

    self.gradient_checkpointing = False

hunyuanvae ¶

Classes¶

fastvideo.models.vaes.hunyuanvae.AutoencoderKLHunyuanVideo ¶

Methods:¶

fastvideo.models.vaes.hunyuanvae.AutoencoderKLHunyuanVideo.forward ¶

fastvideo.models.vaes.hunyuanvae.HunyuanVideoDecoder3D ¶

fastvideo.models.vaes.hunyuanvae.HunyuanVideoEncoder3D ¶

Functions:¶