mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-28 01:47:32 +08:00
1308 lines
50 KiB
Python
1308 lines
50 KiB
Python
import threading
|
|
import torch
|
|
from torch import nn
|
|
from functools import partial
|
|
import math
|
|
from einops import rearrange
|
|
from typing import List, Optional, Tuple, Union
|
|
from .conv_nd_factory import make_conv_nd, make_linear_nd
|
|
from .causal_conv3d import CausalConv3d
|
|
from .pixel_norm import PixelNorm
|
|
from ..model import PixArtAlphaCombinedTimestepSizeEmbeddings
|
|
import comfy.ops
|
|
import comfy.model_management
|
|
from comfy.ldm.modules.diffusionmodules.model import torch_cat_if_needed
|
|
|
|
ops = comfy.ops.disable_weight_init
|
|
|
|
def in_meta_context():
|
|
return torch.device("meta") == torch.empty(0).device
|
|
|
|
def mark_conv3d_ended(module):
|
|
tid = threading.get_ident()
|
|
for _, m in module.named_modules():
|
|
if isinstance(m, CausalConv3d):
|
|
current = m.temporal_cache_state.get(tid, (None, False))
|
|
m.temporal_cache_state[tid] = (current[0], True)
|
|
|
|
def split2(tensor, split_point, dim=2):
|
|
return torch.split(tensor, [split_point, tensor.shape[dim] - split_point], dim=dim)
|
|
|
|
def add_exchange_cache(dest, cache_in, new_input, dim=2):
|
|
if dest is not None:
|
|
if cache_in is not None:
|
|
cache_to_dest = min(dest.shape[dim], cache_in.shape[dim])
|
|
lead_in_dest, dest = split2(dest, cache_to_dest, dim=dim)
|
|
lead_in_source, cache_in = split2(cache_in, cache_to_dest, dim=dim)
|
|
lead_in_dest.add_(lead_in_source)
|
|
body, new_input = split2(new_input, dest.shape[dim], dim)
|
|
dest.add_(body)
|
|
return torch_cat_if_needed([cache_in, new_input], dim=dim)
|
|
|
|
class Encoder(nn.Module):
|
|
r"""
|
|
The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.
|
|
|
|
Args:
|
|
dims (`int` or `Tuple[int, int]`, *optional*, defaults to 3):
|
|
The number of dimensions to use in convolutions.
|
|
in_channels (`int`, *optional*, defaults to 3):
|
|
The number of input channels.
|
|
out_channels (`int`, *optional*, defaults to 3):
|
|
The number of output channels.
|
|
blocks (`List[Tuple[str, int]]`, *optional*, defaults to `[("res_x", 1)]`):
|
|
The blocks to use. Each block is a tuple of the block name and the number of layers.
|
|
base_channels (`int`, *optional*, defaults to 128):
|
|
The number of output channels for the first convolutional layer.
|
|
norm_num_groups (`int`, *optional*, defaults to 32):
|
|
The number of groups for normalization.
|
|
patch_size (`int`, *optional*, defaults to 1):
|
|
The patch size to use. Should be a power of 2.
|
|
norm_layer (`str`, *optional*, defaults to `group_norm`):
|
|
The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
|
latent_log_var (`str`, *optional*, defaults to `per_channel`):
|
|
The number of channels for the log variance. Can be either `per_channel`, `uniform`, `constant` or `none`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dims: Union[int, Tuple[int, int]] = 3,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
blocks: List[Tuple[str, int | dict]] = [("res_x", 1)],
|
|
base_channels: int = 128,
|
|
norm_num_groups: int = 32,
|
|
patch_size: Union[int, Tuple[int]] = 1,
|
|
norm_layer: str = "group_norm", # group_norm, pixel_norm
|
|
latent_log_var: str = "per_channel",
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.norm_layer = norm_layer
|
|
self.latent_channels = out_channels
|
|
self.latent_log_var = latent_log_var
|
|
self.blocks_desc = blocks
|
|
|
|
in_channels = in_channels * patch_size**2
|
|
output_channel = base_channels
|
|
|
|
self.conv_in = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=in_channels,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.down_blocks = nn.ModuleList([])
|
|
|
|
for block_name, block_params in blocks:
|
|
input_channel = output_channel
|
|
if isinstance(block_params, int):
|
|
block_params = {"num_layers": block_params}
|
|
|
|
if block_name == "res_x":
|
|
block = UNetMidBlock3D(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
num_layers=block_params["num_layers"],
|
|
resnet_eps=1e-6,
|
|
resnet_groups=norm_num_groups,
|
|
norm_layer=norm_layer,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "res_x_y":
|
|
output_channel = block_params.get("multiplier", 2) * output_channel
|
|
block = ResnetBlock3D(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
eps=1e-6,
|
|
groups=norm_num_groups,
|
|
norm_layer=norm_layer,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_time":
|
|
block = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=(2, 1, 1),
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_space":
|
|
block = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=(1, 2, 2),
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_all":
|
|
block = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=(2, 2, 2),
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_all_x_y":
|
|
output_channel = block_params.get("multiplier", 2) * output_channel
|
|
block = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=(2, 2, 2),
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_all_res":
|
|
output_channel = block_params.get("multiplier", 2) * output_channel
|
|
block = SpaceToDepthDownsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
stride=(2, 2, 2),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_space_res":
|
|
output_channel = block_params.get("multiplier", 2) * output_channel
|
|
block = SpaceToDepthDownsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
stride=(1, 2, 2),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_time_res":
|
|
output_channel = block_params.get("multiplier", 2) * output_channel
|
|
block = SpaceToDepthDownsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
stride=(2, 1, 1),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
else:
|
|
raise ValueError(f"unknown block: {block_name}")
|
|
|
|
self.down_blocks.append(block)
|
|
|
|
# out
|
|
if norm_layer == "group_norm":
|
|
self.conv_norm_out = nn.GroupNorm(
|
|
num_channels=output_channel, num_groups=norm_num_groups, eps=1e-6
|
|
)
|
|
elif norm_layer == "pixel_norm":
|
|
self.conv_norm_out = PixelNorm()
|
|
elif norm_layer == "layer_norm":
|
|
self.conv_norm_out = LayerNorm(output_channel, eps=1e-6)
|
|
|
|
self.conv_act = nn.SiLU()
|
|
|
|
conv_out_channels = out_channels
|
|
if latent_log_var == "per_channel":
|
|
conv_out_channels *= 2
|
|
elif latent_log_var == "uniform":
|
|
conv_out_channels += 1
|
|
elif latent_log_var == "constant":
|
|
conv_out_channels += 1
|
|
elif latent_log_var != "none":
|
|
raise ValueError(f"Invalid latent_log_var: {latent_log_var}")
|
|
self.conv_out = make_conv_nd(
|
|
dims,
|
|
output_channel,
|
|
conv_out_channels,
|
|
3,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def _forward_chunk(self, sample: torch.FloatTensor) -> Optional[torch.FloatTensor]:
|
|
sample = self.conv_in(sample)
|
|
|
|
checkpoint_fn = (
|
|
partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)
|
|
if self.gradient_checkpointing and self.training
|
|
else lambda x: x
|
|
)
|
|
|
|
for down_block in self.down_blocks:
|
|
sample = checkpoint_fn(down_block)(sample)
|
|
if sample is None or sample.shape[2] == 0:
|
|
return None
|
|
|
|
sample = self.conv_norm_out(sample)
|
|
sample = self.conv_act(sample)
|
|
sample = self.conv_out(sample)
|
|
if sample is None or sample.shape[2] == 0:
|
|
return None
|
|
|
|
if self.latent_log_var == "uniform":
|
|
last_channel = sample[:, -1:, ...]
|
|
num_dims = sample.dim()
|
|
|
|
if num_dims == 4:
|
|
# For shape (B, C, H, W)
|
|
repeated_last_channel = last_channel.repeat(
|
|
1, sample.shape[1] - 2, 1, 1
|
|
)
|
|
sample = torch.cat([sample, repeated_last_channel], dim=1)
|
|
elif num_dims == 5:
|
|
# For shape (B, C, F, H, W)
|
|
repeated_last_channel = last_channel.repeat(
|
|
1, sample.shape[1] - 2, 1, 1, 1
|
|
)
|
|
sample = torch.cat([sample, repeated_last_channel], dim=1)
|
|
else:
|
|
raise ValueError(f"Invalid input shape: {sample.shape}")
|
|
elif self.latent_log_var == "constant":
|
|
sample = sample[:, :-1, ...]
|
|
approx_ln_0 = (
|
|
-30
|
|
) # this is the minimal clamp value in DiagonalGaussianDistribution objects
|
|
sample = torch.cat(
|
|
[sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
|
|
dim=1,
|
|
)
|
|
|
|
return sample
|
|
|
|
def forward_orig(self, sample: torch.FloatTensor, device=None) -> torch.FloatTensor:
|
|
r"""The forward method of the `Encoder` class."""
|
|
|
|
max_chunk_size = get_max_chunk_size(sample.device if device is None else device) * 2 # encoder is more memory-efficient than decoder
|
|
frame_size = sample[:, :, :1, :, :].numel() * sample.element_size()
|
|
frame_size = int(frame_size * (self.conv_in.out_channels / self.conv_in.in_channels))
|
|
|
|
outputs = []
|
|
samples = [sample[:, :, :1, :, :]]
|
|
if sample.shape[2] > 1:
|
|
chunk_t = max(2, max_chunk_size // frame_size)
|
|
if chunk_t < 4:
|
|
chunk_t = 2
|
|
elif chunk_t < 8:
|
|
chunk_t = 4
|
|
else:
|
|
chunk_t = (chunk_t // 8) * 8
|
|
samples += list(torch.split(sample[:, :, 1:, :, :], chunk_t, dim=2))
|
|
for chunk_idx, chunk in enumerate(samples):
|
|
if chunk_idx == len(samples) - 1:
|
|
mark_conv3d_ended(self)
|
|
chunk = patchify(chunk, patch_size_hw=self.patch_size, patch_size_t=1).to(device=device)
|
|
output = self._forward_chunk(chunk)
|
|
if output is not None:
|
|
outputs.append(output)
|
|
|
|
return torch_cat_if_needed(outputs, dim=2)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
try:
|
|
return self.forward_orig(*args, **kwargs)
|
|
finally:
|
|
tid = threading.get_ident()
|
|
for _, module in self.named_modules():
|
|
# ComfyUI doesn't thread this kind of stuff today, but just in case
|
|
# we key on the thread to make it thread safe.
|
|
tid = threading.get_ident()
|
|
if hasattr(module, "temporal_cache_state"):
|
|
module.temporal_cache_state.pop(tid, None)
|
|
|
|
|
|
MIN_VRAM_FOR_CHUNK_SCALING = 6 * 1024 ** 3
|
|
MAX_VRAM_FOR_CHUNK_SCALING = 24 * 1024 ** 3
|
|
MIN_CHUNK_SIZE = 32 * 1024 ** 2
|
|
MAX_CHUNK_SIZE = 128 * 1024 ** 2
|
|
|
|
def get_max_chunk_size(device: torch.device) -> int:
|
|
total_memory = comfy.model_management.get_total_memory(dev=device)
|
|
|
|
if total_memory <= MIN_VRAM_FOR_CHUNK_SCALING:
|
|
return MIN_CHUNK_SIZE
|
|
if total_memory >= MAX_VRAM_FOR_CHUNK_SCALING:
|
|
return MAX_CHUNK_SIZE
|
|
|
|
interp = (total_memory - MIN_VRAM_FOR_CHUNK_SCALING) / (
|
|
MAX_VRAM_FOR_CHUNK_SCALING - MIN_VRAM_FOR_CHUNK_SCALING
|
|
)
|
|
return int(MIN_CHUNK_SIZE + interp * (MAX_CHUNK_SIZE - MIN_CHUNK_SIZE))
|
|
|
|
class Decoder(nn.Module):
|
|
r"""
|
|
The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.
|
|
|
|
Args:
|
|
dims (`int` or `Tuple[int, int]`, *optional*, defaults to 3):
|
|
The number of dimensions to use in convolutions.
|
|
in_channels (`int`, *optional*, defaults to 3):
|
|
The number of input channels.
|
|
out_channels (`int`, *optional*, defaults to 3):
|
|
The number of output channels.
|
|
blocks (`List[Tuple[str, int]]`, *optional*, defaults to `[("res_x", 1)]`):
|
|
The blocks to use. Each block is a tuple of the block name and the number of layers.
|
|
base_channels (`int`, *optional*, defaults to 128):
|
|
The number of output channels for the first convolutional layer.
|
|
norm_num_groups (`int`, *optional*, defaults to 32):
|
|
The number of groups for normalization.
|
|
patch_size (`int`, *optional*, defaults to 1):
|
|
The patch size to use. Should be a power of 2.
|
|
norm_layer (`str`, *optional*, defaults to `group_norm`):
|
|
The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
|
causal (`bool`, *optional*, defaults to `True`):
|
|
Whether to use causal convolutions or not.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dims,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
blocks: List[Tuple[str, int | dict]] = [("res_x", 1)],
|
|
base_channels: int = 128,
|
|
layers_per_block: int = 2,
|
|
norm_num_groups: int = 32,
|
|
patch_size: int = 1,
|
|
norm_layer: str = "group_norm",
|
|
causal: bool = True,
|
|
timestep_conditioning: bool = False,
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.layers_per_block = layers_per_block
|
|
out_channels = out_channels * patch_size**2
|
|
self.causal = causal
|
|
self.blocks_desc = blocks
|
|
|
|
# Compute output channel to be product of all channel-multiplier blocks
|
|
output_channel = base_channels
|
|
for block_name, block_params in list(reversed(blocks)):
|
|
block_params = block_params if isinstance(block_params, dict) else {}
|
|
if block_name == "res_x_y":
|
|
output_channel = output_channel * block_params.get("multiplier", 2)
|
|
if block_name == "compress_all":
|
|
output_channel = output_channel * block_params.get("multiplier", 1)
|
|
if block_name == "compress_space":
|
|
output_channel = output_channel * block_params.get("multiplier", 1)
|
|
if block_name == "compress_time":
|
|
output_channel = output_channel * block_params.get("multiplier", 1)
|
|
|
|
self.conv_in = make_conv_nd(
|
|
dims,
|
|
in_channels,
|
|
output_channel,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.up_blocks = nn.ModuleList([])
|
|
|
|
for block_name, block_params in list(reversed(blocks)):
|
|
input_channel = output_channel
|
|
if isinstance(block_params, int):
|
|
block_params = {"num_layers": block_params}
|
|
|
|
if block_name == "res_x":
|
|
block = UNetMidBlock3D(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
num_layers=block_params["num_layers"],
|
|
resnet_eps=1e-6,
|
|
resnet_groups=norm_num_groups,
|
|
norm_layer=norm_layer,
|
|
inject_noise=block_params.get("inject_noise", False),
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "res_x_y":
|
|
output_channel = output_channel // block_params.get("multiplier", 2)
|
|
block = ResnetBlock3D(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
eps=1e-6,
|
|
groups=norm_num_groups,
|
|
norm_layer=norm_layer,
|
|
inject_noise=block_params.get("inject_noise", False),
|
|
timestep_conditioning=False,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_time":
|
|
output_channel = output_channel // block_params.get("multiplier", 1)
|
|
block = DepthToSpaceUpsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
stride=(2, 1, 1),
|
|
out_channels_reduction_factor=block_params.get("multiplier", 1),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_space":
|
|
output_channel = output_channel // block_params.get("multiplier", 1)
|
|
block = DepthToSpaceUpsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
stride=(1, 2, 2),
|
|
out_channels_reduction_factor=block_params.get("multiplier", 1),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
elif block_name == "compress_all":
|
|
output_channel = output_channel // block_params.get("multiplier", 1)
|
|
block = DepthToSpaceUpsample(
|
|
dims=dims,
|
|
in_channels=input_channel,
|
|
stride=(2, 2, 2),
|
|
residual=block_params.get("residual", False),
|
|
out_channels_reduction_factor=block_params.get("multiplier", 1),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
else:
|
|
raise ValueError(f"unknown layer: {block_name}")
|
|
|
|
self.up_blocks.append(block)
|
|
|
|
if norm_layer == "group_norm":
|
|
self.conv_norm_out = nn.GroupNorm(
|
|
num_channels=output_channel, num_groups=norm_num_groups, eps=1e-6
|
|
)
|
|
elif norm_layer == "pixel_norm":
|
|
self.conv_norm_out = PixelNorm()
|
|
elif norm_layer == "layer_norm":
|
|
self.conv_norm_out = LayerNorm(output_channel, eps=1e-6)
|
|
|
|
self.conv_act = nn.SiLU()
|
|
self.conv_out = make_conv_nd(
|
|
dims,
|
|
output_channel,
|
|
out_channels,
|
|
3,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
# Precompute output scale factors: (channels, (t_scale, h_scale, w_scale), t_offset)
|
|
ts, hs, ws, to = 1, 1, 1, 0
|
|
for block in self.up_blocks:
|
|
if isinstance(block, DepthToSpaceUpsample):
|
|
ts *= block.stride[0]
|
|
hs *= block.stride[1]
|
|
ws *= block.stride[2]
|
|
if block.stride[0] > 1:
|
|
to = to * block.stride[0] + 1
|
|
self._output_scale = (out_channels // (patch_size ** 2), (ts, hs * patch_size, ws * patch_size), to)
|
|
|
|
self.timestep_conditioning = timestep_conditioning
|
|
|
|
if timestep_conditioning:
|
|
self.timestep_scale_multiplier = nn.Parameter(
|
|
torch.tensor(1000.0, dtype=torch.float32)
|
|
)
|
|
self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
|
|
output_channel * 2, 0, operations=ops,
|
|
)
|
|
self.last_scale_shift_table = nn.Parameter(torch.empty(2, output_channel))
|
|
else:
|
|
self.register_buffer(
|
|
"last_scale_shift_table",
|
|
torch.tensor(
|
|
[0.0, 0.0],
|
|
device="cpu" if in_meta_context() else None
|
|
).unsqueeze(1).expand(2, output_channel),
|
|
persistent=False,
|
|
)
|
|
|
|
|
|
def decode_output_shape(self, input_shape):
|
|
c, (ts, hs, ws), to = self._output_scale
|
|
return (input_shape[0], c, input_shape[2] * ts - to, input_shape[3] * hs, input_shape[4] * ws)
|
|
|
|
def run_up(self, idx, sample_ref, ended, timestep_shift_scale, scaled_timestep, checkpoint_fn, output_buffer, output_offset, max_chunk_size):
|
|
sample = sample_ref[0]
|
|
sample_ref[0] = None
|
|
if idx >= len(self.up_blocks):
|
|
sample = self.conv_norm_out(sample)
|
|
if timestep_shift_scale is not None:
|
|
shift, scale = timestep_shift_scale
|
|
sample = sample * (1 + scale) + shift
|
|
sample = self.conv_act(sample)
|
|
if ended:
|
|
mark_conv3d_ended(self.conv_out)
|
|
sample = self.conv_out(sample, causal=self.causal)
|
|
if sample is not None and sample.shape[2] > 0:
|
|
sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
|
|
t = sample.shape[2]
|
|
output_buffer[:, :, output_offset[0]:output_offset[0] + t].copy_(sample)
|
|
output_offset[0] += t
|
|
return
|
|
|
|
up_block = self.up_blocks[idx]
|
|
if ended:
|
|
mark_conv3d_ended(up_block)
|
|
if self.timestep_conditioning and isinstance(up_block, UNetMidBlock3D):
|
|
sample = checkpoint_fn(up_block)(
|
|
sample, causal=self.causal, timestep=scaled_timestep
|
|
)
|
|
else:
|
|
sample = checkpoint_fn(up_block)(sample, causal=self.causal)
|
|
|
|
if sample is None or sample.shape[2] == 0:
|
|
return
|
|
|
|
total_bytes = sample.numel() * sample.element_size()
|
|
num_chunks = (total_bytes + max_chunk_size - 1) // max_chunk_size
|
|
|
|
if num_chunks == 1:
|
|
# when we are not chunking, detach our x so the callee can free it as soon as they are done
|
|
next_sample_ref = [sample]
|
|
del sample
|
|
self.run_up(idx + 1, next_sample_ref, ended, timestep_shift_scale, scaled_timestep, checkpoint_fn, output_buffer, output_offset, max_chunk_size)
|
|
return
|
|
else:
|
|
samples = torch.chunk(sample, chunks=num_chunks, dim=2)
|
|
|
|
for chunk_idx, sample1 in enumerate(samples):
|
|
self.run_up(idx + 1, [sample1], ended and chunk_idx == len(samples) - 1, timestep_shift_scale, scaled_timestep, checkpoint_fn, output_buffer, output_offset, max_chunk_size)
|
|
|
|
def forward_orig(
|
|
self,
|
|
sample: torch.FloatTensor,
|
|
timestep: Optional[torch.Tensor] = None,
|
|
output_buffer: Optional[torch.Tensor] = None,
|
|
) -> torch.FloatTensor:
|
|
r"""The forward method of the `Decoder` class."""
|
|
batch_size = sample.shape[0]
|
|
|
|
mark_conv3d_ended(self.conv_in)
|
|
sample = self.conv_in(sample, causal=self.causal)
|
|
|
|
checkpoint_fn = (
|
|
partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)
|
|
if self.gradient_checkpointing and self.training
|
|
else lambda x: x
|
|
)
|
|
|
|
timestep_shift_scale = None
|
|
scaled_timestep = None
|
|
if self.timestep_conditioning:
|
|
assert (
|
|
timestep is not None
|
|
), "should pass timestep with timestep_conditioning=True"
|
|
scaled_timestep = timestep * self.timestep_scale_multiplier.to(dtype=sample.dtype, device=sample.device)
|
|
embedded_timestep = self.last_time_embedder(
|
|
timestep=scaled_timestep.flatten(),
|
|
resolution=None,
|
|
aspect_ratio=None,
|
|
batch_size=sample.shape[0],
|
|
hidden_dtype=sample.dtype,
|
|
)
|
|
embedded_timestep = embedded_timestep.view(
|
|
batch_size, embedded_timestep.shape[-1], 1, 1, 1
|
|
)
|
|
ada_values = self.last_scale_shift_table[
|
|
None, ..., None, None, None
|
|
].to(device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape(
|
|
batch_size,
|
|
2,
|
|
-1,
|
|
embedded_timestep.shape[-3],
|
|
embedded_timestep.shape[-2],
|
|
embedded_timestep.shape[-1],
|
|
)
|
|
timestep_shift_scale = ada_values.unbind(dim=1)
|
|
|
|
if output_buffer is None:
|
|
output_buffer = torch.empty(
|
|
self.decode_output_shape(sample.shape),
|
|
dtype=sample.dtype, device=comfy.model_management.intermediate_device(),
|
|
)
|
|
output_offset = [0]
|
|
|
|
max_chunk_size = get_max_chunk_size(sample.device)
|
|
|
|
self.run_up(0, [sample], True, timestep_shift_scale, scaled_timestep, checkpoint_fn, output_buffer, output_offset, max_chunk_size)
|
|
|
|
return output_buffer
|
|
|
|
def forward(self, *args, **kwargs):
|
|
try:
|
|
return self.forward_orig(*args, **kwargs)
|
|
finally:
|
|
for _, module in self.named_modules():
|
|
#ComfyUI doesn't thread this kind of stuff today, but just incase
|
|
#we key on the thread to make it thread safe.
|
|
tid = threading.get_ident()
|
|
if hasattr(module, "temporal_cache_state"):
|
|
module.temporal_cache_state.pop(tid, None)
|
|
|
|
|
|
class UNetMidBlock3D(nn.Module):
|
|
"""
|
|
A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.
|
|
|
|
Args:
|
|
in_channels (`int`): The number of input channels.
|
|
dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
|
|
num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
|
|
resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
|
|
resnet_groups (`int`, *optional*, defaults to 32):
|
|
The number of groups to use in the group normalization layers of the resnet blocks.
|
|
norm_layer (`str`, *optional*, defaults to `group_norm`):
|
|
The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
|
|
inject_noise (`bool`, *optional*, defaults to `False`):
|
|
Whether to inject noise into the hidden states.
|
|
timestep_conditioning (`bool`, *optional*, defaults to `False`):
|
|
Whether to condition the hidden states on the timestep.
|
|
|
|
Returns:
|
|
`torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
|
|
in_channels, height, width)`.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dims: Union[int, Tuple[int, int]],
|
|
in_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_groups: int = 32,
|
|
norm_layer: str = "group_norm",
|
|
inject_noise: bool = False,
|
|
timestep_conditioning: bool = False,
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
resnet_groups = (
|
|
resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
)
|
|
|
|
self.timestep_conditioning = timestep_conditioning
|
|
|
|
if timestep_conditioning:
|
|
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
|
|
in_channels * 4, 0, operations=ops,
|
|
)
|
|
|
|
self.res_blocks = nn.ModuleList(
|
|
[
|
|
ResnetBlock3D(
|
|
dims=dims,
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
norm_layer=norm_layer,
|
|
inject_noise=inject_noise,
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.FloatTensor,
|
|
causal: bool = True,
|
|
timestep: Optional[torch.Tensor] = None,
|
|
) -> torch.FloatTensor:
|
|
timestep_embed = None
|
|
if self.timestep_conditioning:
|
|
assert (
|
|
timestep is not None
|
|
), "should pass timestep with timestep_conditioning=True"
|
|
batch_size = hidden_states.shape[0]
|
|
timestep_embed = self.time_embedder(
|
|
timestep=timestep.flatten(),
|
|
resolution=None,
|
|
aspect_ratio=None,
|
|
batch_size=batch_size,
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
timestep_embed = timestep_embed.view(
|
|
batch_size, timestep_embed.shape[-1], 1, 1, 1
|
|
)
|
|
|
|
for resnet in self.res_blocks:
|
|
hidden_states = resnet(hidden_states, causal=causal, timestep=timestep_embed)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class SpaceToDepthDownsample(nn.Module):
|
|
def __init__(self, dims, in_channels, out_channels, stride, spatial_padding_mode):
|
|
super().__init__()
|
|
self.stride = stride
|
|
self.group_size = in_channels * math.prod(stride) // out_channels
|
|
self.conv = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels // math.prod(stride),
|
|
kernel_size=3,
|
|
stride=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
self.temporal_cache_state = {}
|
|
|
|
def forward(self, x, causal: bool = True):
|
|
tid = threading.get_ident()
|
|
cached, pad_first, cached_x, cached_input = self.temporal_cache_state.get(tid, (None, True, None, None))
|
|
if cached_input is not None:
|
|
x = torch_cat_if_needed([cached_input, x], dim=2)
|
|
cached_input = None
|
|
|
|
if self.stride[0] == 2 and pad_first:
|
|
x = torch.cat(
|
|
[x[:, :, :1, :, :], x], dim=2
|
|
) # duplicate first frames for padding
|
|
pad_first = False
|
|
|
|
if x.shape[2] < self.stride[0]:
|
|
cached_input = x
|
|
self.temporal_cache_state[tid] = (cached, pad_first, cached_x, cached_input)
|
|
return None
|
|
|
|
# skip connection
|
|
x_in = rearrange(
|
|
x,
|
|
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
|
|
p1=self.stride[0],
|
|
p2=self.stride[1],
|
|
p3=self.stride[2],
|
|
)
|
|
x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
|
|
x_in = x_in.mean(dim=2)
|
|
|
|
# conv
|
|
x = self.conv(x, causal=causal)
|
|
if self.stride[0] == 2 and x.shape[2] == 1:
|
|
if cached_x is not None:
|
|
x = torch_cat_if_needed([cached_x, x], dim=2)
|
|
cached_x = None
|
|
else:
|
|
cached_x = x
|
|
x = None
|
|
|
|
if x is not None:
|
|
x = rearrange(
|
|
x,
|
|
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
|
|
p1=self.stride[0],
|
|
p2=self.stride[1],
|
|
p3=self.stride[2],
|
|
)
|
|
|
|
cached = add_exchange_cache(x, cached, x_in, dim=2)
|
|
|
|
self.temporal_cache_state[tid] = (cached, pad_first, cached_x, cached_input)
|
|
|
|
return x
|
|
|
|
|
|
class DepthToSpaceUpsample(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dims,
|
|
in_channels,
|
|
stride,
|
|
residual=False,
|
|
out_channels_reduction_factor=1,
|
|
spatial_padding_mode="zeros",
|
|
):
|
|
super().__init__()
|
|
self.stride = stride
|
|
self.out_channels = (
|
|
math.prod(stride) * in_channels // out_channels_reduction_factor
|
|
)
|
|
self.conv = make_conv_nd(
|
|
dims=dims,
|
|
in_channels=in_channels,
|
|
out_channels=self.out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
self.residual = residual
|
|
self.out_channels_reduction_factor = out_channels_reduction_factor
|
|
self.temporal_cache_state = {}
|
|
|
|
def forward(self, x, causal: bool = True, timestep: Optional[torch.Tensor] = None):
|
|
tid = threading.get_ident()
|
|
cached, drop_first_conv, drop_first_res = self.temporal_cache_state.get(tid, (None, True, True))
|
|
y = self.conv(x, causal=causal)
|
|
y = rearrange(
|
|
y,
|
|
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
|
p1=self.stride[0],
|
|
p2=self.stride[1],
|
|
p3=self.stride[2],
|
|
)
|
|
if self.stride[0] == 2 and y.shape[2] > 0 and drop_first_conv:
|
|
y = y[:, :, 1:, :, :]
|
|
drop_first_conv = False
|
|
if self.residual:
|
|
# Reshape and duplicate the input to match the output shape
|
|
x_in = rearrange(
|
|
x,
|
|
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
|
p1=self.stride[0],
|
|
p2=self.stride[1],
|
|
p3=self.stride[2],
|
|
)
|
|
num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
|
|
x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
|
|
if self.stride[0] == 2 and x_in.shape[2] > 0 and drop_first_res:
|
|
x_in = x_in[:, :, 1:, :, :]
|
|
drop_first_res = False
|
|
|
|
if y.shape[2] == 0:
|
|
y = None
|
|
|
|
cached = add_exchange_cache(y, cached, x_in, dim=2)
|
|
self.temporal_cache_state[tid] = (cached, drop_first_conv, drop_first_res)
|
|
|
|
else:
|
|
self.temporal_cache_state[tid] = (None, drop_first_conv, False)
|
|
|
|
return y
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, dim, eps, elementwise_affine=True) -> None:
|
|
super().__init__()
|
|
self.norm = ops.LayerNorm(dim, eps=eps, elementwise_affine=elementwise_affine)
|
|
|
|
def forward(self, x):
|
|
x = rearrange(x, "b c d h w -> b d h w c")
|
|
x = self.norm(x)
|
|
x = rearrange(x, "b d h w c -> b c d h w")
|
|
return x
|
|
|
|
|
|
class ResnetBlock3D(nn.Module):
|
|
r"""
|
|
A Resnet block.
|
|
|
|
Parameters:
|
|
in_channels (`int`): The number of channels in the input.
|
|
out_channels (`int`, *optional*, default to be `None`):
|
|
The number of output channels for the first conv layer. If None, same as `in_channels`.
|
|
dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
|
|
groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
|
|
eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dims: Union[int, Tuple[int, int]],
|
|
in_channels: int,
|
|
out_channels: Optional[int] = None,
|
|
dropout: float = 0.0,
|
|
groups: int = 32,
|
|
eps: float = 1e-6,
|
|
norm_layer: str = "group_norm",
|
|
inject_noise: bool = False,
|
|
timestep_conditioning: bool = False,
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
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.inject_noise = inject_noise
|
|
|
|
if norm_layer == "group_norm":
|
|
self.norm1 = nn.GroupNorm(
|
|
num_groups=groups, num_channels=in_channels, eps=eps, affine=True
|
|
)
|
|
elif norm_layer == "pixel_norm":
|
|
self.norm1 = PixelNorm()
|
|
elif norm_layer == "layer_norm":
|
|
self.norm1 = LayerNorm(in_channels, eps=eps, elementwise_affine=True)
|
|
|
|
self.non_linearity = nn.SiLU()
|
|
|
|
self.conv1 = make_conv_nd(
|
|
dims,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
if inject_noise:
|
|
self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
|
|
|
|
if norm_layer == "group_norm":
|
|
self.norm2 = nn.GroupNorm(
|
|
num_groups=groups, num_channels=out_channels, eps=eps, affine=True
|
|
)
|
|
elif norm_layer == "pixel_norm":
|
|
self.norm2 = PixelNorm()
|
|
elif norm_layer == "layer_norm":
|
|
self.norm2 = LayerNorm(out_channels, eps=eps, elementwise_affine=True)
|
|
|
|
self.dropout = torch.nn.Dropout(dropout)
|
|
|
|
self.conv2 = make_conv_nd(
|
|
dims,
|
|
out_channels,
|
|
out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
causal=True,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
if inject_noise:
|
|
self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
|
|
|
|
self.conv_shortcut = (
|
|
make_linear_nd(
|
|
dims=dims, in_channels=in_channels, out_channels=out_channels
|
|
)
|
|
if in_channels != out_channels
|
|
else nn.Identity()
|
|
)
|
|
|
|
self.norm3 = (
|
|
LayerNorm(in_channels, eps=eps, elementwise_affine=True)
|
|
if in_channels != out_channels
|
|
else nn.Identity()
|
|
)
|
|
|
|
self.timestep_conditioning = timestep_conditioning
|
|
|
|
if timestep_conditioning:
|
|
self.scale_shift_table = nn.Parameter(
|
|
torch.randn(4, in_channels) / in_channels**0.5
|
|
)
|
|
else:
|
|
self.register_buffer(
|
|
"scale_shift_table",
|
|
torch.tensor(
|
|
[0.0, 0.0, 0.0, 0.0],
|
|
device="cpu" if in_meta_context() else None
|
|
).unsqueeze(1).expand(4, in_channels),
|
|
persistent=False,
|
|
)
|
|
|
|
self.temporal_cache_state={}
|
|
|
|
def _feed_spatial_noise(
|
|
self, hidden_states: torch.FloatTensor, per_channel_scale: torch.FloatTensor
|
|
) -> torch.FloatTensor:
|
|
spatial_shape = hidden_states.shape[-2:]
|
|
device = hidden_states.device
|
|
dtype = hidden_states.dtype
|
|
|
|
# similar to the "explicit noise inputs" method in style-gan
|
|
spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype)[None]
|
|
scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
|
|
hidden_states = hidden_states + scaled_noise
|
|
|
|
return hidden_states
|
|
|
|
def forward(
|
|
self,
|
|
input_tensor: torch.FloatTensor,
|
|
causal: bool = True,
|
|
timestep: Optional[torch.Tensor] = None,
|
|
) -> torch.FloatTensor:
|
|
hidden_states = input_tensor
|
|
batch_size = hidden_states.shape[0]
|
|
|
|
hidden_states = self.norm1(hidden_states)
|
|
if self.timestep_conditioning:
|
|
assert (
|
|
timestep is not None
|
|
), "should pass timestep with timestep_conditioning=True"
|
|
ada_values = self.scale_shift_table[
|
|
None, ..., None, None, None
|
|
].to(device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape(
|
|
batch_size,
|
|
4,
|
|
-1,
|
|
timestep.shape[-3],
|
|
timestep.shape[-2],
|
|
timestep.shape[-1],
|
|
)
|
|
shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
|
|
|
|
hidden_states = hidden_states * (1 + scale1) + shift1
|
|
|
|
hidden_states = self.non_linearity(hidden_states)
|
|
|
|
hidden_states = self.conv1(hidden_states, causal=causal)
|
|
|
|
if self.inject_noise:
|
|
hidden_states = self._feed_spatial_noise(
|
|
hidden_states, self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype)
|
|
)
|
|
|
|
hidden_states = self.norm2(hidden_states)
|
|
|
|
if self.timestep_conditioning:
|
|
hidden_states = hidden_states * (1 + scale2) + shift2
|
|
|
|
hidden_states = self.non_linearity(hidden_states)
|
|
|
|
hidden_states = self.dropout(hidden_states)
|
|
|
|
hidden_states = self.conv2(hidden_states, causal=causal)
|
|
|
|
if self.inject_noise:
|
|
hidden_states = self._feed_spatial_noise(
|
|
hidden_states, self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype)
|
|
)
|
|
|
|
input_tensor = self.norm3(input_tensor)
|
|
|
|
batch_size = input_tensor.shape[0]
|
|
|
|
input_tensor = self.conv_shortcut(input_tensor)
|
|
|
|
tid = threading.get_ident()
|
|
cached = self.temporal_cache_state.get(tid, None)
|
|
cached = add_exchange_cache(hidden_states, cached, input_tensor, dim=2)
|
|
self.temporal_cache_state[tid] = cached
|
|
|
|
return hidden_states
|
|
|
|
|
|
def patchify(x, patch_size_hw, patch_size_t=1):
|
|
if patch_size_hw == 1 and patch_size_t == 1:
|
|
return x
|
|
if x.dim() == 4:
|
|
x = rearrange(
|
|
x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw
|
|
)
|
|
elif x.dim() == 5:
|
|
x = rearrange(
|
|
x,
|
|
"b c (f p) (h q) (w r) -> b (c p r q) f h w",
|
|
p=patch_size_t,
|
|
q=patch_size_hw,
|
|
r=patch_size_hw,
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid input shape: {x.shape}")
|
|
|
|
return x
|
|
|
|
|
|
def unpatchify(x, patch_size_hw, patch_size_t=1):
|
|
if patch_size_hw == 1 and patch_size_t == 1:
|
|
return x
|
|
|
|
if x.dim() == 4:
|
|
x = rearrange(
|
|
x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw
|
|
)
|
|
elif x.dim() == 5:
|
|
x = rearrange(
|
|
x,
|
|
"b (c p r q) f h w -> b c (f p) (h q) (w r)",
|
|
p=patch_size_t,
|
|
q=patch_size_hw,
|
|
r=patch_size_hw,
|
|
)
|
|
|
|
return x
|
|
|
|
class processor(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.register_buffer("std-of-means", torch.empty(128))
|
|
self.register_buffer("mean-of-means", torch.empty(128))
|
|
|
|
def un_normalize(self, x):
|
|
return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
|
|
|
|
def normalize(self, x):
|
|
return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
|
|
|
|
class VideoVAE(nn.Module):
|
|
comfy_has_chunked_io = True
|
|
|
|
def __init__(self, version=0, config=None):
|
|
super().__init__()
|
|
|
|
if config is None:
|
|
config = self.get_default_config(version)
|
|
|
|
self.config = config
|
|
self.timestep_conditioning = config.get("timestep_conditioning", False)
|
|
self.decode_noise_scale = config.get("decode_noise_scale", 0.025)
|
|
self.decode_timestep = config.get("decode_timestep", 0.05)
|
|
double_z = config.get("double_z", True)
|
|
latent_log_var = config.get(
|
|
"latent_log_var", "per_channel" if double_z else "none"
|
|
)
|
|
|
|
self.encoder = Encoder(
|
|
dims=config["dims"],
|
|
in_channels=config.get("in_channels", 3),
|
|
out_channels=config["latent_channels"],
|
|
blocks=config.get("encoder_blocks", config.get("encoder_blocks", config.get("blocks"))),
|
|
patch_size=config.get("patch_size", 1),
|
|
latent_log_var=latent_log_var,
|
|
norm_layer=config.get("norm_layer", "group_norm"),
|
|
spatial_padding_mode=config.get("spatial_padding_mode", "zeros"),
|
|
base_channels=config.get("encoder_base_channels", 128),
|
|
)
|
|
|
|
self.decoder = Decoder(
|
|
dims=config["dims"],
|
|
in_channels=config["latent_channels"],
|
|
out_channels=config.get("out_channels", 3),
|
|
blocks=config.get("decoder_blocks", config.get("decoder_blocks", config.get("blocks"))),
|
|
base_channels=config.get("decoder_base_channels", 128),
|
|
patch_size=config.get("patch_size", 1),
|
|
norm_layer=config.get("norm_layer", "group_norm"),
|
|
causal=config.get("causal_decoder", False),
|
|
timestep_conditioning=self.timestep_conditioning,
|
|
spatial_padding_mode=config.get("spatial_padding_mode", "reflect"),
|
|
)
|
|
|
|
self.per_channel_statistics = processor()
|
|
|
|
def get_default_config(self, version):
|
|
if version == 0:
|
|
config = {
|
|
"_class_name": "CausalVideoAutoencoder",
|
|
"dims": 3,
|
|
"in_channels": 3,
|
|
"out_channels": 3,
|
|
"latent_channels": 128,
|
|
"blocks": [
|
|
["res_x", 4],
|
|
["compress_all", 1],
|
|
["res_x_y", 1],
|
|
["res_x", 3],
|
|
["compress_all", 1],
|
|
["res_x_y", 1],
|
|
["res_x", 3],
|
|
["compress_all", 1],
|
|
["res_x", 3],
|
|
["res_x", 4],
|
|
],
|
|
"scaling_factor": 1.0,
|
|
"norm_layer": "pixel_norm",
|
|
"patch_size": 4,
|
|
"latent_log_var": "uniform",
|
|
"use_quant_conv": False,
|
|
"causal_decoder": False,
|
|
}
|
|
elif version == 1:
|
|
config = {
|
|
"_class_name": "CausalVideoAutoencoder",
|
|
"dims": 3,
|
|
"in_channels": 3,
|
|
"out_channels": 3,
|
|
"latent_channels": 128,
|
|
"decoder_blocks": [
|
|
["res_x", {"num_layers": 5, "inject_noise": True}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 6, "inject_noise": True}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 7, "inject_noise": True}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 8, "inject_noise": False}]
|
|
],
|
|
"encoder_blocks": [
|
|
["res_x", {"num_layers": 4}],
|
|
["compress_all", {}],
|
|
["res_x_y", 1],
|
|
["res_x", {"num_layers": 3}],
|
|
["compress_all", {}],
|
|
["res_x_y", 1],
|
|
["res_x", {"num_layers": 3}],
|
|
["compress_all", {}],
|
|
["res_x", {"num_layers": 3}],
|
|
["res_x", {"num_layers": 4}]
|
|
],
|
|
"scaling_factor": 1.0,
|
|
"norm_layer": "pixel_norm",
|
|
"patch_size": 4,
|
|
"latent_log_var": "uniform",
|
|
"use_quant_conv": False,
|
|
"causal_decoder": False,
|
|
"timestep_conditioning": True,
|
|
}
|
|
else:
|
|
config = {
|
|
"_class_name": "CausalVideoAutoencoder",
|
|
"dims": 3,
|
|
"in_channels": 3,
|
|
"out_channels": 3,
|
|
"latent_channels": 128,
|
|
"encoder_blocks": [
|
|
["res_x", {"num_layers": 4}],
|
|
["compress_space_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 6}],
|
|
["compress_time_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 6}],
|
|
["compress_all_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 2}],
|
|
["compress_all_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 2}]
|
|
],
|
|
"decoder_blocks": [
|
|
["res_x", {"num_layers": 5, "inject_noise": False}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 5, "inject_noise": False}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 5, "inject_noise": False}],
|
|
["compress_all", {"residual": True, "multiplier": 2}],
|
|
["res_x", {"num_layers": 5, "inject_noise": False}]
|
|
],
|
|
"scaling_factor": 1.0,
|
|
"norm_layer": "pixel_norm",
|
|
"patch_size": 4,
|
|
"latent_log_var": "uniform",
|
|
"use_quant_conv": False,
|
|
"causal_decoder": False,
|
|
"timestep_conditioning": True
|
|
}
|
|
return config
|
|
|
|
def encode(self, x, device=None):
|
|
x = x[:, :, :max(1, 1 + ((x.shape[2] - 1) // 8) * 8), :, :]
|
|
means, logvar = torch.chunk(self.encoder(x, device=device), 2, dim=1)
|
|
return self.per_channel_statistics.normalize(means)
|
|
|
|
def decode_output_shape(self, input_shape):
|
|
return self.decoder.decode_output_shape(input_shape)
|
|
|
|
def decode(self, x, output_buffer=None):
|
|
if self.timestep_conditioning: #TODO: seed
|
|
x = torch.randn_like(x) * self.decode_noise_scale + (1.0 - self.decode_noise_scale) * x
|
|
return self.decoder(self.per_channel_statistics.un_normalize(x), timestep=self.decode_timestep, output_buffer=output_buffer)
|