--- title: hyworldvae ---

fastvideo.models.vaes.hyworldvae.get_act_fn

get_act_fn(act_fn_name: str) -> nn.Module

Get an activation function by name.

fastvideo.models.vaes.hyworldvae.Hunyuan15VAEConfig

Hunyuan15VAEConfig(arch_config: VAEArchConfig = Hunyuan15VAEArchConfig(), load_encoder: bool = True, load_decoder: bool = True, tile_sample_min_height: int = 256, tile_sample_min_width: int = 256, tile_sample_min_num_frames: int = 16, tile_sample_stride_height: int = 192, tile_sample_stride_width: int = 192, tile_sample_stride_num_frames: int = 12, blend_num_frames: int = 0, use_tiling: bool = True, use_temporal_tiling: bool = True, use_parallel_tiling: bool = True, use_temporal_scaling_frames: bool = True, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE

ParallelTiledVAE(config: VAEConfig, **kwargs)

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.to

to(device) -> ParallelTiledVAE

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.encode

encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.decode

decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_v

blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_h

blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_t

blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.spatial_tiled_encode

spatial_tiled_encode(x: torch.Tensor) -> torch.Tensor

Encode a batch of images using a tiled encoder.

Parameters

x: `torch.Tensor`

Input batch of videos.

Returns

torch.Tensor

torch.Tensor: The latent representation of the encoded videos.

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.parallel_tiled_decode

parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensor

Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.spatial_tiled_decode

spatial_tiled_decode(z: torch.Tensor) -> torch.Tensor

Decode a batch of images using a tiled decoder.

Parameters

z: `torch.Tensor`

Input batch of latent vectors.

Returns

torch.Tensor

torch.Tensor: The decoded images.

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.tiled_encode

tiled_encode(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.tiled_decode

tiled_decode(z: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.enable_tiling

enable_tiling(tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None, tile_sample_min_num_frames: int | None = None, tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None, tile_sample_stride_num_frames: int | None = None, blend_num_frames: int | None = None, use_tiling: bool | None = None, use_temporal_tiling: bool | None = None, use_parallel_tiling: bool | None = None) -> None

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

Parameters

tile_sample_min_height: `int`, *optional*

The minimum height required for a sample to be separated into tiles across the height dimension.

tile_sample_min_width: `int`, *optional*

The minimum width required for a sample to be separated into tiles across the width dimension.

tile_sample_min_num_frames: `int`, *optional*

The minimum number of frames required for a sample to be separated into tiles across the frame dimension.

tile_sample_stride_height: `int`, *optional*

The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension.

tile_sample_stride_width: `int`, *optional*

The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension.

tile_sample_stride_num_frames: `int`, *optional*

The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts produced across the frame dimension.

fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.disable_tiling

disable_tiling() -> None

Disable tiled VAE decoding. If enable_tiling was previously enabled, this method will go back to computing decoding in one step.

fastvideo.models.vaes.hyworldvae.HYWorldRMS_norm

HYWorldRMS_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.

fastvideo.models.vaes.hyworldvae.HYWorldRMS_norm.forward

forward(x)

fastvideo.models.vaes.hyworldvae.HYWorldAttnBlock

HYWorldAttnBlock(in_channels: int)

fastvideo.models.vaes.hyworldvae.HYWorldAttnBlock.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.

fastvideo.models.vaes.hyworldvae.HYWorldAttnBlock.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.hyworldvae.HYWorldCausalConv3d

HYWorldCausalConv3d(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')

Causal Conv3d with optional cache support for temporal feature caching.

Source
class HYWorldCausalConv3d(nn.Module):
    """Causal Conv3d with optional cache support for temporal feature caching."""

    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
        # Padding format: (W_left, W_right, H_left, H_right, T_left, T_right)
        self.time_causal_padding = (
            kernel_size[0] // 2,  # W_left (spatial)
            kernel_size[0] // 2,  # W_right (spatial)
            kernel_size[1] // 2,  # H_left (spatial)
            kernel_size[1] // 2,  # H_right (spatial)
            kernel_size[2] - 1,  # T_left (temporal causal padding)
            0,  # T_right (no future padding for causal)
        )

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

    def forward(self, hidden_states: torch.Tensor, cache_x: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            cache_x: Optional cached frames from previous chunk, shape (B, C, CACHE_T, H, W)
                    When provided, uses cached frames instead of padding for temporal dimension.
        """
        padding = list(self.time_causal_padding)

        if cache_x is not None and self.time_causal_padding[4] > 0:  # Has temporal padding and cache
            cache_x = cache_x.to(hidden_states.device)
            # Concatenate cached frames with current input on temporal dimension
            hidden_states = torch.cat([cache_x, hidden_states], dim=2)
            # Reduce temporal padding since we have cached frames
            padding[4] -= cache_x.shape[2]

        hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode)
        return self.conv(hidden_states)

fastvideo.models.vaes.hyworldvae.HYWorldCausalConv3d.forward

forward(hidden_states: torch.Tensor, cache_x: Optional[torch.Tensor] = None) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

cache_x: Optional[torch.Tensor]

Optional cached frames from previous chunk, shape (B, C, CACHE_T, H, W) When provided, uses cached frames instead of padding for temporal dimension.

Source
def forward(self, hidden_states: torch.Tensor, cache_x: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        cache_x: Optional cached frames from previous chunk, shape (B, C, CACHE_T, H, W)
                When provided, uses cached frames instead of padding for temporal dimension.
    """
    padding = list(self.time_causal_padding)

    if cache_x is not None and self.time_causal_padding[4] > 0:  # Has temporal padding and cache
        cache_x = cache_x.to(hidden_states.device)
        # Concatenate cached frames with current input on temporal dimension
        hidden_states = torch.cat([cache_x, hidden_states], dim=2)
        # Reduce temporal padding since we have cached frames
        padding[4] -= cache_x.shape[2]

    hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode)
    return self.conv(hidden_states)

fastvideo.models.vaes.hyworldvae.HYWorldUpsample

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

Hierarchical upsampling with temporal/spatial support and optional caching.

Source
class HYWorldUpsample(nn.Module):
    """Hierarchical upsampling with temporal/spatial support and optional caching."""

    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 = HYWorldCausalConv3d(in_channels, out_channels * factor, kernel_size=3)
        self.add_temporal_upsample = add_temporal_upsample
        self.repeats = factor * out_channels // in_channels

    def forward(
        self,
        x: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
        first_chunk: bool = False,
    ):
        """
        Forward pass with optional temporal caching.

        Args:
            x: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
            first_chunk: Whether this is the first chunk (affects temporal upsample behavior)
        """
        r1 = 2 if self.add_temporal_upsample else 1
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            h = self.conv(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            h = self.conv(x)

        if self.add_temporal_upsample:
            if first_chunk:
                # First chunk: only spatial upsample
                h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
                h = h[:, :h.shape[1] // 2]
                # Compute the shortcut part
                shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
                shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1)
            elif feat_cache is None and x.shape[2] > 1:
                # No cache and multiple frames: separate first frame and rest
                h_first = h[:, :, :1, :, :]
                h_rest = h[:, :, 1:, :, :]
                x_first = x[:, :, :1, :, :]
                x_rest = x[:, :, 1:, :, :]

                # First frame: only spatial upsample
                h_first = rearrange(h_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
                h_first = h_first[:, :h_first.shape[1] // 2]
                shortcut_first = rearrange(x_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
                shortcut_first = shortcut_first.repeat_interleave(repeats=self.repeats // 2, dim=1)
                out_first = h_first + shortcut_first

                # Remaining frames: spatio-temporal upsample
                h_rest = rearrange(h_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
                shortcut_rest = rearrange(x_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
                shortcut_rest = shortcut_rest.repeat_interleave(repeats=self.repeats, dim=1)
                out_rest = h_rest + shortcut_rest

                return torch.cat([out_first, out_rest], dim=2)
            else:
                # Subsequent chunks with cache: spatio-temporal upsample
                h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
                shortcut = rearrange(x, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
                shortcut = shortcut.repeat_interleave(repeats=self.repeats, dim=1)
        else:
            h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
            shortcut = x.repeat_interleave(repeats=self.repeats, dim=1)
            shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)

        return h + shortcut

fastvideo.models.vaes.hyworldvae.HYWorldUpsample.forward

forward(x: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, first_chunk: bool = False)

Forward pass with optional temporal caching.

Parameters

x: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

first_chunk: bool

Whether this is the first chunk (affects temporal upsample behavior)

Source
def forward(
    self,
    x: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
    first_chunk: bool = False,
):
    """
    Forward pass with optional temporal caching.

    Args:
        x: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
        first_chunk: Whether this is the first chunk (affects temporal upsample behavior)
    """
    r1 = 2 if self.add_temporal_upsample else 1
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = x[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        h = self.conv(x, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        h = self.conv(x)

    if self.add_temporal_upsample:
        if first_chunk:
            # First chunk: only spatial upsample
            h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
            h = h[:, :h.shape[1] // 2]
            # Compute the shortcut part
            shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
            shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1)
        elif feat_cache is None and x.shape[2] > 1:
            # No cache and multiple frames: separate first frame and rest
            h_first = h[:, :, :1, :, :]
            h_rest = h[:, :, 1:, :, :]
            x_first = x[:, :, :1, :, :]
            x_rest = x[:, :, 1:, :, :]

            # First frame: only spatial upsample
            h_first = rearrange(h_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
            h_first = h_first[:, :h_first.shape[1] // 2]
            shortcut_first = rearrange(x_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2)
            shortcut_first = shortcut_first.repeat_interleave(repeats=self.repeats // 2, dim=1)
            out_first = h_first + shortcut_first

            # Remaining frames: spatio-temporal upsample
            h_rest = rearrange(h_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
            shortcut_rest = rearrange(x_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
            shortcut_rest = shortcut_rest.repeat_interleave(repeats=self.repeats, dim=1)
            out_rest = h_rest + shortcut_rest

            return torch.cat([out_first, out_rest], dim=2)
        else:
            # Subsequent chunks with cache: spatio-temporal upsample
            h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
            shortcut = rearrange(x, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
            shortcut = shortcut.repeat_interleave(repeats=self.repeats, dim=1)
    else:
        h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)
        shortcut = x.repeat_interleave(repeats=self.repeats, dim=1)
        shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2)

    return h + shortcut

fastvideo.models.vaes.hyworldvae.HYWorldDownsample

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

Hierarchical downsampling with temporal/spatial support and optional caching.

Source
class HYWorldDownsample(nn.Module):
    """Hierarchical downsampling with temporal/spatial support and optional caching."""

    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 = HYWorldCausalConv3d(in_channels, out_channels // factor, kernel_size=3)

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

    def forward(
        self,
        x: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
    ):
        """
        Forward pass with optional temporal caching.

        Args:
            x: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
        """
        r1 = 2 if self.add_temporal_downsample else 1

        # Apply conv with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = x[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            h = self.conv(x, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            h = self.conv(x)  # Change the channel, ready for spatial or temporal downsample

        if self.add_temporal_downsample:
            if x.shape[2] == 1:
                # Single frame: only spatial downsample
                h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2)
                h = torch.cat([h, h], dim=1)
                # Compute the shortcut part
                shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2)
                B, C, T, H, W = shortcut.shape
                shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2)
            else:
                # Multiple frames: full spatio-temporal downsample
                h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
                # Shortcut computation
                shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
                B, C, T, H, W = shortcut.shape
                shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)
        else:
            h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
            shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
            B, C, T, H, W = shortcut.shape
            shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)

        return h + shortcut

fastvideo.models.vaes.hyworldvae.HYWorldDownsample.forward

forward(x: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None)

Forward pass with optional temporal caching.

Parameters

x: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

Source
def forward(
    self,
    x: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
):
    """
    Forward pass with optional temporal caching.

    Args:
        x: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
    """
    r1 = 2 if self.add_temporal_downsample else 1

    # Apply conv with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = x[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        h = self.conv(x, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        h = self.conv(x)  # Change the channel, ready for spatial or temporal downsample

    if self.add_temporal_downsample:
        if x.shape[2] == 1:
            # Single frame: only spatial downsample
            h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2)
            h = torch.cat([h, h], dim=1)
            # Compute the shortcut part
            shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2)
            B, C, T, H, W = shortcut.shape
            shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2)
        else:
            # Multiple frames: full spatio-temporal downsample
            h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
            # Shortcut computation
            shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
            B, C, T, H, W = shortcut.shape
            shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)
    else:
        h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
        shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2)
        B, C, T, H, W = shortcut.shape
        shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2)

    return h + shortcut

fastvideo.models.vaes.hyworldvae.HYWorldResnetBlock

HYWorldResnetBlock(in_channels: int, out_channels: Optional[int] = None, non_linearity: str = 'swish')
Source
class HYWorldResnetBlock(nn.Module):

    def __init__(
        self,
        in_channels: int,
        out_channels: Optional[int] = None,
        non_linearity: str = "swish",
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels

        self.nonlinearity = get_act_fn(non_linearity)

        self.norm1 = HYWorldRMS_norm(in_channels, images=False)
        self.conv1 = HYWorldCausalConv3d(in_channels, out_channels, kernel_size=3)

        self.norm2 = HYWorldRMS_norm(out_channels, images=False)
        self.conv2 = HYWorldCausalConv3d(out_channels, out_channels, kernel_size=3)
        if in_channels != out_channels:
            self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)

    def forward(
        self,
        hidden_states: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal feature caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
        """
        residual = hidden_states

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

        # apply the feature cacheing mechanism
        if feat_cache is not None and feat_idx is not None:
            # Retrieve the current layer index.
            idx = feat_idx[0]

            # Clone the last CACHE_T frames from the current input to store for the next step.
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()

            # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames)
            # and we have a previous cache, prepend the last frame of the previous cache.
            # This ensures sufficient temporal context for the convolution kernel.
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                # actually this means chunk 1 after chunk 0
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )

            # Apply the convolution layer using `hidden_states`and the *previous* cached state.
            hidden_states = self.conv1(hidden_states, feat_cache[idx])

            # Update the cache for this layer with the newly prepared context ('cache_x').
            feat_cache[idx] = cache_x

            # Increment the global layer index.
            feat_idx[0] += 1
        else:
            hidden_states = self.conv1(hidden_states)

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

        # Second conv with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            hidden_states = self.conv2(hidden_states, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            hidden_states = self.conv2(hidden_states)

        if self.in_channels != self.out_channels:
            residual = self.nin_shortcut(residual)

        return hidden_states + residual

fastvideo.models.vaes.hyworldvae.HYWorldResnetBlock.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None) -> torch.Tensor

Forward pass with optional temporal feature caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
) -> torch.Tensor:
    """
    Forward pass with optional temporal feature caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
    """
    residual = hidden_states

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

    # apply the feature cacheing mechanism
    if feat_cache is not None and feat_idx is not None:
        # Retrieve the current layer index.
        idx = feat_idx[0]

        # Clone the last CACHE_T frames from the current input to store for the next step.
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()

        # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames)
        # and we have a previous cache, prepend the last frame of the previous cache.
        # This ensures sufficient temporal context for the convolution kernel.
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            # actually this means chunk 1 after chunk 0
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )

        # Apply the convolution layer using `hidden_states`and the *previous* cached state.
        hidden_states = self.conv1(hidden_states, feat_cache[idx])

        # Update the cache for this layer with the newly prepared context ('cache_x').
        feat_cache[idx] = cache_x

        # Increment the global layer index.
        feat_idx[0] += 1
    else:
        hidden_states = self.conv1(hidden_states)

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

    # Second conv with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        hidden_states = self.conv2(hidden_states, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        hidden_states = self.conv2(hidden_states)

    if self.in_channels != self.out_channels:
        residual = self.nin_shortcut(residual)

    return hidden_states + residual

fastvideo.models.vaes.hyworldvae.HYWorldMidBlock

HYWorldMidBlock(in_channels: int, num_layers: int = 1, add_attention: bool = True)

Mid block with attention and resnet blocks, with optional caching support.

Source
class HYWorldMidBlock(nn.Module):
    """Mid block with attention and resnet blocks, with optional caching support."""

    def __init__(
        self,
        in_channels: int,
        num_layers: int = 1,
        add_attention: bool = True,
    ) -> None:
        super().__init__()
        self.add_attention = add_attention

        # There is always at least one resnet
        resnets = [HYWorldResnetBlock(
            in_channels=in_channels,
            out_channels=in_channels,
        )]
        attentions = []

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

            resnets.append(HYWorldResnetBlock(
                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,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
        """
        hidden_states = self.resnets[0](hidden_states, feat_cache, feat_idx)

        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            if attn is not None:
                hidden_states = attn(hidden_states)
            hidden_states = resnet(hidden_states, feat_cache, feat_idx)

        return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldMidBlock.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
    """
    hidden_states = self.resnets[0](hidden_states, feat_cache, feat_idx)

    for attn, resnet in zip(self.attentions, self.resnets[1:]):
        if attn is not None:
            hidden_states = attn(hidden_states)
        hidden_states = resnet(hidden_states, feat_cache, feat_idx)

    return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldDownBlock3D

HYWorldDownBlock3D(in_channels: int, out_channels: int, num_layers: int = 1, downsample_out_channels: Optional[int] = None, add_temporal_downsample: int = True)

Down block with resnet blocks and optional downsampling, with caching support.

Source
class HYWorldDownBlock3D(nn.Module):
    """Down block with resnet blocks and optional downsampling, with caching support."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_layers: int = 1,
        downsample_out_channels: Optional[int] = None,
        add_temporal_downsample: int = True,
    ) -> None:
        super().__init__()
        resnets = []

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

        self.resnets = nn.ModuleList(resnets)

        if downsample_out_channels is not None:
            self.downsamplers = nn.ModuleList([
                HYWorldDownsample(
                    out_channels,
                    out_channels=downsample_out_channels,
                    add_temporal_downsample=add_temporal_downsample,
                )
            ])
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
        """
        for resnet in self.resnets:
            hidden_states = resnet(hidden_states, feat_cache, feat_idx)

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

        return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldDownBlock3D.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
    """
    for resnet in self.resnets:
        hidden_states = resnet(hidden_states, feat_cache, feat_idx)

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

    return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldUpBlock3D

HYWorldUpBlock3D(in_channels: int, out_channels: int, num_layers: int = 1, upsample_out_channels: Optional[int] = None, add_temporal_upsample: bool = True)

Up block with resnet blocks and optional upsampling, with caching support.

Source
class HYWorldUpBlock3D(nn.Module):
    """Up block with resnet blocks and optional upsampling, with caching support."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_layers: int = 1,
        upsample_out_channels: Optional[int] = None,
        add_temporal_upsample: bool = True,
    ) -> None:
        super().__init__()
        resnets = []

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

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

        self.resnets = nn.ModuleList(resnets)

        if upsample_out_channels is not None:
            self.upsamplers = nn.ModuleList([
                HYWorldUpsample(
                    out_channels,
                    out_channels=upsample_out_channels,
                    add_temporal_upsample=add_temporal_upsample,
                )
            ])
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
        first_chunk: bool = False,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
            first_chunk: Whether this is the first chunk (for upsampling behavior)
        """
        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, feat_cache, feat_idx)

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states, feat_cache, feat_idx, first_chunk)

        return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldUpBlock3D.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, first_chunk: bool = False) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

first_chunk: bool

Whether this is the first chunk (for upsampling behavior)

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
    first_chunk: bool = False,
) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
        first_chunk: Whether this is the first chunk (for upsampling behavior)
    """
    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, feat_cache, feat_idx)

    if self.upsamplers is not None:
        for upsampler in self.upsamplers:
            hidden_states = upsampler(hidden_states, feat_cache, feat_idx, first_chunk)

    return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldEncoder3D

HYWorldEncoder3D(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 with optional temporal caching.

Source
class HYWorldEncoder3D(nn.Module):
    r"""
    3D vae encoder for HunyuanImageRefiner with optional temporal caching.
    """

    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.block_out_channels = block_out_channels

        self.conv_in = HYWorldCausalConv3d(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 = HYWorldDownBlock3D(
                    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 = HYWorldDownBlock3D(
                    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 = HYWorldMidBlock(in_channels=block_out_channels[-1])

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

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
        """
        # conv_in with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            hidden_states = self.conv_in(hidden_states, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            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, feat_cache, feat_idx)
            hidden_states = self.mid_block(hidden_states, feat_cache, feat_idx)

        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)

        # conv_out with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            hidden_states = self.conv_out(hidden_states, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            hidden_states = self.conv_out(hidden_states)

        hidden_states += short_cut

        return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldEncoder3D.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
    """
    # conv_in with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        hidden_states = self.conv_in(hidden_states, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        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, feat_cache, feat_idx)
        hidden_states = self.mid_block(hidden_states, feat_cache, feat_idx)

    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)

    # conv_out with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        hidden_states = self.conv_out(hidden_states, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        hidden_states = self.conv_out(hidden_states)

    hidden_states += short_cut

    return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldDecoder3D

HYWorldDecoder3D(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 with optional temporal caching.

Source
class HYWorldDecoder3D(nn.Module):
    r"""
    Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner with optional temporal caching.
    """

    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.block_out_channels = block_out_channels
        self.repeat = block_out_channels[0] // self.in_channels

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

        # mid
        self.mid_block = HYWorldMidBlock(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 = HYWorldUpBlock3D(
                    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 = HYWorldUpBlock3D(
                    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 = HYWorldRMS_norm(block_out_channels[-1], images=False)
        self.conv_act = nn.SiLU()
        self.conv_out = HYWorldCausalConv3d(block_out_channels[-1], out_channels, kernel_size=3)

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
        feat_idx: Optional[List[int]] = None,
        first_chunk: bool = False,
    ) -> torch.Tensor:
        """
        Forward pass with optional temporal caching.

        Args:
            hidden_states: Input tensor of shape (B, C, T, H, W)
            feat_cache: List of cached features for each conv layer
            feat_idx: List containing current cache index [idx]
            first_chunk: Whether this is the first chunk (for upsampling behavior)
        """
        # conv_in with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            hidden_states = self.conv_in(hidden_states, feat_cache[idx]) + hidden_states.repeat_interleave(
                repeats=self.repeat, dim=1)
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            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, feat_cache, feat_idx)
            for up_block in self.up_blocks:
                hidden_states = up_block(hidden_states, feat_cache, feat_idx, first_chunk)

        # post-process
        hidden_states = self.norm_out(hidden_states)
        hidden_states = self.conv_act(hidden_states)

        # conv_out with caching
        if feat_cache is not None and feat_idx is not None:
            idx = feat_idx[0]
            cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                cache_x = torch.cat(
                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                    dim=2,
                )
            hidden_states = self.conv_out(hidden_states, feat_cache[idx])
            feat_cache[idx] = cache_x
            feat_idx[0] += 1
        else:
            hidden_states = self.conv_out(hidden_states)

        return hidden_states

fastvideo.models.vaes.hyworldvae.HYWorldDecoder3D.forward

forward(hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, first_chunk: bool = False) -> torch.Tensor

Forward pass with optional temporal caching.

Parameters

hidden_states: torch.Tensor

Input tensor of shape (B, C, T, H, W)

feat_cache: Optional[List[Optional[torch.Tensor]]]

List of cached features for each conv layer

feat_idx: Optional[List[int]]

List containing current cache index [idx]

first_chunk: bool

Whether this is the first chunk (for upsampling behavior)

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    feat_cache: Optional[List[Optional[torch.Tensor]]] = None,
    feat_idx: Optional[List[int]] = None,
    first_chunk: bool = False,
) -> torch.Tensor:
    """
    Forward pass with optional temporal caching.

    Args:
        hidden_states: Input tensor of shape (B, C, T, H, W)
        feat_cache: List of cached features for each conv layer
        feat_idx: List containing current cache index [idx]
        first_chunk: Whether this is the first chunk (for upsampling behavior)
    """
    # conv_in with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        hidden_states = self.conv_in(hidden_states, feat_cache[idx]) + hidden_states.repeat_interleave(
            repeats=self.repeat, dim=1)
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        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, feat_cache, feat_idx)
        for up_block in self.up_blocks:
            hidden_states = up_block(hidden_states, feat_cache, feat_idx, first_chunk)

    # post-process
    hidden_states = self.norm_out(hidden_states)
    hidden_states = self.conv_act(hidden_states)

    # conv_out with caching
    if feat_cache is not None and feat_idx is not None:
        idx = feat_idx[0]
        cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone()
        if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
            cache_x = torch.cat(
                [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
                dim=2,
            )
        hidden_states = self.conv_out(hidden_states, feat_cache[idx])
        feat_cache[idx] = cache_x
        feat_idx[0] += 1
    else:
        hidden_states = self.conv_out(hidden_states)

    return hidden_states

fastvideo.models.vaes.hyworldvae.AutoencoderKLHYWorld

AutoencoderKLHYWorld(config: Hunyuan15VAEConfig)

A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Revised from HunyuanVideo-1.5 with temporal caching support for HY-WorldPlay integration.

Source
class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE):
    r"""
    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
    Revised from HunyuanVideo-1.5 with temporal caching support for HY-WorldPlay integration.
    """

    _supports_gradient_checkpointing = True

    def __init__(
            self,
            config: Hunyuan15VAEConfig,  # HYWorld use the same VAE architecture as HunyuanVideo-1.5
    ) -> None:
        nn.Module.__init__(self)
        ParallelTiledVAE.__init__(self, config)

        self.encoder: Optional[HYWorldEncoder3D] = None
        self.decoder: Optional[HYWorldDecoder3D] = None

        if config.load_encoder:
            self.encoder = HYWorldEncoder3D(
                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 = HYWorldDecoder3D(
                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,
            )

        # TODO: Add spatial tiling.
        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

        # Cache-related attributes (initialized in clear_cache)
        self._conv_num: int = 0
        self._conv_idx: List[int] = [0]
        self._feat_map: List[Optional[torch.Tensor]] = []
        self._enc_conv_num: int = 0
        self._enc_conv_idx: List[int] = [0]
        self._enc_feat_map: List[Optional[torch.Tensor]] = []

        # Precompute and cache conv counts for encoder and decoder for clear_cache speedup
        self._cached_conv_counts = {
            "decoder": (sum(1 for m in self.decoder.modules()
                            if isinstance(m, HYWorldCausalConv3d)) if self.decoder is not None else 0),
            "encoder": (sum(1 for m in self.encoder.modules()
                            if isinstance(m, HYWorldCausalConv3d)) if self.encoder is not None else 0),
        }

    def clear_cache(self) -> None:
        """
        Initialize/clear the feature cache for chunk-based encoding/decoding.

        This should be called before starting a new encode/decode sequence to ensure
        the cache is properly initialized.
        """
        # Cache for decoder
        self._conv_num = self._cached_conv_counts["decoder"]
        self._conv_idx = [0]
        self._feat_map: List[Optional[torch.Tensor]] = [None] * self._conv_num

        # Cache for encoder
        self._enc_conv_num = self._cached_conv_counts["encoder"]
        self._enc_conv_idx = [0]
        self._enc_feat_map: List[Optional[torch.Tensor]] = [None] * self._enc_conv_num

    def _encode(self, x: torch.Tensor) -> torch.Tensor:
        """
        Encode video with temporal caching for chunk-based processing.

        This processes video in chunks (first frame separately, then 4 frames at a time)
        while maintaining temporal context through caching. This matches the HY-WorldPlay
        behavior for memory-efficient long video encoding.

        Args:
            x: Input video tensor of shape (B, C, T, H, W)

        Returns:
            Encoded latent tensor
        """
        assert self.encoder is not None, "Encoder not loaded"
        _, _, num_frame, _, _ = x.shape

        self.clear_cache()

        # Process in chunks: first frame alone, then groups of 4 frames
        iter_ = 1 + (num_frame - 1) // 4
        for i in range(iter_):
            self._enc_conv_idx = [0]
            if i == 0:
                # First frame
                out = self.encoder(
                    x[:, :, :1, :, :],
                    feat_cache=self._enc_feat_map,
                    feat_idx=self._enc_conv_idx,
                )
            else:
                # Subsequent frames in groups of 4
                out_ = self.encoder(
                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
                    feat_cache=self._enc_feat_map,
                    feat_idx=self._enc_conv_idx,
                )
                out = torch.cat([out, out_], dim=2)

        self.clear_cache()
        return out

    def _decode(self, z: torch.Tensor) -> torch.Tensor:
        """
        Decode latents with temporal caching for chunk-based processing.

        This processes latents one frame at a time while maintaining temporal context
        through caching. This matches the HY-WorldPlay behavior for memory-efficient
        long video decoding.

        Args:
            z: Latent tensor of shape (B, C, T, H, W)

        Returns:
            Decoded video tensor
        """
        assert self.decoder is not None, "Decoder not loaded"
        _, _, num_frame, _, _ = z.shape

        self.clear_cache()

        # Process one frame at a time with caching
        for i in range(num_frame):
            self._conv_idx = [0]
            if i == 0:
                # First frame with first_chunk=True
                out = self.decoder(
                    z[:, :, i:i + 1, :, :],
                    feat_cache=self._feat_map,
                    feat_idx=self._conv_idx,
                    first_chunk=True,
                )
            else:
                # Subsequent frames
                out_ = self.decoder(
                    z[:, :, i:i + 1, :, :],
                    feat_cache=self._feat_map,
                    feat_idx=self._conv_idx,
                    first_chunk=False,
                )
                out = torch.cat([out, out_], dim=2)

        self.clear_cache()
        return out

    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.
            generator (`torch.Generator`, *optional*):
                Generator for sampling.
        """
        x = sample

        # encode() uses temporal caching by default (via _encode)
        posterior = self.encode(x)

        z = posterior.sample(generator=generator) if sample_posterior else posterior.mode()

        # decode() uses temporal caching by default (via _decode)
        dec = self.decode(z)

        return dec

fastvideo.models.vaes.hyworldvae.AutoencoderKLHYWorld.clear_cache

clear_cache() -> None

Initialize/clear the feature cache for chunk-based encoding/decoding.

This should be called before starting a new encode/decode sequence to ensure the cache is properly initialized.

Source
def clear_cache(self) -> None:
    """
    Initialize/clear the feature cache for chunk-based encoding/decoding.

    This should be called before starting a new encode/decode sequence to ensure
    the cache is properly initialized.
    """
    # Cache for decoder
    self._conv_num = self._cached_conv_counts["decoder"]
    self._conv_idx = [0]
    self._feat_map: List[Optional[torch.Tensor]] = [None] * self._conv_num

    # Cache for encoder
    self._enc_conv_num = self._cached_conv_counts["encoder"]
    self._enc_conv_idx = [0]
    self._enc_feat_map: List[Optional[torch.Tensor]] = [None] * self._enc_conv_num

fastvideo.models.vaes.hyworldvae.AutoencoderKLHYWorld.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.

generator: `torch.Generator`, *optional*

Generator for sampling.

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.
        generator (`torch.Generator`, *optional*):
            Generator for sampling.
    """
    x = sample

    # encode() uses temporal caching by default (via _encode)
    posterior = self.encode(x)

    z = posterior.sample(generator=generator) if sample_posterior else posterior.mode()

    # decode() uses temporal caching by default (via _decode)
    dec = self.decode(z)

    return dec