Spaces:
Paused
Paused
import json | |
import os | |
from functools import partial | |
from types import SimpleNamespace | |
from typing import Any, Mapping, Optional, Tuple, Union, List | |
import torch | |
import numpy as np | |
from einops import rearrange | |
from torch import nn | |
from diffusers.utils import logging | |
from xora.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd | |
from xora.models.autoencoders.pixel_norm import PixelNorm | |
from xora.models.autoencoders.vae import AutoencoderKLWrapper | |
logger = logging.get_logger(__name__) # pylint: disable=invalid-name | |
class CausalVideoAutoencoder(AutoencoderKLWrapper): | |
def from_pretrained( | |
cls, | |
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], | |
*args, | |
**kwargs, | |
): | |
config_local_path = pretrained_model_name_or_path / "config.json" | |
config = cls.load_config(config_local_path, **kwargs) | |
video_vae = cls.from_config(config) | |
video_vae.to(kwargs["torch_dtype"]) | |
model_local_path = pretrained_model_name_or_path / "autoencoder.pth" | |
ckpt_state_dict = torch.load(model_local_path, map_location=torch.device("cpu")) | |
video_vae.load_state_dict(ckpt_state_dict) | |
statistics_local_path = ( | |
pretrained_model_name_or_path / "per_channel_statistics.json" | |
) | |
if statistics_local_path.exists(): | |
with open(statistics_local_path, "r") as file: | |
data = json.load(file) | |
transposed_data = list(zip(*data["data"])) | |
data_dict = { | |
col: torch.tensor(vals) | |
for col, vals in zip(data["columns"], transposed_data) | |
} | |
video_vae.register_buffer("std_of_means", data_dict["std-of-means"]) | |
video_vae.register_buffer( | |
"mean_of_means", | |
data_dict.get( | |
"mean-of-means", torch.zeros_like(data_dict["std-of-means"]) | |
), | |
) | |
return video_vae | |
def from_config(config): | |
assert ( | |
config["_class_name"] == "CausalVideoAutoencoder" | |
), "config must have _class_name=CausalVideoAutoencoder" | |
if isinstance(config["dims"], list): | |
config["dims"] = tuple(config["dims"]) | |
assert config["dims"] in [2, 3, (2, 1)], "dims must be 2, 3 or (2, 1)" | |
double_z = config.get("double_z", True) | |
latent_log_var = config.get( | |
"latent_log_var", "per_channel" if double_z else "none" | |
) | |
use_quant_conv = config.get("use_quant_conv", True) | |
if use_quant_conv and latent_log_var == "uniform": | |
raise ValueError("uniform latent_log_var requires use_quant_conv=False") | |
encoder = Encoder( | |
dims=config["dims"], | |
in_channels=config.get("in_channels", 3), | |
out_channels=config["latent_channels"], | |
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"), | |
) | |
decoder = Decoder( | |
dims=config["dims"], | |
in_channels=config["latent_channels"], | |
out_channels=config.get("out_channels", 3), | |
blocks=config.get("decoder_blocks", config.get("blocks")), | |
patch_size=config.get("patch_size", 1), | |
norm_layer=config.get("norm_layer", "group_norm"), | |
causal=config.get("causal_decoder", False), | |
) | |
dims = config["dims"] | |
return CausalVideoAutoencoder( | |
encoder=encoder, | |
decoder=decoder, | |
latent_channels=config["latent_channels"], | |
dims=dims, | |
use_quant_conv=use_quant_conv, | |
) | |
def config(self): | |
return SimpleNamespace( | |
_class_name="CausalVideoAutoencoder", | |
dims=self.dims, | |
in_channels=self.encoder.conv_in.in_channels // self.encoder.patch_size**2, | |
out_channels=self.decoder.conv_out.out_channels | |
// self.decoder.patch_size**2, | |
latent_channels=self.decoder.conv_in.in_channels, | |
encoder_blocks=self.encoder.blocks_desc, | |
decoder_blocks=self.decoder.blocks_desc, | |
scaling_factor=1.0, | |
norm_layer=self.encoder.norm_layer, | |
patch_size=self.encoder.patch_size, | |
latent_log_var=self.encoder.latent_log_var, | |
use_quant_conv=self.use_quant_conv, | |
causal_decoder=self.decoder.causal, | |
) | |
def is_video_supported(self): | |
""" | |
Check if the model supports video inputs of shape (B, C, F, H, W). Otherwise, the model only supports 2D images. | |
""" | |
return self.dims != 2 | |
def spatial_downscale_factor(self): | |
return ( | |
2 | |
** len( | |
[ | |
block | |
for block in self.encoder.blocks_desc | |
if block[0] in ["compress_space", "compress_all"] | |
] | |
) | |
* self.encoder.patch_size | |
) | |
def temporal_downscale_factor(self): | |
return 2 ** len( | |
[ | |
block | |
for block in self.encoder.blocks_desc | |
if block[0] in ["compress_time", "compress_all"] | |
] | |
) | |
def to_json_string(self) -> str: | |
import json | |
return json.dumps(self.config.__dict__) | |
def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): | |
per_channel_statistics_prefix = "per_channel_statistics." | |
ckpt_state_dict = { | |
key: value | |
for key, value in state_dict.items() | |
if not key.startswith(per_channel_statistics_prefix) | |
} | |
model_keys = set(name for name, _ in self.named_parameters()) | |
key_mapping = { | |
".resnets.": ".res_blocks.", | |
"downsamplers.0": "downsample", | |
"upsamplers.0": "upsample", | |
} | |
converted_state_dict = {} | |
for key, value in ckpt_state_dict.items(): | |
for k, v in key_mapping.items(): | |
key = key.replace(k, v) | |
if "norm" in key and key not in model_keys: | |
logger.info( | |
f"Removing key {key} from state_dict as it is not present in the model" | |
) | |
continue | |
converted_state_dict[key] = value | |
super().load_state_dict(converted_state_dict, strict=strict) | |
data_dict = { | |
key.removeprefix(per_channel_statistics_prefix): value | |
for key, value in state_dict.items() | |
if key.startswith(per_channel_statistics_prefix) | |
} | |
if len(data_dict) > 0: | |
self.register_buffer("std_of_means", data_dict["std-of-means"]) | |
self.register_buffer( | |
"mean_of_means", | |
data_dict.get( | |
"mean-of-means", torch.zeros_like(data_dict["std-of-means"]) | |
), | |
) | |
def last_layer(self): | |
if hasattr(self.decoder, "conv_out"): | |
if isinstance(self.decoder.conv_out, nn.Sequential): | |
last_layer = self.decoder.conv_out[-1] | |
else: | |
last_layer = self.decoder.conv_out | |
else: | |
last_layer = self.decoder.layers[-1] | |
return last_layer | |
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`, 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", | |
): | |
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, | |
) | |
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, | |
) | |
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, | |
) | |
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, | |
) | |
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, | |
) | |
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, | |
) | |
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, | |
) | |
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 != "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 | |
) | |
self.gradient_checkpointing = False | |
def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: | |
r"""The forward method of the `Encoder` class.""" | |
sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) | |
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) | |
sample = self.conv_norm_out(sample) | |
sample = self.conv_act(sample) | |
sample = self.conv_out(sample) | |
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}") | |
return sample | |
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, | |
): | |
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) | |
self.conv_in = make_conv_nd( | |
dims, | |
in_channels, | |
output_channel, | |
kernel_size=3, | |
stride=1, | |
padding=1, | |
causal=True, | |
) | |
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, | |
) | |
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, | |
) | |
elif block_name == "compress_time": | |
block = DepthToSpaceUpsample( | |
dims=dims, in_channels=input_channel, stride=(2, 1, 1) | |
) | |
elif block_name == "compress_space": | |
block = DepthToSpaceUpsample( | |
dims=dims, in_channels=input_channel, stride=(1, 2, 2) | |
) | |
elif block_name == "compress_all": | |
block = DepthToSpaceUpsample( | |
dims=dims, | |
in_channels=input_channel, | |
stride=(2, 2, 2), | |
residual=block_params.get("residual", False), | |
) | |
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 | |
) | |
self.gradient_checkpointing = False | |
def forward(self, sample: torch.FloatTensor, target_shape) -> torch.FloatTensor: | |
r"""The forward method of the `Decoder` class.""" | |
assert target_shape is not None, "target_shape must be provided" | |
sample = self.conv_in(sample, causal=self.causal) | |
upscale_dtype = next(iter(self.up_blocks.parameters())).dtype | |
checkpoint_fn = ( | |
partial(torch.utils.checkpoint.checkpoint, use_reentrant=False) | |
if self.gradient_checkpointing and self.training | |
else lambda x: x | |
) | |
sample = sample.to(upscale_dtype) | |
for up_block in self.up_blocks: | |
sample = checkpoint_fn(up_block)(sample, causal=self.causal) | |
sample = self.conv_norm_out(sample) | |
sample = self.conv_act(sample) | |
sample = self.conv_out(sample, causal=self.causal) | |
sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) | |
return sample | |
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. | |
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", | |
): | |
super().__init__() | |
resnet_groups = ( | |
resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) | |
) | |
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, | |
) | |
for _ in range(num_layers) | |
] | |
) | |
def forward( | |
self, hidden_states: torch.FloatTensor, causal: bool = True | |
) -> torch.FloatTensor: | |
for resnet in self.res_blocks: | |
hidden_states = resnet(hidden_states, causal=causal) | |
return hidden_states | |
class DepthToSpaceUpsample(nn.Module): | |
def __init__(self, dims, in_channels, stride, residual=False): | |
super().__init__() | |
self.stride = stride | |
self.out_channels = np.prod(stride) * in_channels | |
self.conv = make_conv_nd( | |
dims=dims, | |
in_channels=in_channels, | |
out_channels=self.out_channels, | |
kernel_size=3, | |
stride=1, | |
causal=True, | |
) | |
self.residual = residual | |
def forward(self, x, causal: bool = True): | |
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], | |
) | |
x_in = x_in.repeat(1, np.prod(self.stride), 1, 1, 1) | |
if self.stride[0] == 2: | |
x_in = x_in[:, :, 1:, :, :] | |
x = self.conv(x, causal=causal) | |
x = 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], | |
) | |
if self.stride[0] == 2: | |
x = x[:, :, 1:, :, :] | |
if self.residual: | |
x = x + x_in | |
return x | |
class LayerNorm(nn.Module): | |
def __init__(self, dim, eps, elementwise_affine=True) -> None: | |
super().__init__() | |
self.norm = nn.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", | |
): | |
super().__init__() | |
self.in_channels = in_channels | |
out_channels = in_channels if out_channels is None else out_channels | |
self.out_channels = out_channels | |
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, | |
) | |
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, | |
) | |
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() | |
) | |
def forward( | |
self, | |
input_tensor: torch.FloatTensor, | |
causal: bool = True, | |
) -> torch.FloatTensor: | |
hidden_states = input_tensor | |
hidden_states = self.norm1(hidden_states) | |
hidden_states = self.non_linearity(hidden_states) | |
hidden_states = self.conv1(hidden_states, causal=causal) | |
hidden_states = self.norm2(hidden_states) | |
hidden_states = self.non_linearity(hidden_states) | |
hidden_states = self.dropout(hidden_states) | |
hidden_states = self.conv2(hidden_states, causal=causal) | |
input_tensor = self.norm3(input_tensor) | |
input_tensor = self.conv_shortcut(input_tensor) | |
output_tensor = input_tensor + hidden_states | |
return output_tensor | |
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 | |
def create_video_autoencoder_config( | |
latent_channels: int = 64, | |
): | |
config = { | |
"_class_name": "CausalVideoAutoencoder", | |
"dims": 3, # (2, 1), # 2 for Conv2, 3 for Conv3d, (2, 1) for Conv2d followed by Conv1d | |
"in_channels": 3, # Number of input color channels (e.g., RGB) | |
"out_channels": 3, # Number of output color channels | |
"latent_channels": latent_channels, # Number of channels in the latent space representation | |
"blocks": [ | |
("res_x", 4), | |
("compress_space", 1), | |
("res_x_y", 1), | |
("res_x", 2), | |
("compress_all", 1), | |
("res_x", 3), | |
("compress_all", 1), | |
("res_x_y", 1), | |
("res_x", 2), | |
("compress_time", 1), | |
("res_x", 3), | |
("res_x", 3), | |
], | |
"patch_size": 4, | |
"latent_log_var": "uniform", | |
"use_quant_conv": False, | |
"norm_layer": "layer_norm", | |
"causal_decoder": True, | |
} | |
return config | |
def test_vae_patchify_unpatchify(): | |
import torch | |
x = torch.randn(2, 3, 8, 64, 64) | |
x_patched = patchify(x, patch_size_hw=4, patch_size_t=4) | |
x_unpatched = unpatchify(x_patched, patch_size_hw=4, patch_size_t=4) | |
assert torch.allclose(x, x_unpatched) | |
def demo_video_autoencoder_forward_backward(): | |
# Configuration for the VideoAutoencoder | |
config = create_video_autoencoder_config() | |
# Instantiate the VideoAutoencoder with the specified configuration | |
video_autoencoder = CausalVideoAutoencoder.from_config(config) | |
print(video_autoencoder) | |
video_autoencoder.eval() | |
# Print the total number of parameters in the video autoencoder | |
total_params = sum(p.numel() for p in video_autoencoder.parameters()) | |
print(f"Total number of parameters in VideoAutoencoder: {total_params:,}") | |
# Create a mock input tensor simulating a batch of videos | |
# Shape: (batch_size, channels, depth, height, width) | |
# E.g., 4 videos, each with 3 color channels, 16 frames, and 64x64 pixels per frame | |
input_videos = torch.randn(2, 3, 17, 64, 64) | |
# Forward pass: encode and decode the input videos | |
latent = video_autoencoder.encode(input_videos).latent_dist.mode() | |
print(f"input shape={input_videos.shape}") | |
print(f"latent shape={latent.shape}") | |
reconstructed_videos = video_autoencoder.decode( | |
latent, target_shape=input_videos.shape | |
).sample | |
print(f"reconstructed shape={reconstructed_videos.shape}") | |
# Validate that single image gets treated the same way as first frame | |
input_image = input_videos[:, :, :1, :, :] | |
image_latent = video_autoencoder.encode(input_image).latent_dist.mode() | |
reconstructed_image = video_autoencoder.decode( | |
image_latent, target_shape=image_latent.shape | |
).sample | |
first_frame_latent = latent[:, :, :1, :, :] | |
# assert torch.allclose(image_latent, first_frame_latent, atol=1e-6) | |
# assert torch.allclose(reconstructed_image, reconstructed_videos[:, :, :1, :, :], atol=1e-6) | |
assert (image_latent == first_frame_latent).all() | |
assert (reconstructed_image == reconstructed_videos[:, :, :1, :, :]).all() | |
# Calculate the loss (e.g., mean squared error) | |
loss = torch.nn.functional.mse_loss(input_videos, reconstructed_videos) | |
# Perform backward pass | |
loss.backward() | |
print(f"Demo completed with loss: {loss.item()}") | |
# Ensure to call the demo function to execute the forward and backward pass | |
if __name__ == "__main__": | |
demo_video_autoencoder_forward_backward() | |