hunyuanvae ¶

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)

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

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