hyworldvae ¶

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

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)

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

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

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

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

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

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

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