Skip to content

vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia

NVIDIA (CUDA) implementations.

Modules:

Name Description
autoencoder_blocks

ResnetBlock for the HunyuanImage3 autoencoder — NVIDIA CUDA + Triton implementation.
transformer_blocks

ResBlock for HunyuanImage3 -- NVIDIA CUDA + Triton implementation.

ResBlock

Bases: Module

A residual block that can optionally change the number of channels. Args: in_channels (int): The number of input channels. emb_channels (int): The number of timestep embedding channels. dropout (float): The rate of dropout. out_channels (int, optional): If specified, the number of output channels. use_conv (bool, optional): If True and out_channels is specified, use a spatial convolution instead of a smaller 1x1 convolution to change the channels in the skip connection. dims (int, optional): Determines if the signal is 1D, 2D, or 3D. up (bool, optional): If True, use this block for upsampling. down (bool, optional): If True, use this block for downsampling.

dropout instance-attribute 

dropout = dropout

emb_layers instance-attribute 

emb_layers = nn.Sequential(
    nn.SiLU(),
    nn.Linear(
        emb_channels,
        2 * self.out_channels,
        **factory_kwargs,
    ),
)

h_upd instance-attribute 

h_upd = nn.Identity()

in_channels instance-attribute 

in_channels = in_channels

in_layers instance-attribute 

in_layers = nn.Sequential(
    normalization(self.in_channels, **factory_kwargs),
    nn.SiLU(),
    conv_nd(
        dims,
        self.in_channels,
        self.out_channels,
        3,
        padding=1,
        **factory_kwargs,
    ),
)

out_channels instance-attribute 

out_channels = out_channels or self.in_channels

out_layers instance-attribute 

out_layers = nn.Sequential(
    normalization(self.out_channels, **factory_kwargs),
    nn.SiLU(),
    nn.Dropout(p=dropout),
    zero_module(
        conv_nd(
            dims,
            self.out_channels,
            self.out_channels,
            3,
            padding=1,
            **factory_kwargs,
        )
    ),
)

skip_connection instance-attribute 

skip_connection = nn.Identity()

updown instance-attribute 

updown = up or down

use_conv instance-attribute 

use_conv = use_conv

x_upd instance-attribute 

x_upd = nn.Identity()

forward 

forward(x, emb) -> Tensor

ResnetBlock

Bases: Module

conv1 instance-attribute 

conv1 = nn.Conv3d(
    in_channels,
    out_channels,
    kernel_size=3,
    stride=1,
    padding=1,
)

conv2 instance-attribute 

conv2 = nn.Conv3d(
    out_channels,
    out_channels,
    kernel_size=3,
    stride=1,
    padding=1,
)

in_channels instance-attribute 

in_channels = in_channels

nin_shortcut instance-attribute 

nin_shortcut = nn.Conv3d(
    in_channels,
    out_channels,
    kernel_size=1,
    stride=1,
    padding=0,
)

norm1 instance-attribute 

norm1 = nn.GroupNorm(
    num_groups=32,
    num_channels=in_channels,
    eps=1e-06,
    affine=True,
)

norm2 instance-attribute 

norm2 = nn.GroupNorm(
    num_groups=32,
    num_channels=out_channels,
    eps=1e-06,
    affine=True,
)

out_channels instance-attribute 

out_channels = out_channels

forward 

forward(x)