|
|
from typing import Optional, Union
|
|
|
from pathlib import Path
|
|
|
import os
|
|
|
import json
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from einops import rearrange
|
|
|
from diffusers import ConfigMixin, ModelMixin
|
|
|
from safetensors.torch import safe_open
|
|
|
|
|
|
from ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND
|
|
|
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
|
def __init__(
|
|
|
self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
|
|
|
):
|
|
|
super().__init__()
|
|
|
if mid_channels is None:
|
|
|
mid_channels = channels
|
|
|
|
|
|
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
|
|
|
|
|
|
self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)
|
|
|
self.norm1 = nn.GroupNorm(32, mid_channels)
|
|
|
self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)
|
|
|
self.norm2 = nn.GroupNorm(32, channels)
|
|
|
self.activation = nn.SiLU()
|
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
residual = x
|
|
|
x = self.conv1(x)
|
|
|
x = self.norm1(x)
|
|
|
x = self.activation(x)
|
|
|
x = self.conv2(x)
|
|
|
x = self.norm2(x)
|
|
|
x = self.activation(x + residual)
|
|
|
return x
|
|
|
|
|
|
|
|
|
class LatentUpsampler(ModelMixin, ConfigMixin):
|
|
|
"""
|
|
|
Model to spatially upsample VAE latents.
|
|
|
|
|
|
Args:
|
|
|
in_channels (`int`): Number of channels in the input latent
|
|
|
mid_channels (`int`): Number of channels in the middle layers
|
|
|
num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
|
|
|
dims (`int`): Number of dimensions for convolutions (2 or 3)
|
|
|
spatial_upsample (`bool`): Whether to spatially upsample the latent
|
|
|
temporal_upsample (`bool`): Whether to temporally upsample the latent
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
in_channels: int = 128,
|
|
|
mid_channels: int = 512,
|
|
|
num_blocks_per_stage: int = 4,
|
|
|
dims: int = 3,
|
|
|
spatial_upsample: bool = True,
|
|
|
temporal_upsample: bool = False,
|
|
|
):
|
|
|
super().__init__()
|
|
|
|
|
|
self.in_channels = in_channels
|
|
|
self.mid_channels = mid_channels
|
|
|
self.num_blocks_per_stage = num_blocks_per_stage
|
|
|
self.dims = dims
|
|
|
self.spatial_upsample = spatial_upsample
|
|
|
self.temporal_upsample = temporal_upsample
|
|
|
|
|
|
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
|
|
|
|
|
|
self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)
|
|
|
self.initial_norm = nn.GroupNorm(32, mid_channels)
|
|
|
self.initial_activation = nn.SiLU()
|
|
|
|
|
|
self.res_blocks = nn.ModuleList(
|
|
|
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
|
|
|
)
|
|
|
|
|
|
if spatial_upsample and temporal_upsample:
|
|
|
self.upsampler = nn.Sequential(
|
|
|
nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
|
|
|
PixelShuffleND(3),
|
|
|
)
|
|
|
elif spatial_upsample:
|
|
|
self.upsampler = nn.Sequential(
|
|
|
nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
|
|
|
PixelShuffleND(2),
|
|
|
)
|
|
|
elif temporal_upsample:
|
|
|
self.upsampler = nn.Sequential(
|
|
|
nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
|
|
|
PixelShuffleND(1),
|
|
|
)
|
|
|
else:
|
|
|
raise ValueError(
|
|
|
"Either spatial_upsample or temporal_upsample must be True"
|
|
|
)
|
|
|
|
|
|
self.post_upsample_res_blocks = nn.ModuleList(
|
|
|
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
|
|
|
)
|
|
|
|
|
|
self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)
|
|
|
|
|
|
def forward(self, latent: torch.Tensor) -> torch.Tensor:
|
|
|
b, c, f, h, w = latent.shape
|
|
|
|
|
|
if self.dims == 2:
|
|
|
x = rearrange(latent, "b c f h w -> (b f) c h w")
|
|
|
x = self.initial_conv(x)
|
|
|
x = self.initial_norm(x)
|
|
|
x = self.initial_activation(x)
|
|
|
|
|
|
for block in self.res_blocks:
|
|
|
x = block(x)
|
|
|
|
|
|
x = self.upsampler(x)
|
|
|
|
|
|
for block in self.post_upsample_res_blocks:
|
|
|
x = block(x)
|
|
|
|
|
|
x = self.final_conv(x)
|
|
|
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
|
|
else:
|
|
|
x = self.initial_conv(latent)
|
|
|
x = self.initial_norm(x)
|
|
|
x = self.initial_activation(x)
|
|
|
|
|
|
for block in self.res_blocks:
|
|
|
x = block(x)
|
|
|
|
|
|
if self.temporal_upsample:
|
|
|
x = self.upsampler(x)
|
|
|
x = x[:, :, 1:, :, :]
|
|
|
else:
|
|
|
x = rearrange(x, "b c f h w -> (b f) c h w")
|
|
|
x = self.upsampler(x)
|
|
|
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
|
|
|
|
|
for block in self.post_upsample_res_blocks:
|
|
|
x = block(x)
|
|
|
|
|
|
x = self.final_conv(x)
|
|
|
|
|
|
return x
|
|
|
|
|
|
@classmethod
|
|
|
def from_config(cls, config):
|
|
|
return cls(
|
|
|
in_channels=config.get("in_channels", 4),
|
|
|
mid_channels=config.get("mid_channels", 128),
|
|
|
num_blocks_per_stage=config.get("num_blocks_per_stage", 4),
|
|
|
dims=config.get("dims", 2),
|
|
|
spatial_upsample=config.get("spatial_upsample", True),
|
|
|
temporal_upsample=config.get("temporal_upsample", False),
|
|
|
)
|
|
|
|
|
|
def config(self):
|
|
|
return {
|
|
|
"_class_name": "LatentUpsampler",
|
|
|
"in_channels": self.in_channels,
|
|
|
"mid_channels": self.mid_channels,
|
|
|
"num_blocks_per_stage": self.num_blocks_per_stage,
|
|
|
"dims": self.dims,
|
|
|
"spatial_upsample": self.spatial_upsample,
|
|
|
"temporal_upsample": self.temporal_upsample,
|
|
|
}
|
|
|
|
|
|
@classmethod
|
|
|
def from_pretrained(
|
|
|
cls,
|
|
|
pretrained_model_path: Optional[Union[str, os.PathLike]],
|
|
|
*args,
|
|
|
**kwargs,
|
|
|
):
|
|
|
pretrained_model_path = Path(pretrained_model_path)
|
|
|
if pretrained_model_path.is_file() and str(pretrained_model_path).endswith(
|
|
|
".safetensors"
|
|
|
):
|
|
|
state_dict = {}
|
|
|
with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
|
|
|
metadata = f.metadata()
|
|
|
for k in f.keys():
|
|
|
state_dict[k] = f.get_tensor(k)
|
|
|
config = json.loads(metadata["config"])
|
|
|
with torch.device("meta"):
|
|
|
latent_upsampler = LatentUpsampler.from_config(config)
|
|
|
latent_upsampler.load_state_dict(state_dict, assign=True)
|
|
|
return latent_upsampler
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
latent_upsampler = LatentUpsampler(num_blocks_per_stage=4, dims=3)
|
|
|
print(latent_upsampler)
|
|
|
total_params = sum(p.numel() for p in latent_upsampler.parameters())
|
|
|
print(f"Total number of parameters: {total_params:,}")
|
|
|
latent = torch.randn(1, 128, 9, 16, 16)
|
|
|
upsampled_latent = latent_upsampler(latent)
|
|
|
print(f"Upsampled latent shape: {upsampled_latent.shape}")
|
|
|
|
