hunyuan15vae ¶

Classes¶

fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15 ¶

AutoencoderKLHunyuanVideo15(config: Hunyuan15VAEConfig)

Bases: Module, ParallelTiledVAE

A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for HunyuanVideo-1.5.

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/hunyuan15vae.py

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

    if config.load_encoder:
        self.encoder = HunyuanVideo15Encoder3D(
            in_channels=config.in_channels,
            out_channels=config.latent_channels * 2,
            block_out_channels=config.block_out_channels,
            layers_per_block=config.layers_per_block,
            temporal_compression_ratio=config.temporal_compression_ratio,
            spatial_compression_ratio=config.spatial_compression_ratio,
            downsample_match_channel=config.downsample_match_channel,
        )

    if config.load_decoder:
        self.decoder = HunyuanVideo15Decoder3D(
            in_channels=config.latent_channels,
            out_channels=config.out_channels,
            block_out_channels=list(reversed(config.block_out_channels)),
            layers_per_block=config.layers_per_block,
            temporal_compression_ratio=config.temporal_compression_ratio,
            spatial_compression_ratio=config.spatial_compression_ratio,
            upsample_match_channel=config.upsample_match_channel,
        )

    # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
    # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
    # intermediate tiles together, the memory requirement can be lowered.
    self.use_tiling = False

    # The minimal tile height and width for spatial tiling to be used
    self.tile_sample_min_height = 256
    self.tile_sample_min_width = 256
    self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling

Methods:¶

fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15.forward ¶

forward(sample: Tensor, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[Generator] = 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`
`return_dict`	`bool`, optional, defaults to `True`	Whether or not to return a [`DecoderOutput`] instead of a plain tuple.	`True`

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

def forward(
    self,
    sample: torch.Tensor,
    sample_posterior: bool = False,
    return_dict: bool = True,
    generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
    r"""
    Args:
        sample (`torch.Tensor`): Input sample.
        sample_posterior (`bool`, *optional*, defaults to `False`):
            Whether to sample from the posterior.
        return_dict (`bool`, *optional*, defaults to `True`):
            Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
    """
    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.hunyuan15vae.HunyuanVideo15AttnBlock ¶

HunyuanVideo15AttnBlock(in_channels: int)

Bases: Module

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

def __init__(self, in_channels: int):
    super().__init__()
    self.in_channels = in_channels

    self.norm = HunyuanVideo15RMS_norm(in_channels, images=False)

    self.to_q = nn.Conv3d(in_channels, in_channels, kernel_size=1)
    self.to_k = nn.Conv3d(in_channels, in_channels, kernel_size=1)
    self.to_v = nn.Conv3d(in_channels, in_channels, kernel_size=1)
    self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1)

Methods:¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock.prepare_causal_attention_mask `staticmethod` ¶

prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None)

Prepare a causal attention mask for 3D videos.

Parameters:

Name	Type	Description	Default
`n_frame`	`int`	Number of frames (temporal length).	required
`n_hw`	`int`	Product of height and width.	required
`dtype`		Desired mask dtype.	required
`device`		Device for the mask.	required
`batch_size`	`int`	If set, expands for batch.	`None`

Returns:

Type	Description
	torch.Tensor: Causal attention mask.

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

@staticmethod
def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
    """Prepare a causal attention mask for 3D videos.

    Args:
        n_frame (int): Number of frames (temporal length).
        n_hw (int): Product of height and width.
        dtype: Desired mask dtype.
        device: Device for the mask.
        batch_size (int, optional): If set, expands for batch.

    Returns:
        torch.Tensor: Causal attention mask.
    """
    seq_len = n_frame * n_hw
    mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
    for i in range(seq_len):
        i_frame = i // n_hw
        mask[i, : (i_frame + 1) * n_hw] = 0
    if batch_size is not None:
        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
    return mask

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Decoder3D ¶

HunyuanVideo15Decoder3D(in_channels: int = 32, out_channels: int = 3, block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128), layers_per_block: int = 2, spatial_compression_ratio: int = 16, temporal_compression_ratio: int = 4, upsample_match_channel: bool = True)

Bases: Module

Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner.

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

def __init__(
    self,
    in_channels: int = 32,
    out_channels: int = 3,
    block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128),
    layers_per_block: int = 2,
    spatial_compression_ratio: int = 16,
    temporal_compression_ratio: int = 4,
    upsample_match_channel: bool = True,
):
    super().__init__()
    self.layers_per_block = layers_per_block
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.repeat = block_out_channels[0] // self.in_channels

    self.conv_in = HunyuanVideo15CausalConv3d(self.in_channels, block_out_channels[0], kernel_size=3)
    self.up_blocks = nn.ModuleList([])

    # mid
    self.mid_block = HunyuanVideo15MidBlock(in_channels=block_out_channels[0])

    # up
    input_channel = block_out_channels[0]
    for i in range(len(block_out_channels)):
        output_channel = block_out_channels[i]

        add_spatial_upsample = i < np.log2(spatial_compression_ratio)
        add_temporal_upsample = i < np.log2(temporal_compression_ratio)
        if add_spatial_upsample or add_temporal_upsample:
            upsample_out_channels = block_out_channels[i + 1] if upsample_match_channel else output_channel
            up_block = HunyuanVideo15UpBlock3D(
                num_layers=self.layers_per_block + 1,
                in_channels=input_channel,
                out_channels=output_channel,
                upsample_out_channels=upsample_out_channels,
                add_temporal_upsample=add_temporal_upsample,
            )
            input_channel = upsample_out_channels
        else:
            up_block = HunyuanVideo15UpBlock3D(
                num_layers=self.layers_per_block + 1,
                in_channels=input_channel,
                out_channels=output_channel,
                upsample_out_channels=None,
                add_temporal_upsample=False,
            )
            input_channel = output_channel

        self.up_blocks.append(up_block)

    # out
    self.norm_out = HunyuanVideo15RMS_norm(block_out_channels[-1], images=False)
    self.conv_act = nn.SiLU()
    self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3)

    self.gradient_checkpointing = False

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Downsample ¶

HunyuanVideo15Downsample(in_channels: int, out_channels: int, add_temporal_downsample: bool = True)

Bases: Module

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

def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True):
    super().__init__()
    factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2
    self.conv = HunyuanVideo15CausalConv3d(in_channels, out_channels // factor, kernel_size=3)

    self.add_temporal_downsample = add_temporal_downsample
    self.group_size = factor * in_channels // out_channels

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Encoder3D ¶

HunyuanVideo15Encoder3D(in_channels: int = 3, out_channels: int = 64, block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), layers_per_block: int = 2, temporal_compression_ratio: int = 4, spatial_compression_ratio: int = 16, downsample_match_channel: bool = True)

Bases: Module

3D vae encoder for HunyuanImageRefiner.

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

def __init__(
    self,
    in_channels: int = 3,
    out_channels: int = 64,
    block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024),
    layers_per_block: int = 2,
    temporal_compression_ratio: int = 4,
    spatial_compression_ratio: int = 16,
    downsample_match_channel: bool = True,
) -> None:
    super().__init__()

    self.in_channels = in_channels
    self.out_channels = out_channels
    self.group_size = block_out_channels[-1] // self.out_channels

    self.conv_in = HunyuanVideo15CausalConv3d(in_channels, block_out_channels[0], kernel_size=3)
    self.mid_block = None
    self.down_blocks = nn.ModuleList([])

    input_channel = block_out_channels[0]
    for i in range(len(block_out_channels)):
        add_spatial_downsample = i < np.log2(spatial_compression_ratio)
        output_channel = block_out_channels[i]
        if not add_spatial_downsample:
            down_block = HunyuanVideo15DownBlock3D(
                num_layers=layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                downsample_out_channels=None,
                add_temporal_downsample=False,
            )
            input_channel = output_channel
        else:
            add_temporal_downsample = i >= np.log2(spatial_compression_ratio // temporal_compression_ratio)
            downsample_out_channels = block_out_channels[i + 1] if downsample_match_channel else output_channel
            down_block = HunyuanVideo15DownBlock3D(
                num_layers=layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                downsample_out_channels=downsample_out_channels,
                add_temporal_downsample=add_temporal_downsample,
            )
            input_channel = downsample_out_channels

        self.down_blocks.append(down_block)

    self.mid_block = HunyuanVideo15MidBlock(in_channels=block_out_channels[-1])

    self.norm_out = HunyuanVideo15RMS_norm(block_out_channels[-1], images=False)
    self.conv_act = nn.SiLU()
    self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3)

    self.gradient_checkpointing = False

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15RMS_norm ¶

HunyuanVideo15RMS_norm(dim: int, channel_first: bool = True, images: bool = True, bias: bool = False)

Bases: Module

A custom RMS normalization layer.

Parameters:

Name	Type	Description	Default
`dim`	`int`	The number of dimensions to normalize over.	required
`channel_first`	`bool`	Whether the input tensor has channels as the first dimension. Default is True.	`True`
`images`	`bool`	Whether the input represents image data. Default is True.	`True`
`bias`	`bool`	Whether to include a learnable bias term. Default is False.	`False`

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

def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
    super().__init__()
    broadcastable_dims = (1, 1, 1) if not images else (1, 1)
    shape = (dim, *broadcastable_dims) if channel_first else (dim,)

    self.channel_first = channel_first
    self.scale = dim**0.5
    self.gamma = nn.Parameter(torch.ones(shape))
    self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Upsample ¶

HunyuanVideo15Upsample(in_channels: int, out_channels: int, add_temporal_upsample: bool = True)

Bases: Module

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

def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True):
    super().__init__()
    factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2
    self.conv = HunyuanVideo15CausalConv3d(in_channels, out_channels * factor, kernel_size=3)

    self.add_temporal_upsample = add_temporal_upsample
    self.repeats = factor * out_channels // in_channels

hunyuan15vae ¶

Classes¶

fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15 ¶

Methods:¶

fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15.forward ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock ¶

Methods:¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock.prepare_causal_attention_mask staticmethod ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Decoder3D ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Downsample ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Encoder3D ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15RMS_norm ¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Upsample ¶

Functions:¶

fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock.prepare_causal_attention_mask `staticmethod` ¶