hyworldvae
fastvideo.models.vaes.hyworldvae.get_act_fn
get_act_fn(act_fn_name: str) -> nn.ModuleGet 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) -> ParallelTiledVAEfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.encode
encode(x: torch.Tensor) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.decode
decode(z: torch.Tensor) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_v
blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_h
blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.blend_t
blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.spatial_tiled_encode
spatial_tiled_encode(x: torch.Tensor) -> torch.TensorEncode a batch of images using a tiled encoder.
Parameters
x: `torch.Tensor`Input batch of videos.
Returns
torch.Tensortorch.Tensor: The latent representation of the encoded videos.
fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.parallel_tiled_decode
parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensorParallel 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.TensorDecode a batch of images using a tiled decoder.
Parameters
z: `torch.Tensor`Input batch of latent vectors.
Returns
torch.Tensortorch.Tensor: The decoded images.
fastvideo.models.vaes.hyworldvae.ParallelTiledVAE.tiled_encode
tiled_encode(x: torch.Tensor) -> torch.Tensorfastvideo.models.vaes.hyworldvae.ParallelTiledVAE.tiled_decode
tiled_decode(z: torch.Tensor) -> torch.Tensorfastvideo.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) -> NoneEnable 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() -> NoneDisable 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: intThe number of dimensions to normalize over.
channel_first: boolWhether the input tensor has channels as the first dimension. Default is True.
images: boolWhether the input represents image data. Default is True.
bias: boolWhether 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: intNumber of frames (temporal length).
n_hw: intProduct of height and width.
dtypeDesired mask dtype.
deviceDevice for the mask.
batch_size: intIf set, expands for batch.
Returns
torch.Tensor: Causal attention mask.
fastvideo.models.vaes.hyworldvae.HYWorldAttnBlock.forward
forward(x: torch.Tensor) -> torch.Tensorfastvideo.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)</code></pre></details>
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.TensorInput 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)</code></pre></details>
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</code></pre></details>
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.TensorInput 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: boolWhether 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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: boolWhether 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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</code></pre></details>
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</code></pre></details>
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.TensorInput 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: boolWhether 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</code></pre></details>
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</code></pre></details>
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</code></pre></details>
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</code></pre></details>