hunyuan15vae

class HunyuanVideo15CausalConv3d(nn.Module):
def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: Union[int, Tuple[int, int, int]] = 3,
    stride: Union[int, Tuple[int, int, int]] = 1,
    padding: Union[int, Tuple[int, int, int]] = 0,
    dilation: Union[int, Tuple[int, int, int]] = 1,
    bias: bool = True,
    pad_mode: str = "replicate",
) -> None:
    super().__init__()

    kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size

    self.pad_mode = pad_mode
    self.time_causal_padding = (
        kernel_size[0] // 2,
        kernel_size[0] // 2,
        kernel_size[1] // 2,
        kernel_size[1] // 2,
        kernel_size[2] - 1,
        0,
    )

    self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode)
    return self.conv(hidden_states)</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15CausalConv3d.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode)
    return self.conv(hidden_states)
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15RMS_norm
HunyuanVideo15RMS_norm(dim: int, channel_first: bool = True, images: bool = True, bias: bool = False)
A custom RMS normalization layer.
Parameters

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

Source
class HunyuanVideo15RMS_norm(nn.Module):
    r"""
    A custom RMS normalization layer.
Args:
    dim (int): The number of dimensions to normalize over.
    channel_first (bool, optional): Whether the input tensor has channels as the first dimension.
        Default is True.
    images (bool, optional): Whether the input represents image data. Default is True.
    bias (bool, optional): Whether to include a learnable bias term. Default is False.
&quot;&quot;&quot;

def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -&gt; 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

def forward(self, x):
    return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15RMS_norm.forward
forward(x)
Source
def forward(self, x):
    return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock
HunyuanVideo15AttnBlock(in_channels: int)
Source
class HunyuanVideo15AttnBlock(nn.Module):
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)

@staticmethod
def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
    &quot;&quot;&quot;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.
    &quot;&quot;&quot;
    seq_len = n_frame * n_hw
    mask = torch.full((seq_len, seq_len), float(&quot;-inf&quot;), dtype=dtype, device=device)
    for i in range(seq_len):
        i_frame = i // n_hw
        mask&#91;i, :(i_frame + 1) * n_hw&#93; = 0
    if batch_size is not None:
        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
    return mask

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    identity = x

    x = self.norm(x)

    query = self.to_q(x)
    key = self.to_k(x)
    value = self.to_v(x)

    batch_size, channels, frames, height, width = query.shape

    query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
    key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
    value = value.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()

    attention_mask = self.prepare_causal_attention_mask(frames,
                                                        height * width,
                                                        query.dtype,
                                                        query.device,
                                                        batch_size=batch_size)

    x = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)

    # batch_size, 1, frames * height * width, channels

    x = x.squeeze(1).reshape(batch_size, frames, height, width, channels).permute(0, 4, 1, 2, 3)
    x = self.proj_out(x)

    return x + identity</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock.prepare_causal_attention_mask
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

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
If set, expands for batch.

Returns

torch.Tensor: Causal attention mask.

Source
@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.
&quot;&quot;&quot;
seq_len = n_frame * n_hw
mask = torch.full((seq_len, seq_len), float(&quot;-inf&quot;), dtype=dtype, device=device)
for i in range(seq_len):
    i_frame = i // n_hw
    mask&#91;i, :(i_frame + 1) * n_hw&#93; = 0
if batch_size is not None:
    mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
return mask</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15AttnBlock.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    identity = x
x = self.norm(x)

query = self.to_q(x)
key = self.to_k(x)
value = self.to_v(x)

batch_size, channels, frames, height, width = query.shape

query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
value = value.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()

attention_mask = self.prepare_causal_attention_mask(frames,
                                                    height * width,
                                                    query.dtype,
                                                    query.device,
                                                    batch_size=batch_size)

x = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)

# batch_size, 1, frames * height * width, channels

x = x.squeeze(1).reshape(batch_size, frames, height, width, channels).permute(0, 4, 1, 2, 3)
x = self.proj_out(x)

return x + identity</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Upsample
HunyuanVideo15Upsample(in_channels: int, out_channels: int, add_temporal_upsample: bool = True)
Source
class HunyuanVideo15Upsample(nn.Module):
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

@staticmethod
def _dcae_upsample_rearrange(tensor, r1=1, r2=2, r3=2):
    &quot;&quot;&quot;
    Convert (b, r1*r2*r3*c, f, h, w) -&gt; (b, c, r1*f, r2*h, r3*w)

    Args:
        tensor: Input tensor of shape (b, r1*r2*r3*c, f, h, w)
        r1: temporal upsampling factor
        r2: height upsampling factor
        r3: width upsampling factor
    &quot;&quot;&quot;
    b, packed_c, f, h, w = tensor.shape
    factor = r1 * r2 * r3
    c = packed_c // factor

    tensor = tensor.view(b, r1, r2, r3, c, f, h, w)
    tensor = tensor.permute(0, 4, 5, 1, 6, 2, 7, 3)
    return tensor.reshape(b, c, f * r1, h * r2, w * r3)

def forward(self, x: torch.Tensor):
    r1 = 2 if self.add_temporal_upsample else 1
    h = self.conv(x)
    if self.add_temporal_upsample:
        h_first = h&#91;:, :, :1, :, :&#93;
        h_first = self._dcae_upsample_rearrange(h_first, r1=1, r2=2, r3=2)
        h_first = h_first&#91;:, :h_first.shape&#91;1&#93; // 2&#93;
        h_next = h&#91;:, :, 1:, :, :&#93;
        h_next = self._dcae_upsample_rearrange(h_next, r1=r1, r2=2, r3=2)
        h = torch.cat(&#91;h_first, h_next&#93;, dim=2)

        # shortcut computation
        x_first = x&#91;:, :, :1, :, :&#93;
        x_first = self._dcae_upsample_rearrange(x_first, r1=1, r2=2, r3=2)
        x_first = x_first.repeat_interleave(repeats=self.repeats // 2, dim=1)

        x_next = x&#91;:, :, 1:, :, :&#93;
        x_next = self._dcae_upsample_rearrange(x_next, r1=r1, r2=2, r3=2)
        x_next = x_next.repeat_interleave(repeats=self.repeats, dim=1)
        shortcut = torch.cat(&#91;x_first, x_next&#93;, dim=2)

    else:
        h = self._dcae_upsample_rearrange(h, r1=r1, r2=2, r3=2)
        shortcut = x.repeat_interleave(repeats=self.repeats, dim=1)
        shortcut = self._dcae_upsample_rearrange(shortcut, r1=r1, r2=2, r3=2)
    return h + shortcut</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Upsample.forward
forward(x: torch.Tensor)
Source
def forward(self, x: torch.Tensor):
    r1 = 2 if self.add_temporal_upsample else 1
    h = self.conv(x)
    if self.add_temporal_upsample:
        h_first = h[:, :, :1, :, :]
        h_first = self._dcae_upsample_rearrange(h_first, r1=1, r2=2, r3=2)
        h_first = h_first[:, :h_first.shape[1] // 2]
        h_next = h[:, :, 1:, :, :]
        h_next = self._dcae_upsample_rearrange(h_next, r1=r1, r2=2, r3=2)
        h = torch.cat([h_first, h_next], dim=2)
    # shortcut computation
    x_first = x&#91;:, :, :1, :, :&#93;
    x_first = self._dcae_upsample_rearrange(x_first, r1=1, r2=2, r3=2)
    x_first = x_first.repeat_interleave(repeats=self.repeats // 2, dim=1)

    x_next = x&#91;:, :, 1:, :, :&#93;
    x_next = self._dcae_upsample_rearrange(x_next, r1=r1, r2=2, r3=2)
    x_next = x_next.repeat_interleave(repeats=self.repeats, dim=1)
    shortcut = torch.cat(&#91;x_first, x_next&#93;, dim=2)

else:
    h = self._dcae_upsample_rearrange(h, r1=r1, r2=2, r3=2)
    shortcut = x.repeat_interleave(repeats=self.repeats, dim=1)
    shortcut = self._dcae_upsample_rearrange(shortcut, r1=r1, r2=2, r3=2)
return h + shortcut</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Downsample
HunyuanVideo15Downsample(in_channels: int, out_channels: int, add_temporal_downsample: bool = True)
Source
class HunyuanVideo15Downsample(nn.Module):
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

@staticmethod
def _dcae_downsample_rearrange(tensor, r1=1, r2=2, r3=2):
    &quot;&quot;&quot;
    Convert (b, c, r1*f, r2*h, r3*w) -&gt; (b, r1*r2*r3*c, f, h, w)

    This packs spatial/temporal dimensions into channels (opposite of upsample)
    &quot;&quot;&quot;
    b, c, packed_f, packed_h, packed_w = tensor.shape
    f, h, w = packed_f // r1, packed_h // r2, packed_w // r3

    tensor = tensor.view(b, c, f, r1, h, r2, w, r3)
    tensor = tensor.permute(0, 3, 5, 7, 1, 2, 4, 6)
    return tensor.reshape(b, r1 * r2 * r3 * c, f, h, w)

def forward(self, x: torch.Tensor):
    r1 = 2 if self.add_temporal_downsample else 1
    h = self.conv(x)
    if self.add_temporal_downsample:
        h_first = h&#91;:, :, :1, :, :&#93;
        h_first = self._dcae_downsample_rearrange(h_first, r1=1, r2=2, r3=2)
        h_first = torch.cat(&#91;h_first, h_first&#93;, dim=1)
        h_next = h&#91;:, :, 1:, :, :&#93;
        h_next = self._dcae_downsample_rearrange(h_next, r1=r1, r2=2, r3=2)
        h = torch.cat(&#91;h_first, h_next&#93;, dim=2)

        # shortcut computation
        x_first = x&#91;:, :, :1, :, :&#93;
        x_first = self._dcae_downsample_rearrange(x_first, r1=1, r2=2, r3=2)
        B, C, T, H, W = x_first.shape
        x_first = x_first.view(B, h.shape&#91;1&#93;, self.group_size // 2, T, H, W).mean(dim=2)
        x_next = x&#91;:, :, 1:, :, :&#93;
        x_next = self._dcae_downsample_rearrange(x_next, r1=r1, r2=2, r3=2)
        B, C, T, H, W = x_next.shape
        x_next = x_next.view(B, h.shape&#91;1&#93;, self.group_size, T, H, W).mean(dim=2)
        shortcut = torch.cat(&#91;x_first, x_next&#93;, dim=2)
    else:
        h = self._dcae_downsample_rearrange(h, r1=r1, r2=2, r3=2)
        shortcut = self._dcae_downsample_rearrange(x, r1=r1, r2=2, r3=2)
        B, C, T, H, W = shortcut.shape
        shortcut = shortcut.view(B, h.shape&#91;1&#93;, self.group_size, T, H, W).mean(dim=2)

    return h + shortcut</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Downsample.forward
forward(x: torch.Tensor)
Source
def forward(self, x: torch.Tensor):
    r1 = 2 if self.add_temporal_downsample else 1
    h = self.conv(x)
    if self.add_temporal_downsample:
        h_first = h[:, :, :1, :, :]
        h_first = self._dcae_downsample_rearrange(h_first, r1=1, r2=2, r3=2)
        h_first = torch.cat([h_first, h_first], dim=1)
        h_next = h[:, :, 1:, :, :]
        h_next = self._dcae_downsample_rearrange(h_next, r1=r1, r2=2, r3=2)
        h = torch.cat([h_first, h_next], dim=2)
    # shortcut computation
    x_first = x&#91;:, :, :1, :, :&#93;
    x_first = self._dcae_downsample_rearrange(x_first, r1=1, r2=2, r3=2)
    B, C, T, H, W = x_first.shape
    x_first = x_first.view(B, h.shape&#91;1&#93;, self.group_size // 2, T, H, W).mean(dim=2)
    x_next = x&#91;:, :, 1:, :, :&#93;
    x_next = self._dcae_downsample_rearrange(x_next, r1=r1, r2=2, r3=2)
    B, C, T, H, W = x_next.shape
    x_next = x_next.view(B, h.shape&#91;1&#93;, self.group_size, T, H, W).mean(dim=2)
    shortcut = torch.cat(&#91;x_first, x_next&#93;, dim=2)
else:
    h = self._dcae_downsample_rearrange(h, r1=r1, r2=2, r3=2)
    shortcut = self._dcae_downsample_rearrange(x, r1=r1, r2=2, r3=2)
    B, C, T, H, W = shortcut.shape
    shortcut = shortcut.view(B, h.shape&#91;1&#93;, self.group_size, T, H, W).mean(dim=2)

return h + shortcut</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15ResnetBlock
HunyuanVideo15ResnetBlock(in_channels: int, out_channels: Optional[int] = None, non_linearity: str = 'swish')
Source
class HunyuanVideo15ResnetBlock(nn.Module):
def __init__(
    self,
    in_channels: int,
    out_channels: Optional&#91;int&#93; = None,
    non_linearity: str = &quot;swish&quot;,
) -&gt; None:
    super().__init__()
    out_channels = out_channels or in_channels

    self.nonlinearity = get_act_fn(non_linearity)

    self.norm1 = HunyuanVideo15RMS_norm(in_channels, images=False)
    self.conv1 = HunyuanVideo15CausalConv3d(in_channels, out_channels, kernel_size=3)

    self.norm2 = HunyuanVideo15RMS_norm(out_channels, images=False)
    self.conv2 = HunyuanVideo15CausalConv3d(out_channels, out_channels, kernel_size=3)

    self.conv_shortcut = None
    if in_channels != out_channels:
        self.conv_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    residual = hidden_states

    hidden_states = self.norm1(hidden_states)
    hidden_states = self.nonlinearity(hidden_states)
    hidden_states = self.conv1(hidden_states)

    hidden_states = self.norm2(hidden_states)
    hidden_states = self.nonlinearity(hidden_states)
    hidden_states = self.conv2(hidden_states)

    if self.conv_shortcut is not None:
        residual = self.conv_shortcut(residual)

    return hidden_states + residual</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15ResnetBlock.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    residual = hidden_states
hidden_states = self.norm1(hidden_states)
hidden_states = self.nonlinearity(hidden_states)
hidden_states = self.conv1(hidden_states)

hidden_states = self.norm2(hidden_states)
hidden_states = self.nonlinearity(hidden_states)
hidden_states = self.conv2(hidden_states)

if self.conv_shortcut is not None:
    residual = self.conv_shortcut(residual)

return hidden_states + residual</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15MidBlock
HunyuanVideo15MidBlock(in_channels: int, num_layers: int = 1, add_attention: bool = True)
Source
class HunyuanVideo15MidBlock(nn.Module):
def __init__(
    self,
    in_channels: int,
    num_layers: int = 1,
    add_attention: bool = True,
) -&gt; None:
    super().__init__()
    self.add_attention = add_attention

    # There is always at least one resnet
    resnets = &#91;HunyuanVideo15ResnetBlock(
        in_channels=in_channels,
        out_channels=in_channels,
    )&#93;
    attentions = &#91;&#93;

    for _ in range(num_layers):
        if self.add_attention:
            attentions.append(HunyuanVideo15AttnBlock(in_channels))
        else:
            attentions.append(None)

        resnets.append(HunyuanVideo15ResnetBlock(
            in_channels=in_channels,
            out_channels=in_channels,
        ))

    self.attentions = nn.ModuleList(attentions)
    self.resnets = nn.ModuleList(resnets)

    self.gradient_checkpointing = False

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    hidden_states = self.resnets&#91;0&#93;(hidden_states)

    for attn, resnet in zip(self.attentions, self.resnets&#91;1:&#93;):
        if attn is not None:
            hidden_states = attn(hidden_states)
        hidden_states = resnet(hidden_states)

    return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15MidBlock.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = self.resnets[0](hidden_states)
for attn, resnet in zip(self.attentions, self.resnets&#91;1:&#93;):
    if attn is not None:
        hidden_states = attn(hidden_states)
    hidden_states = resnet(hidden_states)

return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15DownBlock3D
HunyuanVideo15DownBlock3D(in_channels: int, out_channels: int, num_layers: int = 1, downsample_out_channels: Optional[int] = None, add_temporal_downsample: int = True)
Source
class HunyuanVideo15DownBlock3D(nn.Module):
def __init__(
    self,
    in_channels: int,
    out_channels: int,
    num_layers: int = 1,
    downsample_out_channels: Optional&#91;int&#93; = None,
    add_temporal_downsample: int = True,
) -&gt; None:
    super().__init__()
    resnets = &#91;&#93;

    for i in range(num_layers):
        in_channels = in_channels if i == 0 else out_channels
        resnets.append(HunyuanVideo15ResnetBlock(
            in_channels=in_channels,
            out_channels=out_channels,
        ))

    self.resnets = nn.ModuleList(resnets)

    if downsample_out_channels is not None:
        self.downsamplers = nn.ModuleList(&#91;
            HunyuanVideo15Downsample(
                out_channels,
                out_channels=downsample_out_channels,
                add_temporal_downsample=add_temporal_downsample,
            )
        &#93;)
    else:
        self.downsamplers = None

    self.gradient_checkpointing = False

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    for resnet in self.resnets:
        hidden_states = resnet(hidden_states)

    if self.downsamplers is not None:
        for downsampler in self.downsamplers:
            hidden_states = downsampler(hidden_states)

    return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15DownBlock3D.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    for resnet in self.resnets:
        hidden_states = resnet(hidden_states)
if self.downsamplers is not None:
    for downsampler in self.downsamplers:
        hidden_states = downsampler(hidden_states)

return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15UpBlock3D
HunyuanVideo15UpBlock3D(in_channels: int, out_channels: int, num_layers: int = 1, upsample_out_channels: Optional[int] = None, add_temporal_upsample: bool = True)
Source
class HunyuanVideo15UpBlock3D(nn.Module):
def __init__(
    self,
    in_channels: int,
    out_channels: int,
    num_layers: int = 1,
    upsample_out_channels: Optional&#91;int&#93; = None,
    add_temporal_upsample: bool = True,
) -&gt; None:
    super().__init__()
    resnets = &#91;&#93;

    for i in range(num_layers):
        input_channels = in_channels if i == 0 else out_channels

        resnets.append(HunyuanVideo15ResnetBlock(
            in_channels=input_channels,
            out_channels=out_channels,
        ))

    self.resnets = nn.ModuleList(resnets)

    if upsample_out_channels is not None:
        self.upsamplers = nn.ModuleList(&#91;
            HunyuanVideo15Upsample(
                out_channels,
                out_channels=upsample_out_channels,
                add_temporal_upsample=add_temporal_upsample,
            )
        &#93;)
    else:
        self.upsamplers = None

    self.gradient_checkpointing = False

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for resnet in self.resnets:
            hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)

    else:
        for resnet in self.resnets:
            hidden_states = resnet(hidden_states)

    if self.upsamplers is not None:
        for upsampler in self.upsamplers:
            hidden_states = upsampler(hidden_states)

    return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15UpBlock3D.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for resnet in self.resnets:
            hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
else:
    for resnet in self.resnets:
        hidden_states = resnet(hidden_states)

if self.upsamplers is not None:
    for upsampler in self.upsamplers:
        hidden_states = upsampler(hidden_states)

return hidden_states</code></pre></details>
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)
3D vae encoder for HunyuanImageRefiner.
Source
class HunyuanVideo15Encoder3D(nn.Module):
    r"""
    3D vae encoder for HunyuanImageRefiner.
    """
def __init__(
    self,
    in_channels: int = 3,
    out_channels: int = 64,
    block_out_channels: Tuple&#91;int, ...&#93; = (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,
) -&gt; None:
    super().__init__()

    self.in_channels = in_channels
    self.out_channels = out_channels
    self.group_size = block_out_channels&#91;-1&#93; // self.out_channels

    self.conv_in = HunyuanVideo15CausalConv3d(in_channels, block_out_channels&#91;0&#93;, kernel_size=3)
    self.mid_block = None
    self.down_blocks = nn.ModuleList(&#91;&#93;)

    input_channel = block_out_channels&#91;0&#93;
    for i in range(len(block_out_channels)):
        add_spatial_downsample = i &lt; np.log2(spatial_compression_ratio)
        output_channel = block_out_channels&#91;i&#93;
        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 &gt;= np.log2(spatial_compression_ratio // temporal_compression_ratio)
            downsample_out_channels = block_out_channels&#91;i + 1&#93; 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&#91;-1&#93;)

    self.norm_out = HunyuanVideo15RMS_norm(block_out_channels&#91;-1&#93;, images=False)
    self.conv_act = nn.SiLU()
    self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels&#91;-1&#93;, out_channels, kernel_size=3)

    self.gradient_checkpointing = False

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    hidden_states = self.conv_in(hidden_states)

    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for down_block in self.down_blocks:
            hidden_states = self._gradient_checkpointing_func(down_block, hidden_states)

        hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
    else:
        for down_block in self.down_blocks:
            hidden_states = down_block(hidden_states)

        hidden_states = self.mid_block(hidden_states)

    batch_size, _, frame, height, width = hidden_states.shape
    short_cut = hidden_states.view(batch_size, -1, self.group_size, frame, height, width).mean(dim=2)

    hidden_states = self.norm_out(hidden_states)
    hidden_states = self.conv_act(hidden_states)
    hidden_states = self.conv_out(hidden_states)

    hidden_states += short_cut

    return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Encoder3D.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = self.conv_in(hidden_states)
if torch.is_grad_enabled() and self.gradient_checkpointing:
    for down_block in self.down_blocks:
        hidden_states = self._gradient_checkpointing_func(down_block, hidden_states)

    hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
else:
    for down_block in self.down_blocks:
        hidden_states = down_block(hidden_states)

    hidden_states = self.mid_block(hidden_states)

batch_size, _, frame, height, width = hidden_states.shape
short_cut = hidden_states.view(batch_size, -1, self.group_size, frame, height, width).mean(dim=2)

hidden_states = self.norm_out(hidden_states)
hidden_states = self.conv_act(hidden_states)
hidden_states = self.conv_out(hidden_states)

hidden_states += short_cut

return hidden_states</code></pre></details>
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)
Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner.
Source
class HunyuanVideo15Decoder3D(nn.Module):
    r"""
    Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner.
    """
def __init__(
    self,
    in_channels: int = 32,
    out_channels: int = 3,
    block_out_channels: Tuple&#91;int, ...&#93; = (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&#91;0&#93; // self.in_channels

    self.conv_in = HunyuanVideo15CausalConv3d(self.in_channels, block_out_channels&#91;0&#93;, kernel_size=3)
    self.up_blocks = nn.ModuleList(&#91;&#93;)

    # mid
    self.mid_block = HunyuanVideo15MidBlock(in_channels=block_out_channels&#91;0&#93;)

    # up
    input_channel = block_out_channels&#91;0&#93;
    for i in range(len(block_out_channels)):
        output_channel = block_out_channels&#91;i&#93;

        add_spatial_upsample = i &lt; np.log2(spatial_compression_ratio)
        add_temporal_upsample = i &lt; np.log2(temporal_compression_ratio)
        if add_spatial_upsample or add_temporal_upsample:
            upsample_out_channels = block_out_channels&#91;i + 1&#93; 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&#91;-1&#93;, images=False)
    self.conv_act = nn.SiLU()
    self.conv_out = HunyuanVideo15CausalConv3d(block_out_channels&#91;-1&#93;, out_channels, kernel_size=3)

    self.gradient_checkpointing = False

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1)

    if torch.is_grad_enabled() and self.gradient_checkpointing:
        hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)

        for up_block in self.up_blocks:
            hidden_states = self._gradient_checkpointing_func(up_block, hidden_states)
    else:
        hidden_states = self.mid_block(hidden_states)

        for up_block in self.up_blocks:
            hidden_states = up_block(hidden_states)

    # post-process
    hidden_states = self.norm_out(hidden_states)
    hidden_states = self.conv_act(hidden_states)
    hidden_states = self.conv_out(hidden_states)
    return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.HunyuanVideo15Decoder3D.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1)
if torch.is_grad_enabled() and self.gradient_checkpointing:
    hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)

    for up_block in self.up_blocks:
        hidden_states = self._gradient_checkpointing_func(up_block, hidden_states)
else:
    hidden_states = self.mid_block(hidden_states)

    for up_block in self.up_blocks:
        hidden_states = up_block(hidden_states)

# post-process
hidden_states = self.norm_out(hidden_states)
hidden_states = self.conv_act(hidden_states)
hidden_states = self.conv_out(hidden_states)
return hidden_states</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15
AutoencoderKLHunyuanVideo15(config: Hunyuan15VAEConfig)
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
class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE):
    r"""
    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 &#91;`ModelMixin`&#93;. Check the superclass documentation for it&#x27;s generic methods implemented
for all models (such as downloading or saving).
&quot;&quot;&quot;

_supports_gradient_checkpointing = True

def __init__(
    self,
    config: Hunyuan15VAEConfig,
) -&gt; 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

def _encode(self, x: torch.Tensor) -&gt; torch.Tensor:
    x = self.encoder(x)
    return x

def _decode(self, z: torch.Tensor) -&gt; torch.Tensor:
    dec = self.decoder(z)
    return dec

def forward(
    self,
    sample: torch.Tensor,
    sample_posterior: bool = False,
    return_dict: bool = True,
    generator: Optional&#91;torch.Generator&#93; = None,
) -&gt; torch.Tensor:
    r&quot;&quot;&quot;
    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 &#91;`DecoderOutput`&#93; instead of a plain tuple.
    &quot;&quot;&quot;
    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</code></pre></details>
fastvideo.models.vaes.hunyuan15vae.AutoencoderKLHunyuanVideo15.forward
forward(sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[torch.Generator] = None) -> torch.Tensor
Parameters

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.

Source
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