5f95115e70
Signed-off-by: leffff <levnovitskiy@gmail.com> Signed-off-by: Artem Borisov <divotionchannel@gmail.com> Co-authored-by: leffff <levnovitskiy@gmail.com> Co-authored-by: Artem Borisov <divotionchannel@gmail.com>
1190 lines
45 KiB
Python
1190 lines
45 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Generated standalone Kandinsky 6 video VAE component for vLLM-Omni.
|
|
|
|
Identical to the Diffusers port's VAE module: ``AutoencoderKLHunyuanVideo``
|
|
only uses ``diffusers`` as a modeling-utility library (ModelMixin/
|
|
ConfigMixin checkpoint serialization), not for pipeline coupling, so it
|
|
needs no vLLM-specific override — vLLM-Omni's own native models (e.g.
|
|
Wan2.2) already import from ``diffusers`` the same way.
|
|
|
|
HunyuanVideo VAE with free-VRAM-aware spatial/temporal tiling, ported from
|
|
kandinsky-5-inference ``models/vae_orig.py`` (not the Diffusers
|
|
``AutoencoderKLHunyuanVideo``, which uses fixed tile sizes and peaks much higher).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from math import ceil, sqrt
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
|
from diffusers.models.activations import get_activation
|
|
from diffusers.models.attention_processor import Attention
|
|
from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution
|
|
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
|
from diffusers.models.modeling_utils import ModelMixin
|
|
from diffusers.utils.accelerate_utils import apply_forward_hook
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
|
DistributedOperator,
|
|
DistributedVaeMixin,
|
|
GridSpec,
|
|
TileTask,
|
|
)
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
|
|
def prepare_causal_attention_mask(f: int, s: int, dtype: torch.dtype, device: torch.device, b: int) -> torch.Tensor:
|
|
return (
|
|
torch.ones((f, f), dtype=dtype, device=device)
|
|
.tril_()
|
|
.log_()
|
|
.repeat_interleave(s, dim=0)
|
|
.repeat_interleave(s, dim=1)
|
|
.unsqueeze(0)
|
|
.expand(b, -1, -1)
|
|
.contiguous()
|
|
)
|
|
|
|
|
|
class HunyuanVideoCausalConv3d(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int | tuple[int, int, int] = 3,
|
|
stride: int | tuple[int, int, int] = 1,
|
|
padding: int | tuple[int, int, int] = 0,
|
|
dilation: int | tuple[int, int, int] = 1,
|
|
bias: bool = True,
|
|
pad_mode: str = "replicate",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
|
|
|
|
self.pad_mode = pad_mode
|
|
self.time_causal_padding = (
|
|
kernel_size[0] // 2,
|
|
kernel_size[0] // 2,
|
|
kernel_size[1] // 2,
|
|
kernel_size[1] // 2,
|
|
kernel_size[2] - 1,
|
|
0,
|
|
)
|
|
|
|
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode)
|
|
return self.conv(hidden_states)
|
|
|
|
|
|
class HunyuanVideoUpsampleCausal3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int | None = None,
|
|
kernel_size: int = 3,
|
|
stride: int = 1,
|
|
bias: bool = True,
|
|
upsample_factor: tuple[float, float, float] = (2, 2, 2),
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
out_channels = out_channels or in_channels
|
|
self.upsample_factor = upsample_factor
|
|
|
|
self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
num_frames = hidden_states.size(2)
|
|
dtp = hidden_states.dtype
|
|
first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2)
|
|
first_frame = (
|
|
F.interpolate(
|
|
first_frame.squeeze(2),
|
|
scale_factor=self.upsample_factor[1:],
|
|
mode="nearest",
|
|
)
|
|
.unsqueeze(2)
|
|
.to(dtp)
|
|
) # force cast
|
|
|
|
if num_frames > 1:
|
|
other_frames = other_frames.contiguous()
|
|
other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest").to(
|
|
dtp
|
|
) # force cast
|
|
hidden_states = torch.cat((first_frame, other_frames), dim=2)
|
|
del first_frame
|
|
del other_frames
|
|
if current_omni_platform.is_available():
|
|
current_omni_platform.empty_cache()
|
|
else:
|
|
hidden_states = first_frame
|
|
|
|
hidden_states = self.conv(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoDownsampleCausal3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
out_channels: int | None = None,
|
|
padding: int = 1,
|
|
kernel_size: int = 3,
|
|
bias: bool = True,
|
|
stride=2,
|
|
) -> None:
|
|
super().__init__()
|
|
out_channels = out_channels or channels
|
|
|
|
self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states = self.conv(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoResnetBlockCausal3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int | None = None,
|
|
dropout: float = 0.0,
|
|
groups: int = 32,
|
|
eps: float = 1e-6,
|
|
non_linearity: str = "swish",
|
|
) -> None:
|
|
super().__init__()
|
|
out_channels = out_channels or in_channels
|
|
|
|
self.nonlinearity = get_activation(non_linearity)
|
|
|
|
self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True)
|
|
self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0)
|
|
|
|
self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True)
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0)
|
|
|
|
self.conv_shortcut = None
|
|
if in_channels != out_channels:
|
|
self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
dtp = hidden_states.dtype
|
|
hidden_states = hidden_states.contiguous()
|
|
residual = hidden_states
|
|
|
|
hidden_states = self.norm1(hidden_states).to(dtp) # force cast
|
|
hidden_states = self.nonlinearity(hidden_states)
|
|
hidden_states = self.conv1(hidden_states)
|
|
|
|
hidden_states = self.norm2(hidden_states).to(dtp) # force cast
|
|
hidden_states = self.nonlinearity(hidden_states)
|
|
hidden_states = self.dropout(hidden_states)
|
|
hidden_states = self.conv2(hidden_states)
|
|
|
|
if self.conv_shortcut is not None:
|
|
residual = self.conv_shortcut(residual)
|
|
|
|
hidden_states = hidden_states + residual
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoMidBlock3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
add_attention: bool = True,
|
|
attention_head_dim: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
self.add_attention = add_attention
|
|
|
|
# There is always at least one resnet
|
|
resnets = [
|
|
HunyuanVideoResnetBlockCausal3D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
non_linearity=resnet_act_fn,
|
|
)
|
|
]
|
|
attentions = []
|
|
|
|
for _ in range(num_layers):
|
|
if self.add_attention:
|
|
attentions.append(
|
|
Attention(
|
|
in_channels,
|
|
heads=in_channels // attention_head_dim,
|
|
dim_head=attention_head_dim,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
residual_connection=True,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
_from_deprecated_attn_block=True,
|
|
)
|
|
)
|
|
else:
|
|
attentions.append(None)
|
|
|
|
resnets.append(
|
|
HunyuanVideoResnetBlockCausal3D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
non_linearity=resnet_act_fn,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states = self.resnets[0](hidden_states)
|
|
|
|
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
|
if attn is not None:
|
|
batch_size, _, num_frames, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3)
|
|
mask = prepare_causal_attention_mask(
|
|
num_frames,
|
|
height * width,
|
|
hidden_states.dtype,
|
|
hidden_states.device,
|
|
batch_size,
|
|
)
|
|
hidden_states = attn(hidden_states, attention_mask=mask)
|
|
hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3)
|
|
|
|
hidden_states = resnet(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoDownBlock3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
add_downsample: bool = True,
|
|
downsample_stride: int = 2,
|
|
downsample_padding: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
resnets = []
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
HunyuanVideoResnetBlockCausal3D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
non_linearity=resnet_act_fn,
|
|
)
|
|
)
|
|
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
HunyuanVideoDownsampleCausal3D(
|
|
out_channels,
|
|
out_channels=out_channels,
|
|
padding=downsample_padding,
|
|
stride=downsample_stride,
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
for resnet in self.resnets:
|
|
hidden_states = resnet(hidden_states)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoUpBlock3D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
add_upsample: bool = True,
|
|
upsample_scale_factor: tuple[int, int, int] = (2, 2, 2),
|
|
) -> None:
|
|
super().__init__()
|
|
resnets = []
|
|
|
|
for i in range(num_layers):
|
|
input_channels = in_channels if i == 0 else out_channels
|
|
|
|
resnets.append(
|
|
HunyuanVideoResnetBlockCausal3D(
|
|
in_channels=input_channels,
|
|
out_channels=out_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
non_linearity=resnet_act_fn,
|
|
)
|
|
)
|
|
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_upsample:
|
|
self.upsamplers = nn.ModuleList(
|
|
[
|
|
HunyuanVideoUpsampleCausal3D(
|
|
out_channels,
|
|
out_channels=out_channels,
|
|
upsample_factor=upsample_scale_factor,
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
for resnet in self.resnets:
|
|
hidden_states = resnet(hidden_states)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoEncoder3D(nn.Module):
|
|
r"""
|
|
Causal encoder for 3D video-like data introduced
|
|
in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
down_block_types: tuple[str, ...] = (
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
),
|
|
block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
|
|
layers_per_block: int = 2,
|
|
norm_num_groups: int = 32,
|
|
act_fn: str = "silu",
|
|
double_z: bool = True,
|
|
mid_block_add_attention=True,
|
|
temporal_compression_ratio: int = 4,
|
|
spatial_compression_ratio: int = 8,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
|
|
self.mid_block = None
|
|
self.down_blocks = nn.ModuleList([])
|
|
|
|
output_channel = block_out_channels[0]
|
|
for i, down_block_type in enumerate(down_block_types):
|
|
if down_block_type != "HunyuanVideoDownBlock3D":
|
|
raise ValueError(f"Unsupported down_block_type: {down_block_type}")
|
|
|
|
input_channel = output_channel
|
|
output_channel = block_out_channels[i]
|
|
is_final_block = i == len(block_out_channels) - 1
|
|
num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
|
|
num_time_downsample_layers = int(np.log2(temporal_compression_ratio))
|
|
|
|
if temporal_compression_ratio == 4:
|
|
add_spatial_downsample = bool(i < num_spatial_downsample_layers)
|
|
add_time_downsample = bool(
|
|
i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block
|
|
)
|
|
elif temporal_compression_ratio == 8:
|
|
add_spatial_downsample = bool(i < num_spatial_downsample_layers)
|
|
add_time_downsample = bool(i < num_time_downsample_layers)
|
|
else:
|
|
raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}")
|
|
|
|
downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
|
|
downsample_stride_T = (2,) if add_time_downsample else (1,)
|
|
downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
|
|
|
|
down_block = HunyuanVideoDownBlock3D(
|
|
num_layers=layers_per_block,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
add_downsample=bool(add_spatial_downsample or add_time_downsample),
|
|
resnet_eps=1e-6,
|
|
resnet_act_fn=act_fn,
|
|
resnet_groups=norm_num_groups,
|
|
downsample_stride=downsample_stride,
|
|
downsample_padding=0,
|
|
)
|
|
|
|
self.down_blocks.append(down_block)
|
|
|
|
self.mid_block = HunyuanVideoMidBlock3D(
|
|
in_channels=block_out_channels[-1],
|
|
resnet_eps=1e-6,
|
|
resnet_act_fn=act_fn,
|
|
attention_head_dim=block_out_channels[-1],
|
|
resnet_groups=norm_num_groups,
|
|
add_attention=mid_block_add_attention,
|
|
)
|
|
|
|
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
|
|
self.conv_act = nn.SiLU()
|
|
|
|
conv_out_channels = 2 * out_channels if double_z else out_channels
|
|
self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states = self.conv_in(hidden_states)
|
|
|
|
for down_block in self.down_blocks:
|
|
hidden_states = down_block(hidden_states)
|
|
|
|
hidden_states = self.mid_block(hidden_states)
|
|
|
|
hidden_states = self.conv_norm_out(hidden_states)
|
|
hidden_states = self.conv_act(hidden_states)
|
|
hidden_states = self.conv_out(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class HunyuanVideoDecoder3D(nn.Module):
|
|
r"""
|
|
Causal decoder for 3D video-like data introduced
|
|
in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
up_block_types: tuple[str, ...] = (
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
),
|
|
block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
|
|
layers_per_block: int = 2,
|
|
norm_num_groups: int = 32,
|
|
act_fn: str = "silu",
|
|
mid_block_add_attention=True,
|
|
time_compression_ratio: int = 4,
|
|
spatial_compression_ratio: int = 8,
|
|
):
|
|
super().__init__()
|
|
self.layers_per_block = layers_per_block
|
|
|
|
self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
|
|
self.up_blocks = nn.ModuleList([])
|
|
|
|
# mid
|
|
self.mid_block = HunyuanVideoMidBlock3D(
|
|
in_channels=block_out_channels[-1],
|
|
resnet_eps=1e-6,
|
|
resnet_act_fn=act_fn,
|
|
attention_head_dim=block_out_channels[-1],
|
|
resnet_groups=norm_num_groups,
|
|
add_attention=mid_block_add_attention,
|
|
)
|
|
|
|
# up
|
|
reversed_block_out_channels = list(reversed(block_out_channels))
|
|
output_channel = reversed_block_out_channels[0]
|
|
for i, up_block_type in enumerate(up_block_types):
|
|
if up_block_type != "HunyuanVideoUpBlock3D":
|
|
raise ValueError(f"Unsupported up_block_type: {up_block_type}")
|
|
|
|
prev_output_channel = output_channel
|
|
output_channel = reversed_block_out_channels[i]
|
|
is_final_block = i == len(block_out_channels) - 1
|
|
num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
|
|
num_time_upsample_layers = int(np.log2(time_compression_ratio))
|
|
|
|
if time_compression_ratio == 4:
|
|
add_spatial_upsample = bool(i < num_spatial_upsample_layers)
|
|
add_time_upsample = bool(
|
|
i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}")
|
|
|
|
upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
|
|
upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
|
|
upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
|
|
|
|
up_block = HunyuanVideoUpBlock3D(
|
|
num_layers=self.layers_per_block + 1,
|
|
in_channels=prev_output_channel,
|
|
out_channels=output_channel,
|
|
add_upsample=bool(add_spatial_upsample or add_time_upsample),
|
|
upsample_scale_factor=upsample_scale_factor,
|
|
resnet_eps=1e-6,
|
|
resnet_act_fn=act_fn,
|
|
resnet_groups=norm_num_groups,
|
|
)
|
|
|
|
self.up_blocks.append(up_block)
|
|
prev_output_channel = output_channel
|
|
|
|
# out
|
|
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
|
|
self.conv_act = nn.SiLU()
|
|
self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
dtp = hidden_states.dtype
|
|
hidden_states = self.conv_in(hidden_states)
|
|
|
|
hidden_states = self.mid_block(hidden_states)
|
|
|
|
for up_block in self.up_blocks:
|
|
hidden_states = up_block(hidden_states)
|
|
|
|
hidden_states = self.conv_norm_out(hidden_states)
|
|
hidden_states = self.conv_act(hidden_states).to(dtp) # force cast
|
|
hidden_states = self.conv_out(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AutoencoderKLHunyuanVideo(ModelMixin, ConfigMixin, DistributedVaeMixin):
|
|
r"""
|
|
A VAE model with KL loss for encoding videos into latents
|
|
and decoding latent representations into videos.
|
|
Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603).
|
|
|
|
This model inherits from [`ModelMixin`]. Check the superclass
|
|
documentation for it's generic methods implemented
|
|
for all models (such as downloading or saving).
|
|
"""
|
|
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
latent_channels: int = 16,
|
|
down_block_types: tuple[str, ...] = (
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
"HunyuanVideoDownBlock3D",
|
|
),
|
|
up_block_types: tuple[str, ...] = (
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
"HunyuanVideoUpBlock3D",
|
|
),
|
|
block_out_channels: tuple[int] = (128, 256, 512, 512),
|
|
layers_per_block: int = 2,
|
|
act_fn: str = "silu",
|
|
norm_num_groups: int = 32,
|
|
scaling_factor: float = 0.476986,
|
|
spatial_compression_ratio: int = 8,
|
|
temporal_compression_ratio: int = 4,
|
|
mid_block_add_attention: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.time_compression_ratio = temporal_compression_ratio
|
|
|
|
self.encoder = HunyuanVideoEncoder3D(
|
|
in_channels=in_channels,
|
|
out_channels=latent_channels,
|
|
down_block_types=down_block_types,
|
|
block_out_channels=block_out_channels,
|
|
layers_per_block=layers_per_block,
|
|
norm_num_groups=norm_num_groups,
|
|
act_fn=act_fn,
|
|
double_z=True,
|
|
mid_block_add_attention=mid_block_add_attention,
|
|
temporal_compression_ratio=temporal_compression_ratio,
|
|
spatial_compression_ratio=spatial_compression_ratio,
|
|
)
|
|
|
|
self.decoder = HunyuanVideoDecoder3D(
|
|
in_channels=latent_channels,
|
|
out_channels=out_channels,
|
|
up_block_types=up_block_types,
|
|
block_out_channels=block_out_channels,
|
|
layers_per_block=layers_per_block,
|
|
norm_num_groups=norm_num_groups,
|
|
act_fn=act_fn,
|
|
time_compression_ratio=temporal_compression_ratio,
|
|
spatial_compression_ratio=spatial_compression_ratio,
|
|
mid_block_add_attention=mid_block_add_attention,
|
|
)
|
|
|
|
self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
|
|
self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
|
|
|
|
self.spatial_compression_ratio = spatial_compression_ratio
|
|
self.temporal_compression_ratio = temporal_compression_ratio
|
|
|
|
self.use_slicing = False
|
|
|
|
self.use_tiling = True
|
|
|
|
self.use_framewise_encoding = True
|
|
self.use_framewise_decoding = True
|
|
|
|
self.tile_sample_min_height = 256
|
|
self.tile_sample_min_width = 256
|
|
self.tile_sample_min_num_frames = 16
|
|
|
|
self.tile_sample_stride_height = 192
|
|
self.tile_sample_stride_width = 192
|
|
self.tile_sample_stride_num_frames = 12
|
|
|
|
self.tile_size = None
|
|
self.distributed_executor = None
|
|
|
|
def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
|
|
if self.distributed_executor is None:
|
|
self.init_distributed()
|
|
super().set_parallel_size(parallel_size, mode=mode)
|
|
|
|
def _decode_tile_params(self) -> tuple[int, int, int, int]:
|
|
ratio = int(self.spatial_compression_ratio)
|
|
return (
|
|
self.tile_sample_min_height // ratio,
|
|
self.tile_sample_min_width // ratio,
|
|
self.tile_sample_stride_height // ratio,
|
|
self.tile_sample_stride_width // ratio,
|
|
)
|
|
|
|
def _decode_tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
|
"""Same spatial tiles as ``tiled_decode``, one task per tile."""
|
|
_, _, _, height, width = z.shape
|
|
tile_h, tile_w, stride_h, stride_w = self._decode_tile_params()
|
|
h_starts = list(range(0, height - tile_h + 1, stride_h))
|
|
w_starts = list(range(0, width - tile_w + 1, stride_w))
|
|
tasks: list[TileTask] = []
|
|
for row, i in enumerate(h_starts):
|
|
for col, j in enumerate(w_starts):
|
|
tile = z[:, :, :, i : i + tile_h, j : j + tile_w]
|
|
tasks.append(
|
|
TileTask(
|
|
tile_id=len(tasks),
|
|
grid_coord=(row, col),
|
|
tensor=tile,
|
|
workload=int(tile.shape[-2] * tile.shape[-1]),
|
|
)
|
|
)
|
|
spec = GridSpec(
|
|
split_dims=(3, 4),
|
|
grid_shape=(len(h_starts), len(w_starts)),
|
|
tile_spec={
|
|
"sample_height": height * int(self.spatial_compression_ratio),
|
|
"sample_width": width * int(self.spatial_compression_ratio),
|
|
},
|
|
output_dtype=z.dtype,
|
|
)
|
|
return tasks, spec
|
|
|
|
def _decode_tile_exec(self, task: TileTask) -> torch.Tensor:
|
|
return self.decoder(self.post_quant_conv(task.tensor))
|
|
|
|
def _decode_tile_merge(
|
|
self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
|
|
) -> torch.Tensor:
|
|
n_rows, n_cols = grid_spec.grid_shape
|
|
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
|
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
|
rows = [[coord_tensor_map[(i, j)] for j in range(n_cols)] for i in range(n_rows)]
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
|
|
if j > 0:
|
|
tile = self.blend_h(row[j - 1], tile, blend_width)
|
|
height_lim = self.tile_sample_min_height if i == n_rows - 1 else self.tile_sample_stride_height
|
|
width_lim = self.tile_sample_min_width if j == n_cols - 1 else self.tile_sample_stride_width
|
|
result_row.append(tile[:, :, :, :height_lim, :width_lim])
|
|
result_rows.append(torch.cat(result_row, dim=-1))
|
|
dec = torch.cat(result_rows, dim=3)
|
|
return dec[:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]]
|
|
|
|
def _encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
_, _, num_frames, height, width = x.shape
|
|
|
|
if self.use_framewise_decoding and num_frames > (self.tile_sample_min_num_frames + 1):
|
|
return self._temporal_tiled_encode(x)
|
|
|
|
if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
|
|
return self.tiled_encode(x)
|
|
|
|
x = self.encoder(x)
|
|
enc = self.quant_conv(x)
|
|
return enc
|
|
|
|
@apply_forward_hook
|
|
def encode(
|
|
self, x: torch.Tensor, opt_tiling: bool = True, return_dict: bool = True
|
|
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
|
|
r"""
|
|
Encode a batch of images into latents.
|
|
|
|
Args:
|
|
x (`torch.Tensor`): Input batch of images.
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`]
|
|
instead of a plain tuple.
|
|
|
|
Returns:
|
|
The latent representations of the encoded videos. If `return_dict` is True, a
|
|
[`~models.autoencoder_kl.AutoencoderKLOutput`] is returned,
|
|
otherwise a plain `tuple` is returned.
|
|
"""
|
|
if opt_tiling:
|
|
tile_size, tile_stride = self.get_enc_optimal_tiling(x.shape, device=x.device)
|
|
else:
|
|
b, _, f, h, w = x.shape
|
|
tile_size, tile_stride = (b, f, h, w), (f, h, w)
|
|
if tile_size != self.tile_size:
|
|
self.tile_size = tile_size
|
|
self.apply_tiling(tile_size, tile_stride)
|
|
|
|
h = self._encode(x)
|
|
|
|
posterior = DiagonalGaussianDistribution(h)
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
|
|
_, _, num_frames, height, width = z.shape
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
|
|
if self.use_framewise_decoding and num_frames > (tile_latent_min_num_frames + 1):
|
|
return self._temporal_tiled_decode(z, return_dict=return_dict)
|
|
|
|
if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
|
|
return self.tiled_decode(z, return_dict=return_dict)
|
|
|
|
z = self.post_quant_conv(z)
|
|
dec = self.decoder(z)
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
|
|
return DecoderOutput(sample=dec)
|
|
|
|
@apply_forward_hook
|
|
def decode(self, z: torch.Tensor, return_dict: bool = True, **ignore_kwargs) -> DecoderOutput | torch.Tensor:
|
|
r"""
|
|
Decode a batch of images.
|
|
|
|
Args:
|
|
z (`torch.Tensor`): Input batch of latent vectors.
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
|
|
|
|
Returns:
|
|
[`~models.vae.DecoderOutput`] or `tuple`:
|
|
If return_dict is True, a [`~models.vae.DecoderOutput`] is returned,
|
|
otherwise a plain `tuple` is returned.
|
|
"""
|
|
tile_size, tile_stride = self.get_dec_optimal_tiling(z.shape, device=z.device)
|
|
if tile_size != self.tile_size:
|
|
self.tile_size = tile_size
|
|
self.apply_tiling(tile_size, tile_stride)
|
|
|
|
decoded = self._decode(z).sample
|
|
|
|
if not return_dict:
|
|
return (decoded,)
|
|
|
|
return DecoderOutput(sample=decoded)
|
|
|
|
def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
|
|
for y in range(blend_extent):
|
|
b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
|
|
y / blend_extent
|
|
)
|
|
return b
|
|
|
|
def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
|
|
for x in range(blend_extent):
|
|
b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
|
|
x / blend_extent
|
|
)
|
|
return b
|
|
|
|
def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
|
|
for x in range(blend_extent):
|
|
b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * (
|
|
x / blend_extent
|
|
)
|
|
return b
|
|
|
|
def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
|
|
r"""Encode a batch of images using a tiled encoder.
|
|
|
|
Args:
|
|
x (`torch.Tensor`): Input batch of videos.
|
|
|
|
Returns:
|
|
`torch.Tensor`:
|
|
The latent representation of the encoded videos.
|
|
"""
|
|
_, _, _, height, width = x.shape
|
|
latent_height = height // self.spatial_compression_ratio
|
|
latent_width = width // self.spatial_compression_ratio
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
|
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
|
|
|
blend_height = tile_latent_min_height - tile_latent_stride_height
|
|
blend_width = tile_latent_min_width - tile_latent_stride_width
|
|
|
|
rows = []
|
|
for i in range(0, height - self.tile_sample_min_height + 1, self.tile_sample_stride_height):
|
|
row = []
|
|
for j in range(0, width - self.tile_sample_min_width + 1, self.tile_sample_stride_width):
|
|
tile = x[
|
|
:,
|
|
:,
|
|
:,
|
|
i : i + self.tile_sample_min_height,
|
|
j : j + self.tile_sample_min_width,
|
|
]
|
|
tile = self.encoder(tile).clone()
|
|
tile = self.quant_conv(tile)
|
|
row.append(tile)
|
|
rows.append(row)
|
|
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
|
|
if j > 0:
|
|
tile = self.blend_h(row[j - 1], tile, blend_width)
|
|
height_lim = tile_latent_min_height if i == len(rows) - 1 else tile_latent_stride_height
|
|
width_lim = tile_latent_min_width if j == len(row) - 1 else tile_latent_stride_width
|
|
result_row.append(tile[:, :, :, :height_lim, :width_lim])
|
|
result_rows.append(torch.cat(result_row, dim=4))
|
|
|
|
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
|
|
return enc
|
|
|
|
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
|
|
r"""
|
|
Decode a batch of images using a tiled decoder.
|
|
|
|
Args:
|
|
z (`torch.Tensor`): Input batch of latent vectors.
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
|
|
|
|
Returns:
|
|
[`~models.vae.DecoderOutput`] or `tuple`:
|
|
If return_dict is True, a [`~models.vae.DecoderOutput`] is returned,
|
|
otherwise a plain `tuple` is returned.
|
|
"""
|
|
|
|
_, _, _, height, width = z.shape
|
|
sample_height = height * self.spatial_compression_ratio
|
|
sample_width = width * self.spatial_compression_ratio
|
|
|
|
if self.distributed_executor is not None and self.is_distributed_enabled():
|
|
decoded = self.distributed_executor.execute(
|
|
z,
|
|
DistributedOperator(
|
|
split=self._decode_tile_split,
|
|
exec=self._decode_tile_exec,
|
|
merge=self._decode_tile_merge,
|
|
),
|
|
broadcast_result=True,
|
|
)
|
|
if not return_dict:
|
|
return (decoded,)
|
|
return DecoderOutput(sample=decoded)
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
|
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
|
|
|
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
|
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
|
|
|
rows = []
|
|
for i in range(0, height - tile_latent_min_height + 1, tile_latent_stride_height):
|
|
row = []
|
|
for j in range(0, width - tile_latent_min_width + 1, tile_latent_stride_width):
|
|
tile = z[
|
|
:,
|
|
:,
|
|
:,
|
|
i : i + tile_latent_min_height,
|
|
j : j + tile_latent_min_width,
|
|
]
|
|
tile = self.post_quant_conv(tile)
|
|
decoded = self.decoder(tile).clone()
|
|
row.append(decoded)
|
|
rows.append(row)
|
|
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
|
|
if j > 0:
|
|
tile = self.blend_h(row[j - 1], tile, blend_width)
|
|
height_lim = self.tile_sample_min_height if i == len(rows) - 1 else self.tile_sample_stride_height
|
|
width_lim = self.tile_sample_min_width if j == len(row) - 1 else self.tile_sample_stride_width
|
|
result_row.append(tile[:, :, :, :height_lim, :width_lim])
|
|
result_rows.append(torch.cat(result_row, dim=-1))
|
|
|
|
dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
return DecoderOutput(sample=dec)
|
|
|
|
def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
|
|
_, _, num_frames, height, width = x.shape
|
|
latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
|
|
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
|
|
blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames
|
|
|
|
row = []
|
|
# for i in range(0, num_frames, self.tile_sample_stride_num_frames):
|
|
for i in range(
|
|
0,
|
|
num_frames - self.tile_sample_min_num_frames + 1,
|
|
self.tile_sample_stride_num_frames,
|
|
):
|
|
tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
|
|
if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width):
|
|
tile = self.tiled_encode(tile)
|
|
else:
|
|
tile = self.encoder(tile).clone()
|
|
tile = self.quant_conv(tile)
|
|
if i > 0:
|
|
tile = tile[:, :, 1:, :, :]
|
|
row.append(tile)
|
|
|
|
result_row = []
|
|
for i, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_t(row[i - 1], tile, blend_num_frames)
|
|
t_lim = tile_latent_min_num_frames if i == len(row) - 1 else tile_latent_stride_num_frames
|
|
result_row.append(tile[:, :, :t_lim, :, :])
|
|
else:
|
|
result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
|
|
|
|
enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
|
|
return enc
|
|
|
|
def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
|
|
_, _, num_frames, _, _ = z.shape
|
|
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
|
|
blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
|
|
|
|
row = []
|
|
for i in range(
|
|
0,
|
|
num_frames - tile_latent_min_num_frames + 1,
|
|
tile_latent_stride_num_frames,
|
|
):
|
|
tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
|
|
if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height):
|
|
decoded = self.tiled_decode(tile, return_dict=True).sample
|
|
else:
|
|
tile = self.post_quant_conv(tile)
|
|
decoded = self.decoder(tile).clone()
|
|
if i > 0:
|
|
decoded = decoded[:, :, 1:, :, :]
|
|
row.append(decoded)
|
|
|
|
result_row = []
|
|
for i, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_t(row[i - 1], tile, blend_num_frames)
|
|
t_lim = self.tile_sample_min_num_frames if i == len(row) - 1 else self.tile_sample_stride_num_frames
|
|
result_row.append(tile[:, :, :t_lim, :, :])
|
|
else:
|
|
result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :])
|
|
|
|
dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames]
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
return DecoderOutput(sample=dec)
|
|
|
|
def forward(
|
|
self,
|
|
sample: torch.Tensor,
|
|
sample_posterior: bool = False,
|
|
return_dict: bool = True,
|
|
generator: torch.Generator | None = None,
|
|
) -> DecoderOutput | torch.Tensor:
|
|
r"""
|
|
Args:
|
|
sample (`torch.Tensor`): Input sample.
|
|
sample_posterior (`bool`, *optional*, defaults to `False`):
|
|
Whether to sample from the posterior.
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
|
|
"""
|
|
x = sample
|
|
posterior = self.encode(x).latent_dist
|
|
if sample_posterior:
|
|
z = posterior.sample(generator=generator)
|
|
else:
|
|
z = posterior.mode()
|
|
dec = self.decode(z, return_dict=return_dict)
|
|
return dec
|
|
|
|
def apply_tiling(self, tile: tuple[int, int, int, int], stride: tuple[int, int, int]):
|
|
"""Applies tiling."""
|
|
_, ft, ht, wt = tile
|
|
fs, hs, ws = stride
|
|
|
|
self.use_tiling = True
|
|
self.tile_sample_min_num_frames = ft - 1
|
|
self.tile_sample_stride_num_frames = fs
|
|
self.tile_sample_min_height = ht
|
|
self.tile_sample_min_width = wt
|
|
self.tile_sample_stride_height = hs
|
|
self.tile_sample_stride_width = ws
|
|
|
|
def get_enc_optimal_tiling(
|
|
self, shape: list[int], device: torch.device | str | None = None
|
|
) -> tuple[tuple[int, int, int, int], tuple[int, int, int]]:
|
|
"""Returns optimal tiling for given shape."""
|
|
h, w = shape[3:]
|
|
|
|
if device is None:
|
|
device = next(self.parameters()).device
|
|
device = torch.device(device)
|
|
free_mem = float(current_omni_platform.get_free_memory(device)) if device.type != "cpu" else float("inf")
|
|
executor = self.distributed_executor
|
|
if (
|
|
device.type == "cuda"
|
|
and executor is not None
|
|
and executor.parallel_size > 1
|
|
and dist.is_initialized()
|
|
and self.use_tiling
|
|
):
|
|
free = torch.tensor([int(free_mem)], device=device, dtype=torch.int64)
|
|
dist.all_reduce(free, op=dist.ReduceOp.MIN, group=executor.group)
|
|
free_mem = float(free.item())
|
|
max_area = free_mem / 256 / 17 / 8
|
|
num_vals = 256 * 17 * (h + 32) * (w + 32)
|
|
|
|
if h * w < max_area and num_vals < 2**31:
|
|
return (1, 17, h, w), (8, h, w)
|
|
|
|
def factorize(n, k):
|
|
a = sqrt(n / k)
|
|
b = sqrt(n * k)
|
|
return ceil(a), ceil(b)
|
|
|
|
k = max(h / w, w / h)
|
|
N = max(ceil(h * w / max_area), ceil(num_vals / 2**31))
|
|
a, b = factorize(N, k)
|
|
if h >= w:
|
|
wn, hn = a, b
|
|
else:
|
|
wn, hn = b, a
|
|
|
|
if wn > 1:
|
|
wt = ceil(w / wn / 8) * 8 + 16
|
|
ws = wt - 32
|
|
else:
|
|
wt = w
|
|
ws = w
|
|
if hn > 1:
|
|
ht = ceil(h / hn / 8) * 8 + 16
|
|
hs = ht - 32
|
|
else:
|
|
ht = h
|
|
hs = h
|
|
return (1, 17, ht, wt), (8, hs, ws)
|
|
|
|
def get_dec_optimal_tiling(
|
|
self, shape: list[int], device: torch.device | str | None = None
|
|
) -> tuple[tuple[int, int, int, int], tuple[int, int, int]]:
|
|
"""Returns optimal tiling for given shape."""
|
|
b, _, f, h, w = shape
|
|
enc_inp_shape = [b, 3, 4 * (f - 1) + 1, 8 * h, 8 * w]
|
|
return self.get_enc_optimal_tiling(enc_inp_shape, device=device)
|
|
|
|
|
|
__all__ = ["AutoencoderKLHunyuanVideo"]
|