v5 release with triple architecture support and prompt enhancer
This commit is contained in:
DeepBeepMeep
0
ltx_video/models/__init__.py
Normal file
0
ltx_video/models/__init__.py
Normal file
0
ltx_video/models/autoencoders/__init__.py
Normal file
0
ltx_video/models/autoencoders/__init__.py
Normal file
63
ltx_video/models/autoencoders/causal_conv3d.py
Normal file
63
ltx_video/models/autoencoders/causal_conv3d.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class CausalConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size: int = 3,
|
||||
stride: Union[int, Tuple[int]] = 1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
spatial_padding_mode: str = "zeros",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
self.time_kernel_size = kernel_size[0]
|
||||
|
||||
dilation = (dilation, 1, 1)
|
||||
|
||||
height_pad = kernel_size[1] // 2
|
||||
width_pad = kernel_size[2] // 2
|
||||
padding = (0, height_pad, width_pad)
|
||||
|
||||
self.conv = nn.Conv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
padding_mode=spatial_padding_mode,
|
||||
groups=groups,
|
||||
)
|
||||
|
||||
def forward(self, x, causal: bool = True):
|
||||
if causal:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat(
|
||||
(1, 1, self.time_kernel_size - 1, 1, 1)
|
||||
)
|
||||
x = torch.concatenate((first_frame_pad, x), dim=2)
|
||||
else:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat(
|
||||
(1, 1, (self.time_kernel_size - 1) // 2, 1, 1)
|
||||
)
|
||||
last_frame_pad = x[:, :, -1:, :, :].repeat(
|
||||
(1, 1, (self.time_kernel_size - 1) // 2, 1, 1)
|
||||
)
|
||||
x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return self.conv.weight
|
||||
1405
ltx_video/models/autoencoders/causal_video_autoencoder.py
Normal file
1405
ltx_video/models/autoencoders/causal_video_autoencoder.py
Normal file
File diff suppressed because it is too large
Load Diff
90
ltx_video/models/autoencoders/conv_nd_factory.py
Normal file
90
ltx_video/models/autoencoders/conv_nd_factory.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_video.models.autoencoders.dual_conv3d import DualConv3d
|
||||
from ltx_video.models.autoencoders.causal_conv3d import CausalConv3d
|
||||
|
||||
|
||||
def make_conv_nd(
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
causal=False,
|
||||
spatial_padding_mode="zeros",
|
||||
temporal_padding_mode="zeros",
|
||||
):
|
||||
if not (spatial_padding_mode == temporal_padding_mode or causal):
|
||||
raise NotImplementedError("spatial and temporal padding modes must be equal")
|
||||
if dims == 2:
|
||||
return torch.nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif dims == 3:
|
||||
if causal:
|
||||
return CausalConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
return torch.nn.Conv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif dims == (2, 1):
|
||||
return DualConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def make_linear_nd(
|
||||
dims: int,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
bias=True,
|
||||
):
|
||||
if dims == 2:
|
||||
return torch.nn.Conv2d(
|
||||
in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
|
||||
)
|
||||
elif dims == 3 or dims == (2, 1):
|
||||
return torch.nn.Conv3d(
|
||||
in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
217
ltx_video/models/autoencoders/dual_conv3d.py
Normal file
217
ltx_video/models/autoencoders/dual_conv3d.py
Normal file
@@ -0,0 +1,217 @@
|
||||
import math
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class DualConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride: Union[int, Tuple[int, int, int]] = 1,
|
||||
padding: Union[int, Tuple[int, int, int]] = 0,
|
||||
dilation: Union[int, Tuple[int, int, int]] = 1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
):
|
||||
super(DualConv3d, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.padding_mode = padding_mode
|
||||
# Ensure kernel_size, stride, padding, and dilation are tuples of length 3
|
||||
if isinstance(kernel_size, int):
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
if kernel_size == (1, 1, 1):
|
||||
raise ValueError(
|
||||
"kernel_size must be greater than 1. Use make_linear_nd instead."
|
||||
)
|
||||
if isinstance(stride, int):
|
||||
stride = (stride, stride, stride)
|
||||
if isinstance(padding, int):
|
||||
padding = (padding, padding, padding)
|
||||
if isinstance(dilation, int):
|
||||
dilation = (dilation, dilation, dilation)
|
||||
|
||||
# Set parameters for convolutions
|
||||
self.groups = groups
|
||||
self.bias = bias
|
||||
|
||||
# Define the size of the channels after the first convolution
|
||||
intermediate_channels = (
|
||||
out_channels if in_channels < out_channels else in_channels
|
||||
)
|
||||
|
||||
# Define parameters for the first convolution
|
||||
self.weight1 = nn.Parameter(
|
||||
torch.Tensor(
|
||||
intermediate_channels,
|
||||
in_channels // groups,
|
||||
1,
|
||||
kernel_size[1],
|
||||
kernel_size[2],
|
||||
)
|
||||
)
|
||||
self.stride1 = (1, stride[1], stride[2])
|
||||
self.padding1 = (0, padding[1], padding[2])
|
||||
self.dilation1 = (1, dilation[1], dilation[2])
|
||||
if bias:
|
||||
self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
|
||||
else:
|
||||
self.register_parameter("bias1", None)
|
||||
|
||||
# Define parameters for the second convolution
|
||||
self.weight2 = nn.Parameter(
|
||||
torch.Tensor(
|
||||
out_channels, intermediate_channels // groups, kernel_size[0], 1, 1
|
||||
)
|
||||
)
|
||||
self.stride2 = (stride[0], 1, 1)
|
||||
self.padding2 = (padding[0], 0, 0)
|
||||
self.dilation2 = (dilation[0], 1, 1)
|
||||
if bias:
|
||||
self.bias2 = nn.Parameter(torch.Tensor(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias2", None)
|
||||
|
||||
# Initialize weights and biases
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.kaiming_uniform_(self.weight1, a=math.sqrt(5))
|
||||
nn.init.kaiming_uniform_(self.weight2, a=math.sqrt(5))
|
||||
if self.bias:
|
||||
fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
|
||||
bound1 = 1 / math.sqrt(fan_in1)
|
||||
nn.init.uniform_(self.bias1, -bound1, bound1)
|
||||
fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
|
||||
bound2 = 1 / math.sqrt(fan_in2)
|
||||
nn.init.uniform_(self.bias2, -bound2, bound2)
|
||||
|
||||
def forward(self, x, use_conv3d=False, skip_time_conv=False):
|
||||
if use_conv3d:
|
||||
return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
|
||||
else:
|
||||
return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
|
||||
|
||||
def forward_with_3d(self, x, skip_time_conv):
|
||||
# First convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight1,
|
||||
self.bias1,
|
||||
self.stride1,
|
||||
self.padding1,
|
||||
self.dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
if skip_time_conv:
|
||||
return x
|
||||
|
||||
# Second convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight2,
|
||||
self.bias2,
|
||||
self.stride2,
|
||||
self.padding2,
|
||||
self.dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_with_2d(self, x, skip_time_conv):
|
||||
b, c, d, h, w = x.shape
|
||||
|
||||
# First 2D convolution
|
||||
x = rearrange(x, "b c d h w -> (b d) c h w")
|
||||
# Squeeze the depth dimension out of weight1 since it's 1
|
||||
weight1 = self.weight1.squeeze(2)
|
||||
# Select stride, padding, and dilation for the 2D convolution
|
||||
stride1 = (self.stride1[1], self.stride1[2])
|
||||
padding1 = (self.padding1[1], self.padding1[2])
|
||||
dilation1 = (self.dilation1[1], self.dilation1[2])
|
||||
x = F.conv2d(
|
||||
x,
|
||||
weight1,
|
||||
self.bias1,
|
||||
stride1,
|
||||
padding1,
|
||||
dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
_, _, h, w = x.shape
|
||||
|
||||
if skip_time_conv:
|
||||
x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
|
||||
return x
|
||||
|
||||
# Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
|
||||
x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
|
||||
|
||||
# Reshape weight2 to match the expected dimensions for conv1d
|
||||
weight2 = self.weight2.squeeze(-1).squeeze(-1)
|
||||
# Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
|
||||
stride2 = self.stride2[0]
|
||||
padding2 = self.padding2[0]
|
||||
dilation2 = self.dilation2[0]
|
||||
x = F.conv1d(
|
||||
x,
|
||||
weight2,
|
||||
self.bias2,
|
||||
stride2,
|
||||
padding2,
|
||||
dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
|
||||
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return self.weight2
|
||||
|
||||
|
||||
def test_dual_conv3d_consistency():
|
||||
# Initialize parameters
|
||||
in_channels = 3
|
||||
out_channels = 5
|
||||
kernel_size = (3, 3, 3)
|
||||
stride = (2, 2, 2)
|
||||
padding = (1, 1, 1)
|
||||
|
||||
# Create an instance of the DualConv3d class
|
||||
dual_conv3d = DualConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
# Example input tensor
|
||||
test_input = torch.randn(1, 3, 10, 10, 10)
|
||||
|
||||
# Perform forward passes with both 3D and 2D settings
|
||||
output_conv3d = dual_conv3d(test_input, use_conv3d=True)
|
||||
output_2d = dual_conv3d(test_input, use_conv3d=False)
|
||||
|
||||
# Assert that the outputs from both methods are sufficiently close
|
||||
assert torch.allclose(
|
||||
output_conv3d, output_2d, atol=1e-6
|
||||
), "Outputs are not consistent between 3D and 2D convolutions."
|
||||
203
ltx_video/models/autoencoders/latent_upsampler.py
Normal file
203
ltx_video/models/autoencoders/latent_upsampler.py
Normal file
@@ -0,0 +1,203 @@
|
||||
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}")
|
||||
12
ltx_video/models/autoencoders/pixel_norm.py
Normal file
12
ltx_video/models/autoencoders/pixel_norm.py
Normal file
@@ -0,0 +1,12 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class PixelNorm(nn.Module):
|
||||
def __init__(self, dim=1, eps=1e-8):
|
||||
super(PixelNorm, self).__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x):
|
||||
return x / torch.sqrt(torch.mean(x**2, dim=self.dim, keepdim=True) + self.eps)
|
||||
33
ltx_video/models/autoencoders/pixel_shuffle.py
Normal file
33
ltx_video/models/autoencoders/pixel_shuffle.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class PixelShuffleND(nn.Module):
|
||||
def __init__(self, dims, upscale_factors=(2, 2, 2)):
|
||||
super().__init__()
|
||||
assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
|
||||
self.dims = dims
|
||||
self.upscale_factors = upscale_factors
|
||||
|
||||
def forward(self, x):
|
||||
if self.dims == 3:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
p3=self.upscale_factors[2],
|
||||
)
|
||||
elif self.dims == 2:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2) h w -> b c (h p1) (w p2)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
)
|
||||
elif self.dims == 1:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1) f h w -> b c (f p1) h w",
|
||||
p1=self.upscale_factors[0],
|
||||
)
|
||||
443
ltx_video/models/autoencoders/vae.py
Normal file
443
ltx_video/models/autoencoders/vae.py
Normal file
@@ -0,0 +1,443 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import inspect
|
||||
import math
|
||||
import torch.nn as nn
|
||||
from diffusers import ConfigMixin, ModelMixin
|
||||
from diffusers.models.autoencoders.vae import (
|
||||
DecoderOutput,
|
||||
DiagonalGaussianDistribution,
|
||||
)
|
||||
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
||||
from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd
|
||||
|
||||
|
||||
class AutoencoderKLWrapper(ModelMixin, ConfigMixin):
|
||||
"""Variational Autoencoder (VAE) model with KL loss.
|
||||
|
||||
VAE from the paper Auto-Encoding Variational Bayes by Diederik P. Kingma and Max Welling.
|
||||
This model is a wrapper around an encoder and a decoder, and it adds a KL loss term to the reconstruction loss.
|
||||
|
||||
Args:
|
||||
encoder (`nn.Module`):
|
||||
Encoder module.
|
||||
decoder (`nn.Module`):
|
||||
Decoder module.
|
||||
latent_channels (`int`, *optional*, defaults to 4):
|
||||
Number of latent channels.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
encoder: nn.Module,
|
||||
decoder: nn.Module,
|
||||
latent_channels: int = 4,
|
||||
dims: int = 2,
|
||||
sample_size=512,
|
||||
use_quant_conv: bool = True,
|
||||
normalize_latent_channels: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
|
||||
self.per_channel_statistics = nn.Module()
|
||||
std_of_means = torch.zeros( (128,), dtype= torch.bfloat16)
|
||||
|
||||
self.per_channel_statistics.register_buffer("std-of-means", std_of_means)
|
||||
self.per_channel_statistics.register_buffer(
|
||||
"mean-of-means",
|
||||
torch.zeros_like(std_of_means)
|
||||
)
|
||||
|
||||
|
||||
|
||||
# pass init params to Encoder
|
||||
self.encoder = encoder
|
||||
self.use_quant_conv = use_quant_conv
|
||||
self.normalize_latent_channels = normalize_latent_channels
|
||||
|
||||
# pass init params to Decoder
|
||||
quant_dims = 2 if dims == 2 else 3
|
||||
self.decoder = decoder
|
||||
if use_quant_conv:
|
||||
self.quant_conv = make_conv_nd(
|
||||
quant_dims, 2 * latent_channels, 2 * latent_channels, 1
|
||||
)
|
||||
self.post_quant_conv = make_conv_nd(
|
||||
quant_dims, latent_channels, latent_channels, 1
|
||||
)
|
||||
else:
|
||||
self.quant_conv = nn.Identity()
|
||||
self.post_quant_conv = nn.Identity()
|
||||
|
||||
if normalize_latent_channels:
|
||||
if dims == 2:
|
||||
self.latent_norm_out = nn.BatchNorm2d(latent_channels, affine=False)
|
||||
else:
|
||||
self.latent_norm_out = nn.BatchNorm3d(latent_channels, affine=False)
|
||||
else:
|
||||
self.latent_norm_out = nn.Identity()
|
||||
self.use_z_tiling = False
|
||||
self.use_hw_tiling = False
|
||||
self.dims = dims
|
||||
self.z_sample_size = 1
|
||||
|
||||
self.decoder_params = inspect.signature(self.decoder.forward).parameters
|
||||
|
||||
# only relevant if vae tiling is enabled
|
||||
self.set_tiling_params(sample_size=sample_size, overlap_factor=0.25)
|
||||
|
||||
@staticmethod
|
||||
def get_VAE_tile_size(vae_config, device_mem_capacity, mixed_precision):
|
||||
|
||||
z_tile = 4
|
||||
# VAE Tiling
|
||||
if vae_config == 0:
|
||||
if mixed_precision:
|
||||
device_mem_capacity = device_mem_capacity / 1.5
|
||||
if device_mem_capacity >= 24000:
|
||||
use_vae_config = 1
|
||||
elif device_mem_capacity >= 8000:
|
||||
use_vae_config = 2
|
||||
else:
|
||||
use_vae_config = 3
|
||||
else:
|
||||
use_vae_config = vae_config
|
||||
|
||||
if use_vae_config == 1:
|
||||
hw_tile = 0
|
||||
elif use_vae_config == 2:
|
||||
hw_tile = 512
|
||||
else:
|
||||
hw_tile = 256
|
||||
|
||||
return (z_tile, hw_tile)
|
||||
|
||||
def set_tiling_params(self, sample_size: int = 512, overlap_factor: float = 0.25):
|
||||
self.tile_sample_min_size = sample_size
|
||||
num_blocks = len(self.encoder.down_blocks)
|
||||
# self.tile_latent_min_size = int(sample_size / (2 ** (num_blocks - 1)))
|
||||
self.tile_latent_min_size = int(sample_size / 32)
|
||||
self.tile_overlap_factor = overlap_factor
|
||||
|
||||
def enable_z_tiling(self, z_sample_size: int = 4):
|
||||
r"""
|
||||
Enable tiling during VAE decoding.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor in tiles to compute decoding in several
|
||||
steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.use_z_tiling = z_sample_size > 1
|
||||
self.z_sample_size = z_sample_size
|
||||
assert (
|
||||
z_sample_size % 4 == 0 or z_sample_size == 1
|
||||
), f"z_sample_size must be a multiple of 4 or 1. Got {z_sample_size}."
|
||||
|
||||
def disable_z_tiling(self):
|
||||
r"""
|
||||
Disable tiling during VAE decoding. If `use_tiling` was previously invoked, this method will go back to computing
|
||||
decoding in one step.
|
||||
"""
|
||||
self.use_z_tiling = False
|
||||
|
||||
def enable_hw_tiling(self):
|
||||
r"""
|
||||
Enable tiling during VAE decoding along the height and width dimension.
|
||||
"""
|
||||
self.use_hw_tiling = True
|
||||
|
||||
def disable_hw_tiling(self):
|
||||
r"""
|
||||
Disable tiling during VAE decoding along the height and width dimension.
|
||||
"""
|
||||
self.use_hw_tiling = False
|
||||
|
||||
def _hw_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True):
|
||||
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_latent_min_size - blend_extent
|
||||
|
||||
# Split the image into 512x512 tiles and encode them separately.
|
||||
rows = []
|
||||
for i in range(0, x.shape[3], overlap_size):
|
||||
row = []
|
||||
for j in range(0, x.shape[4], overlap_size):
|
||||
tile = x[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i : i + self.tile_sample_min_size,
|
||||
j : j + self.tile_sample_min_size,
|
||||
]
|
||||
tile = self.encoder(tile)
|
||||
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):
|
||||
# blend the above tile and the left tile
|
||||
# to the current tile and add the current tile to the result row
|
||||
if i > 0:
|
||||
tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
|
||||
if j > 0:
|
||||
tile = self.blend_h(row[j - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :, :row_limit, :row_limit])
|
||||
result_rows.append(torch.cat(result_row, dim=4))
|
||||
|
||||
moments = torch.cat(result_rows, dim=3)
|
||||
return moments
|
||||
|
||||
def blend_z(
|
||||
self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
|
||||
) -> torch.Tensor:
|
||||
blend_extent = min(a.shape[2], b.shape[2], blend_extent)
|
||||
for z in range(blend_extent):
|
||||
b[:, :, z, :, :] = a[:, :, -blend_extent + z, :, :] * (
|
||||
1 - z / blend_extent
|
||||
) + b[:, :, z, :, :] * (z / blend_extent)
|
||||
return b
|
||||
|
||||
def blend_v(
|
||||
self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
|
||||
) -> torch.Tensor:
|
||||
blend_extent = min(a.shape[3], b.shape[3], 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[4], b.shape[4], 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 _hw_tiled_decode(self, z: torch.FloatTensor, target_shape, timestep = None):
|
||||
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_sample_min_size - blend_extent
|
||||
tile_target_shape = (
|
||||
*target_shape[:3],
|
||||
self.tile_sample_min_size,
|
||||
self.tile_sample_min_size,
|
||||
)
|
||||
# Split z into overlapping 64x64 tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
rows = []
|
||||
for i in range(0, z.shape[3], overlap_size):
|
||||
row = []
|
||||
for j in range(0, z.shape[4], overlap_size):
|
||||
tile = z[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i : i + self.tile_latent_min_size,
|
||||
j : j + self.tile_latent_min_size,
|
||||
]
|
||||
tile = self.post_quant_conv(tile)
|
||||
decoded = self.decoder(tile, target_shape=tile_target_shape, timestep = timestep)
|
||||
row.append(decoded)
|
||||
rows.append(row)
|
||||
result_rows = []
|
||||
for i, row in enumerate(rows):
|
||||
result_row = []
|
||||
for j, tile in enumerate(row):
|
||||
# blend the above tile and the left tile
|
||||
# to the current tile and add the current tile to the result row
|
||||
if i > 0:
|
||||
tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
|
||||
if j > 0:
|
||||
tile = self.blend_h(row[j - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :, :row_limit, :row_limit])
|
||||
result_rows.append(torch.cat(result_row, dim=4))
|
||||
|
||||
dec = torch.cat(result_rows, dim=3)
|
||||
return dec
|
||||
|
||||
def encode(
|
||||
self, z: torch.FloatTensor, return_dict: bool = True
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
|
||||
tile_latent_min_tsize = self.z_sample_size
|
||||
tile_sample_min_tsize = tile_latent_min_tsize * 8
|
||||
tile_overlap_factor = 0.25
|
||||
|
||||
B, C, T, H, W = z.shape
|
||||
overlap_size = int(tile_sample_min_tsize * (1 - tile_overlap_factor))
|
||||
blend_extent = int(tile_latent_min_tsize * tile_overlap_factor)
|
||||
t_limit = tile_latent_min_tsize - blend_extent
|
||||
|
||||
row = []
|
||||
for i in range(0, T, overlap_size):
|
||||
tile = z[:, :, i: i + tile_sample_min_tsize + 1, :, :]
|
||||
if self.use_hw_tiling:
|
||||
tile = self._hw_tiled_encode(tile, return_dict)
|
||||
else:
|
||||
tile = self._encode(tile)
|
||||
if i > 0:
|
||||
tile = tile[:, :, 1:, :, :]
|
||||
row.append(tile)
|
||||
result_row = []
|
||||
for i, tile in enumerate(row):
|
||||
if i > 0:
|
||||
tile = self.blend_z(row[i - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :t_limit, :, :])
|
||||
else:
|
||||
result_row.append(tile[:, :, :t_limit + 1, :, :])
|
||||
|
||||
moments = torch.cat(result_row, dim=2)
|
||||
|
||||
|
||||
else:
|
||||
moments = (
|
||||
self._hw_tiled_encode(z, return_dict)
|
||||
if self.use_hw_tiling and z.shape[2] > 1
|
||||
else self._encode(z)
|
||||
)
|
||||
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
if not return_dict:
|
||||
return (posterior,)
|
||||
|
||||
return AutoencoderKLOutput(latent_dist=posterior)
|
||||
|
||||
def _normalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
_, c, _, _, _ = z.shape
|
||||
z = torch.cat(
|
||||
[
|
||||
self.latent_norm_out(z[:, : c // 2, :, :, :]),
|
||||
z[:, c // 2 :, :, :, :],
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
elif isinstance(self.latent_norm_out, nn.BatchNorm2d):
|
||||
raise NotImplementedError("BatchNorm2d not supported")
|
||||
return z
|
||||
|
||||
def _unnormalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
running_mean = self.latent_norm_out.running_mean.view(1, -1, 1, 1, 1)
|
||||
running_var = self.latent_norm_out.running_var.view(1, -1, 1, 1, 1)
|
||||
eps = self.latent_norm_out.eps
|
||||
|
||||
z = z * torch.sqrt(running_var + eps) + running_mean
|
||||
elif isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
raise NotImplementedError("BatchNorm2d not supported")
|
||||
return z
|
||||
|
||||
def _encode(self, x: torch.FloatTensor) -> AutoencoderKLOutput:
|
||||
h = self.encoder(x)
|
||||
moments = self.quant_conv(h)
|
||||
moments = self._normalize_latent_channels(moments)
|
||||
return moments
|
||||
|
||||
def _decode(
|
||||
self,
|
||||
z: torch.FloatTensor,
|
||||
target_shape=None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
z = self._unnormalize_latent_channels(z)
|
||||
z = self.post_quant_conv(z)
|
||||
if "timestep" in self.decoder_params:
|
||||
dec = self.decoder(z, target_shape=target_shape, timestep=timestep)
|
||||
else:
|
||||
dec = self.decoder(z, target_shape=target_shape)
|
||||
return dec
|
||||
|
||||
def decode(
|
||||
self,
|
||||
z: torch.FloatTensor,
|
||||
return_dict: bool = True,
|
||||
target_shape=None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
assert target_shape is not None, "target_shape must be provided for decoding"
|
||||
if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
tile_latent_min_tsize = self.z_sample_size
|
||||
tile_sample_min_tsize = tile_latent_min_tsize * 8
|
||||
tile_overlap_factor = 0.25
|
||||
|
||||
B, C, T, H, W = z.shape
|
||||
overlap_size = int(tile_latent_min_tsize * (1 - tile_overlap_factor))
|
||||
blend_extent = int(tile_sample_min_tsize * tile_overlap_factor)
|
||||
t_limit = tile_sample_min_tsize - blend_extent
|
||||
|
||||
row = []
|
||||
for i in range(0, T, overlap_size):
|
||||
tile = z[:, :, i: i + tile_latent_min_tsize + 1, :, :]
|
||||
target_shape_split = list(target_shape)
|
||||
target_shape_split[2] = tile.shape[2] * 8
|
||||
if self.use_hw_tiling:
|
||||
decoded = self._hw_tiled_decode(tile, target_shape, timestep)
|
||||
else:
|
||||
decoded = self._decode(tile, target_shape=target_shape, timestep=timestep)
|
||||
|
||||
if i > 0:
|
||||
decoded = decoded[:, :, 1:, :, :]
|
||||
row.append(decoded.to(torch.float16).cpu())
|
||||
decoded = None
|
||||
result_row = []
|
||||
for i, tile in enumerate(row):
|
||||
if i > 0:
|
||||
tile = self.blend_z(row[i - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :t_limit, :, :])
|
||||
else:
|
||||
result_row.append(tile[:, :, :t_limit + 1, :, :])
|
||||
|
||||
dec = torch.cat(result_row, dim=2)
|
||||
if not return_dict:
|
||||
return (dec,)
|
||||
|
||||
return DecoderOutput(sample=dec)
|
||||
else:
|
||||
decoded = (
|
||||
self._hw_tiled_decode(z, target_shape, timestep)
|
||||
if self.use_hw_tiling
|
||||
else self._decode(z, target_shape=target_shape, timestep=timestep)
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return (decoded,)
|
||||
|
||||
return DecoderOutput(sample=decoded)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
sample_posterior: bool = False,
|
||||
return_dict: bool = True,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): Input sample.
|
||||
sample_posterior (`bool`, *optional*, defaults to `False`):
|
||||
Whether to sample from the posterior.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether to return a [`DecoderOutput`] instead of a plain tuple.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
Generator used to sample from the posterior.
|
||||
"""
|
||||
x = sample
|
||||
posterior = self.encode(x).latent_dist
|
||||
if sample_posterior:
|
||||
z = posterior.sample(generator=generator)
|
||||
else:
|
||||
z = posterior.mode()
|
||||
dec = self.decode(z, target_shape=sample.shape).sample
|
||||
|
||||
if not return_dict:
|
||||
return (dec,)
|
||||
|
||||
return DecoderOutput(sample=dec)
|
||||
247
ltx_video/models/autoencoders/vae_encode.py
Normal file
247
ltx_video/models/autoencoders/vae_encode.py
Normal file
@@ -0,0 +1,247 @@
|
||||
from typing import Tuple
|
||||
import torch
|
||||
from diffusers import AutoencoderKL
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
from ltx_video.models.autoencoders.causal_video_autoencoder import (
|
||||
CausalVideoAutoencoder,
|
||||
)
|
||||
from ltx_video.models.autoencoders.video_autoencoder import (
|
||||
Downsample3D,
|
||||
VideoAutoencoder,
|
||||
)
|
||||
|
||||
try:
|
||||
import torch_xla.core.xla_model as xm
|
||||
except ImportError:
|
||||
xm = None
|
||||
|
||||
|
||||
def vae_encode(
|
||||
media_items: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
split_size: int = 1,
|
||||
vae_per_channel_normalize=False,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Encodes media items (images or videos) into latent representations using a specified VAE model.
|
||||
The function supports processing batches of images or video frames and can handle the processing
|
||||
in smaller sub-batches if needed.
|
||||
|
||||
Args:
|
||||
media_items (Tensor): A torch Tensor containing the media items to encode. The expected
|
||||
shape is (batch_size, channels, height, width) for images or (batch_size, channels,
|
||||
frames, height, width) for videos.
|
||||
vae (AutoencoderKL): An instance of the `AutoencoderKL` class from the `diffusers` library,
|
||||
pre-configured and loaded with the appropriate model weights.
|
||||
split_size (int, optional): The number of sub-batches to split the input batch into for encoding.
|
||||
If set to more than 1, the input media items are processed in smaller batches according to
|
||||
this value. Defaults to 1, which processes all items in a single batch.
|
||||
|
||||
Returns:
|
||||
Tensor: A torch Tensor of the encoded latent representations. The shape of the tensor is adjusted
|
||||
to match the input shape, scaled by the model's configuration.
|
||||
|
||||
Examples:
|
||||
>>> import torch
|
||||
>>> from diffusers import AutoencoderKL
|
||||
>>> vae = AutoencoderKL.from_pretrained('your-model-name')
|
||||
>>> images = torch.rand(10, 3, 8 256, 256) # Example tensor with 10 videos of 8 frames.
|
||||
>>> latents = vae_encode(images, vae)
|
||||
>>> print(latents.shape) # Output shape will depend on the model's latent configuration.
|
||||
|
||||
Note:
|
||||
In case of a video, the function encodes the media item frame-by frame.
|
||||
"""
|
||||
is_video_shaped = media_items.dim() == 5
|
||||
batch_size, channels = media_items.shape[0:2]
|
||||
|
||||
if channels != 3:
|
||||
raise ValueError(f"Expects tensors with 3 channels, got {channels}.")
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
|
||||
if split_size > 1:
|
||||
if len(media_items) % split_size != 0:
|
||||
raise ValueError(
|
||||
"Error: The batch size must be divisible by 'train.vae_bs_split"
|
||||
)
|
||||
encode_bs = len(media_items) // split_size
|
||||
# latents = [vae.encode(image_batch).latent_dist.sample() for image_batch in media_items.split(encode_bs)]
|
||||
latents = []
|
||||
if media_items.device.type == "xla":
|
||||
xm.mark_step()
|
||||
for image_batch in media_items.split(encode_bs):
|
||||
latents.append(vae.encode(image_batch).latent_dist.sample())
|
||||
if media_items.device.type == "xla":
|
||||
xm.mark_step()
|
||||
latents = torch.cat(latents, dim=0)
|
||||
else:
|
||||
latents = vae.encode(media_items).latent_dist.sample()
|
||||
|
||||
latents = normalize_latents(latents, vae, vae_per_channel_normalize)
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
latents = rearrange(latents, "(b n) c h w -> b c n h w", b=batch_size)
|
||||
return latents
|
||||
|
||||
|
||||
def vae_decode(
|
||||
latents: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
is_video: bool = True,
|
||||
split_size: int = 1,
|
||||
vae_per_channel_normalize=False,
|
||||
timestep=None,
|
||||
) -> Tensor:
|
||||
is_video_shaped = latents.dim() == 5
|
||||
batch_size = latents.shape[0]
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
latents = rearrange(latents, "b c n h w -> (b n) c h w")
|
||||
if split_size > 1:
|
||||
if len(latents) % split_size != 0:
|
||||
raise ValueError(
|
||||
"Error: The batch size must be divisible by 'train.vae_bs_split"
|
||||
)
|
||||
encode_bs = len(latents) // split_size
|
||||
image_batch = [
|
||||
_run_decoder(
|
||||
latent_batch, vae, is_video, vae_per_channel_normalize, timestep
|
||||
)
|
||||
for latent_batch in latents.split(encode_bs)
|
||||
]
|
||||
images = torch.cat(image_batch, dim=0)
|
||||
else:
|
||||
images = _run_decoder(
|
||||
latents, vae, is_video, vae_per_channel_normalize, timestep
|
||||
)
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
images = rearrange(images, "(b n) c h w -> b c n h w", b=batch_size)
|
||||
return images
|
||||
|
||||
|
||||
def _run_decoder(
|
||||
latents: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
is_video: bool,
|
||||
vae_per_channel_normalize=False,
|
||||
timestep=None,
|
||||
) -> Tensor:
|
||||
if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)):
|
||||
*_, fl, hl, wl = latents.shape
|
||||
temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae)
|
||||
latents = latents.to(vae.dtype)
|
||||
vae_decode_kwargs = {}
|
||||
if timestep is not None:
|
||||
vae_decode_kwargs["timestep"] = timestep
|
||||
image = vae.decode(
|
||||
un_normalize_latents(latents, vae, vae_per_channel_normalize),
|
||||
return_dict=False,
|
||||
target_shape=(
|
||||
1,
|
||||
3,
|
||||
fl * temporal_scale if is_video else 1,
|
||||
hl * spatial_scale,
|
||||
wl * spatial_scale,
|
||||
),
|
||||
**vae_decode_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
image = vae.decode(
|
||||
un_normalize_latents(latents, vae, vae_per_channel_normalize),
|
||||
return_dict=False,
|
||||
)[0]
|
||||
return image
|
||||
|
||||
|
||||
def get_vae_size_scale_factor(vae: AutoencoderKL) -> float:
|
||||
if isinstance(vae, CausalVideoAutoencoder):
|
||||
spatial = vae.spatial_downscale_factor
|
||||
temporal = vae.temporal_downscale_factor
|
||||
else:
|
||||
down_blocks = len(
|
||||
[
|
||||
block
|
||||
for block in vae.encoder.down_blocks
|
||||
if isinstance(block.downsample, Downsample3D)
|
||||
]
|
||||
)
|
||||
spatial = vae.config.patch_size * 2**down_blocks
|
||||
temporal = (
|
||||
vae.config.patch_size_t * 2**down_blocks
|
||||
if isinstance(vae, VideoAutoencoder)
|
||||
else 1
|
||||
)
|
||||
|
||||
return (temporal, spatial, spatial)
|
||||
|
||||
|
||||
def latent_to_pixel_coords(
|
||||
latent_coords: Tensor, vae: AutoencoderKL, causal_fix: bool = False
|
||||
) -> Tensor:
|
||||
"""
|
||||
Converts latent coordinates to pixel coordinates by scaling them according to the VAE's
|
||||
configuration.
|
||||
|
||||
Args:
|
||||
latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents]
|
||||
containing the latent corner coordinates of each token.
|
||||
vae (AutoencoderKL): The VAE model
|
||||
causal_fix (bool): Whether to take into account the different temporal scale
|
||||
of the first frame. Default = False for backwards compatibility.
|
||||
Returns:
|
||||
Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates.
|
||||
"""
|
||||
|
||||
scale_factors = get_vae_size_scale_factor(vae)
|
||||
causal_fix = isinstance(vae, CausalVideoAutoencoder) and causal_fix
|
||||
pixel_coords = latent_to_pixel_coords_from_factors(
|
||||
latent_coords, scale_factors, causal_fix
|
||||
)
|
||||
return pixel_coords
|
||||
|
||||
|
||||
def latent_to_pixel_coords_from_factors(
|
||||
latent_coords: Tensor, scale_factors: Tuple, causal_fix: bool = False
|
||||
) -> Tensor:
|
||||
pixel_coords = (
|
||||
latent_coords
|
||||
* torch.tensor(scale_factors, device=latent_coords.device)[None, :, None]
|
||||
)
|
||||
if causal_fix:
|
||||
# Fix temporal scale for first frame to 1 due to causality
|
||||
pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0)
|
||||
return pixel_coords
|
||||
|
||||
|
||||
def normalize_latents(
|
||||
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
|
||||
) -> Tensor:
|
||||
return (
|
||||
(latents - vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1))
|
||||
/ vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
if vae_per_channel_normalize
|
||||
else latents * vae.config.scaling_factor
|
||||
)
|
||||
|
||||
|
||||
def un_normalize_latents(
|
||||
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
|
||||
) -> Tensor:
|
||||
return (
|
||||
latents * vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
+ vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
if vae_per_channel_normalize
|
||||
else latents / vae.config.scaling_factor
|
||||
)
|
||||
1045
ltx_video/models/autoencoders/video_autoencoder.py
Normal file
1045
ltx_video/models/autoencoders/video_autoencoder.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ltx_video/models/transformers/__init__.py
Normal file
0
ltx_video/models/transformers/__init__.py
Normal file
1323
ltx_video/models/transformers/attention.py
Normal file
1323
ltx_video/models/transformers/attention.py
Normal file
File diff suppressed because it is too large
Load Diff
129
ltx_video/models/transformers/embeddings.py
Normal file
129
ltx_video/models/transformers/embeddings.py
Normal file
@@ -0,0 +1,129 @@
|
||||
# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/embeddings.py
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
|
||||
def get_timestep_embedding(
|
||||
timesteps: torch.Tensor,
|
||||
embedding_dim: int,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
):
|
||||
"""
|
||||
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
|
||||
|
||||
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
|
||||
embeddings. :return: an [N x dim] Tensor of positional embeddings.
|
||||
"""
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
exponent = -math.log(max_period) * torch.arange(
|
||||
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
|
||||
)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
# scale embeddings
|
||||
emb = scale * emb
|
||||
|
||||
# concat sine and cosine embeddings
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
|
||||
# flip sine and cosine embeddings
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
|
||||
# zero pad
|
||||
if embedding_dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
def get_3d_sincos_pos_embed(embed_dim, grid, w, h, f):
|
||||
"""
|
||||
grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
|
||||
[1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
|
||||
"""
|
||||
grid = rearrange(grid, "c (f h w) -> c f h w", h=h, w=w)
|
||||
grid = rearrange(grid, "c f h w -> c h w f", h=h, w=w)
|
||||
grid = grid.reshape([3, 1, w, h, f])
|
||||
pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
pos_embed = pos_embed.transpose(1, 0, 2, 3)
|
||||
return rearrange(pos_embed, "h w f c -> (f h w) c")
|
||||
|
||||
|
||||
def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
|
||||
if embed_dim % 3 != 0:
|
||||
raise ValueError("embed_dim must be divisible by 3")
|
||||
|
||||
# use half of dimensions to encode grid_h
|
||||
emb_f = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (H*W*T, D/3)
|
||||
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (H*W*T, D/3)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (H*W*T, D/3)
|
||||
|
||||
emb = np.concatenate([emb_h, emb_w, emb_f], axis=-1) # (H*W*T, D)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
|
||||
"""
|
||||
if embed_dim % 2 != 0:
|
||||
raise ValueError("embed_dim must be divisible by 2")
|
||||
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.0
|
||||
omega = 1.0 / 10000**omega # (D/2,)
|
||||
|
||||
pos_shape = pos.shape
|
||||
|
||||
pos = pos.reshape(-1)
|
||||
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
|
||||
out = out.reshape([*pos_shape, -1])[0]
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (M, D)
|
||||
return emb
|
||||
|
||||
|
||||
class SinusoidalPositionalEmbedding(nn.Module):
|
||||
"""Apply positional information to a sequence of embeddings.
|
||||
|
||||
Takes in a sequence of embeddings with shape (batch_size, seq_length, embed_dim) and adds positional embeddings to
|
||||
them
|
||||
|
||||
Args:
|
||||
embed_dim: (int): Dimension of the positional embedding.
|
||||
max_seq_length: Maximum sequence length to apply positional embeddings
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, embed_dim: int, max_seq_length: int = 32):
|
||||
super().__init__()
|
||||
position = torch.arange(max_seq_length).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)
|
||||
)
|
||||
pe = torch.zeros(1, max_seq_length, embed_dim)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer("pe", pe)
|
||||
|
||||
def forward(self, x):
|
||||
_, seq_length, _ = x.shape
|
||||
x = x + self.pe[:, :seq_length]
|
||||
return x
|
||||
84
ltx_video/models/transformers/symmetric_patchifier.py
Normal file
84
ltx_video/models/transformers/symmetric_patchifier.py
Normal file
@@ -0,0 +1,84 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class Patchifier(ConfigMixin, ABC):
|
||||
def __init__(self, patch_size: int):
|
||||
super().__init__()
|
||||
self._patch_size = (1, patch_size, patch_size)
|
||||
|
||||
@abstractmethod
|
||||
def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:
|
||||
raise NotImplementedError("Patchify method not implemented")
|
||||
|
||||
@abstractmethod
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: Tensor,
|
||||
output_height: int,
|
||||
output_width: int,
|
||||
out_channels: int,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
def patch_size(self):
|
||||
return self._patch_size
|
||||
|
||||
def get_latent_coords(
|
||||
self, latent_num_frames, latent_height, latent_width, batch_size, device
|
||||
):
|
||||
"""
|
||||
Return a tensor of shape [batch_size, 3, num_patches] containing the
|
||||
top-left corner latent coordinates of each latent patch.
|
||||
The tensor is repeated for each batch element.
|
||||
"""
|
||||
latent_sample_coords = torch.meshgrid(
|
||||
torch.arange(0, latent_num_frames, self._patch_size[0], device=device),
|
||||
torch.arange(0, latent_height, self._patch_size[1], device=device),
|
||||
torch.arange(0, latent_width, self._patch_size[2], device=device),
|
||||
)
|
||||
latent_sample_coords = torch.stack(latent_sample_coords, dim=0)
|
||||
latent_coords = latent_sample_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1)
|
||||
latent_coords = rearrange(
|
||||
latent_coords, "b c f h w -> b c (f h w)", b=batch_size
|
||||
)
|
||||
return latent_coords
|
||||
|
||||
|
||||
class SymmetricPatchifier(Patchifier):
|
||||
def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:
|
||||
b, _, f, h, w = latents.shape
|
||||
latent_coords = self.get_latent_coords(f, h, w, b, latents.device)
|
||||
latents = rearrange(
|
||||
latents,
|
||||
"b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
|
||||
p1=self._patch_size[0],
|
||||
p2=self._patch_size[1],
|
||||
p3=self._patch_size[2],
|
||||
)
|
||||
return latents, latent_coords
|
||||
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: Tensor,
|
||||
output_height: int,
|
||||
output_width: int,
|
||||
out_channels: int,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
output_height = output_height // self._patch_size[1]
|
||||
output_width = output_width // self._patch_size[2]
|
||||
latents = rearrange(
|
||||
latents,
|
||||
"b (f h w) (c p q) -> b c f (h p) (w q)",
|
||||
h=output_height,
|
||||
w=output_width,
|
||||
p=self._patch_size[1],
|
||||
q=self._patch_size[2],
|
||||
)
|
||||
return latents
|
||||
507
ltx_video/models/transformers/transformer3d.py
Normal file
507
ltx_video/models/transformers/transformer3d.py
Normal file
@@ -0,0 +1,507 @@
|
||||
# Adapted from: https://github.com/huggingface/diffusers/blob/v0.26.3/src/diffusers/models/transformers/transformer_2d.py
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import os
|
||||
import json
|
||||
import glob
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import AdaLayerNormSingle
|
||||
from diffusers.utils import BaseOutput, is_torch_version
|
||||
from diffusers.utils import logging
|
||||
from torch import nn
|
||||
from safetensors import safe_open
|
||||
from ltx_video.models.transformers.attention import BasicTransformerBlock, reshape_hidden_states, restore_hidden_states_shape
|
||||
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
|
||||
|
||||
from ltx_video.utils.diffusers_config_mapping import (
|
||||
diffusers_and_ours_config_mapping,
|
||||
make_hashable_key,
|
||||
TRANSFORMER_KEYS_RENAME_DICT,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transformer3DModelOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`Transformer2DModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
|
||||
The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
|
||||
distributions for the unnoised latent pixels.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class Transformer3DModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
out_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
num_vector_embeds: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
adaptive_norm: str = "single_scale_shift", # 'single_scale_shift' or 'single_scale'
|
||||
standardization_norm: str = "layer_norm", # 'layer_norm' or 'rms_norm'
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_eps: float = 1e-5,
|
||||
attention_type: str = "default",
|
||||
caption_channels: int = None,
|
||||
use_tpu_flash_attention: bool = False, # if True uses the TPU attention offload ('flash attention')
|
||||
qk_norm: Optional[str] = None,
|
||||
positional_embedding_type: str = "rope",
|
||||
positional_embedding_theta: Optional[float] = None,
|
||||
positional_embedding_max_pos: Optional[List[int]] = None,
|
||||
timestep_scale_multiplier: Optional[float] = None,
|
||||
causal_temporal_positioning: bool = False, # For backward compatibility, will be deprecated
|
||||
):
|
||||
super().__init__()
|
||||
self.use_tpu_flash_attention = (
|
||||
use_tpu_flash_attention # FIXME: push config down to the attention modules
|
||||
)
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
self.inner_dim = inner_dim
|
||||
self.patchify_proj = nn.Linear(in_channels, inner_dim, bias=True)
|
||||
self.positional_embedding_type = positional_embedding_type
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.positional_embedding_max_pos = positional_embedding_max_pos
|
||||
self.use_rope = self.positional_embedding_type == "rope"
|
||||
self.timestep_scale_multiplier = timestep_scale_multiplier
|
||||
|
||||
if self.positional_embedding_type == "absolute":
|
||||
raise ValueError("Absolute positional embedding is no longer supported")
|
||||
elif self.positional_embedding_type == "rope":
|
||||
if positional_embedding_theta is None:
|
||||
raise ValueError(
|
||||
"If `positional_embedding_type` type is rope, `positional_embedding_theta` must also be defined"
|
||||
)
|
||||
if positional_embedding_max_pos is None:
|
||||
raise ValueError(
|
||||
"If `positional_embedding_type` type is rope, `positional_embedding_max_pos` must also be defined"
|
||||
)
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
double_self_attention=double_self_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
adaptive_norm=adaptive_norm,
|
||||
standardization_norm=standardization_norm,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
norm_eps=norm_eps,
|
||||
attention_type=attention_type,
|
||||
use_tpu_flash_attention=use_tpu_flash_attention,
|
||||
qk_norm=qk_norm,
|
||||
use_rope=self.use_rope,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
|
||||
self.scale_shift_table = nn.Parameter(
|
||||
torch.randn(2, inner_dim) / inner_dim**0.5
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, self.out_channels)
|
||||
|
||||
self.adaln_single = AdaLayerNormSingle(
|
||||
inner_dim, use_additional_conditions=False
|
||||
)
|
||||
if adaptive_norm == "single_scale":
|
||||
self.adaln_single.linear = nn.Linear(inner_dim, 4 * inner_dim, bias=True)
|
||||
|
||||
self.caption_projection = None
|
||||
if caption_channels is not None:
|
||||
self.caption_projection = PixArtAlphaTextProjection(
|
||||
in_features=caption_channels, hidden_size=inner_dim
|
||||
)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def set_use_tpu_flash_attention(self):
|
||||
r"""
|
||||
Function sets the flag in this object and propagates down the children. The flag will enforce the usage of TPU
|
||||
attention kernel.
|
||||
"""
|
||||
logger.info("ENABLE TPU FLASH ATTENTION -> TRUE")
|
||||
self.use_tpu_flash_attention = True
|
||||
# push config down to the attention modules
|
||||
for block in self.transformer_blocks:
|
||||
block.set_use_tpu_flash_attention()
|
||||
|
||||
def create_skip_layer_mask(
|
||||
self,
|
||||
batch_size: int,
|
||||
num_conds: int,
|
||||
ptb_index: int,
|
||||
skip_block_list: Optional[List[int]] = None,
|
||||
):
|
||||
if skip_block_list is None or len(skip_block_list) == 0:
|
||||
return None
|
||||
num_layers = len(self.transformer_blocks)
|
||||
mask = torch.ones(
|
||||
(num_layers, batch_size * num_conds), device=self.device, dtype=self.dtype
|
||||
)
|
||||
for block_idx in skip_block_list:
|
||||
mask[block_idx, ptb_index::num_conds] = 0
|
||||
return mask
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def get_fractional_positions(self, indices_grid):
|
||||
fractional_positions = torch.stack(
|
||||
[
|
||||
indices_grid[:, i] / self.positional_embedding_max_pos[i]
|
||||
for i in range(3)
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
return fractional_positions
|
||||
|
||||
def precompute_freqs_cis(self, indices_grid, spacing="exp"):
|
||||
dtype = torch.float32 # We need full precision in the freqs_cis computation.
|
||||
dim = self.inner_dim
|
||||
theta = self.positional_embedding_theta
|
||||
|
||||
fractional_positions = self.get_fractional_positions(indices_grid)
|
||||
|
||||
start = 1
|
||||
end = theta
|
||||
device = fractional_positions.device
|
||||
if spacing == "exp":
|
||||
indices = theta ** (
|
||||
torch.linspace(
|
||||
math.log(start, theta),
|
||||
math.log(end, theta),
|
||||
dim // 6,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
)
|
||||
indices = indices.to(dtype=dtype)
|
||||
elif spacing == "exp_2":
|
||||
indices = 1.0 / theta ** (torch.arange(0, dim, 6, device=device) / dim)
|
||||
indices = indices.to(dtype=dtype)
|
||||
elif spacing == "linear":
|
||||
indices = torch.linspace(start, end, dim // 6, device=device, dtype=dtype)
|
||||
elif spacing == "sqrt":
|
||||
indices = torch.linspace(
|
||||
start**2, end**2, dim // 6, device=device, dtype=dtype
|
||||
).sqrt()
|
||||
|
||||
indices = indices * math.pi / 2
|
||||
|
||||
if spacing == "exp_2":
|
||||
freqs = (
|
||||
(indices * fractional_positions.unsqueeze(-1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
else:
|
||||
freqs = (
|
||||
(indices * (fractional_positions.unsqueeze(-1) * 2 - 1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
|
||||
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
if dim % 6 != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, : dim % 6])
|
||||
sin_padding = torch.zeros_like(cos_freq[:, :, : dim % 6])
|
||||
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
|
||||
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
|
||||
return cos_freq.to(self.dtype), sin_freq.to(self.dtype)
|
||||
|
||||
def load_state_dict(
|
||||
self,
|
||||
state_dict: Dict,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if any([key.startswith("model.diffusion_model.") for key in state_dict.keys()]):
|
||||
state_dict = {
|
||||
key.replace("model.diffusion_model.", ""): value
|
||||
for key, value in state_dict.items()
|
||||
if key.startswith("model.diffusion_model.")
|
||||
}
|
||||
return super().load_state_dict(state_dict, **kwargs)
|
||||
|
||||
@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_dir():
|
||||
config_path = pretrained_model_path / "transformer" / "config.json"
|
||||
with open(config_path, "r") as f:
|
||||
config = make_hashable_key(json.load(f))
|
||||
|
||||
assert config in diffusers_and_ours_config_mapping, (
|
||||
"Provided diffusers checkpoint config for transformer is not suppported. "
|
||||
"We only support diffusers configs found in Lightricks/LTX-Video."
|
||||
)
|
||||
|
||||
config = diffusers_and_ours_config_mapping[config]
|
||||
state_dict = {}
|
||||
ckpt_paths = (
|
||||
pretrained_model_path
|
||||
/ "transformer"
|
||||
/ "diffusion_pytorch_model*.safetensors"
|
||||
)
|
||||
dict_list = glob.glob(str(ckpt_paths))
|
||||
for dict_path in dict_list:
|
||||
part_dict = {}
|
||||
with safe_open(dict_path, framework="pt", device="cpu") as f:
|
||||
for k in f.keys():
|
||||
part_dict[k] = f.get_tensor(k)
|
||||
state_dict.update(part_dict)
|
||||
|
||||
for key in list(state_dict.keys()):
|
||||
new_key = key
|
||||
for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
|
||||
new_key = new_key.replace(replace_key, rename_key)
|
||||
state_dict[new_key] = state_dict.pop(key)
|
||||
|
||||
with torch.device("meta"):
|
||||
transformer = cls.from_config(config)
|
||||
transformer.load_state_dict(state_dict, assign=True, strict=True)
|
||||
elif pretrained_model_path.is_file() and str(pretrained_model_path).endswith(
|
||||
".safetensors"
|
||||
):
|
||||
comfy_single_file_state_dict = {}
|
||||
with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
|
||||
metadata = f.metadata()
|
||||
for k in f.keys():
|
||||
comfy_single_file_state_dict[k] = f.get_tensor(k)
|
||||
configs = json.loads(metadata["config"])
|
||||
transformer_config = configs["transformer"]
|
||||
with torch.device("meta"):
|
||||
transformer = Transformer3DModel.from_config(transformer_config)
|
||||
transformer.load_state_dict(comfy_single_file_state_dict, assign=True)
|
||||
return transformer
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
freqs_cis: list,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
skip_layer_mask: Optional[torch.Tensor] = None,
|
||||
skip_layer_strategy: Optional[SkipLayerStrategy] = None,
|
||||
latent_shape = None,
|
||||
joint_pass = True,
|
||||
ltxv_model = None,
|
||||
mixed = False,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
indices_grid (`torch.LongTensor` of shape `(batch size, 3, num latent pixels)`):
|
||||
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
|
||||
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
|
||||
self-attention.
|
||||
timestep ( `torch.LongTensor`, *optional*):
|
||||
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
|
||||
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
|
||||
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
|
||||
`AdaLayerZeroNorm`.
|
||||
cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
attention_mask ( `torch.Tensor`, *optional*):
|
||||
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
|
||||
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
|
||||
negative values to the attention scores corresponding to "discard" tokens.
|
||||
encoder_attention_mask ( `torch.Tensor`, *optional*):
|
||||
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
|
||||
|
||||
* Mask `(batch, sequence_length)` True = keep, False = discard.
|
||||
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
|
||||
|
||||
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
|
||||
above. This bias will be added to the cross-attention scores.
|
||||
skip_layer_mask ( `torch.Tensor`, *optional*):
|
||||
A mask of shape `(num_layers, batch)` that indicates which layers to skip. `0` at position
|
||||
`layer, batch_idx` indicates that the layer should be skipped for the corresponding batch index.
|
||||
skip_layer_strategy ( `SkipLayerStrategy`, *optional*, defaults to `None`):
|
||||
Controls which layers are skipped when calculating a perturbed latent for spatiotemporal guidance.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# for tpu attention offload 2d token masks are used. No need to transform.
|
||||
if not self.use_tpu_flash_attention:
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
|
||||
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
|
||||
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
|
||||
# expects mask of shape:
|
||||
# [batch, key_tokens]
|
||||
# adds singleton query_tokens dimension:
|
||||
# [batch, 1, key_tokens]
|
||||
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
|
||||
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
|
||||
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
|
||||
if attention_mask is not None and attention_mask.ndim == 2:
|
||||
# assume that mask is expressed as:
|
||||
# (1 = keep, 0 = discard)
|
||||
# convert mask into a bias that can be added to attention scores:
|
||||
# (keep = +0, discard = -10000.0)
|
||||
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
|
||||
encoder_attention_mask = (
|
||||
1 - encoder_attention_mask.to(hidden_states.dtype)
|
||||
) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# 1. Input
|
||||
hidden_states = self.patchify_proj(hidden_states)
|
||||
|
||||
if self.timestep_scale_multiplier:
|
||||
timestep = self.timestep_scale_multiplier * timestep
|
||||
|
||||
if timestep.shape[-1] > 1:
|
||||
timestep = timestep.reshape(timestep.shape[0], -1, latent_shape[-2] * latent_shape[-1] )
|
||||
timestep = timestep[:, :, 0]
|
||||
|
||||
batch_size = hidden_states.shape[0]
|
||||
timestep, embedded_timestep = self.adaln_single(
|
||||
timestep.flatten(),
|
||||
{"resolution": None, "aspect_ratio": None},
|
||||
batch_size=batch_size,
|
||||
hidden_dtype=hidden_states.dtype,
|
||||
)
|
||||
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
||||
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
||||
embedded_timestep = embedded_timestep.view(
|
||||
batch_size, -1, embedded_timestep.shape[-1]
|
||||
)
|
||||
if mixed:
|
||||
timestep = timestep.float()
|
||||
embedded_timestep = embedded_timestep.float()
|
||||
hidden_states = hidden_states.float()
|
||||
|
||||
|
||||
# 2. Blocks
|
||||
if self.caption_projection is not None:
|
||||
batch_size = hidden_states.shape[0]
|
||||
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states.view(
|
||||
batch_size, -1, hidden_states.shape[-1]
|
||||
)
|
||||
|
||||
|
||||
if joint_pass:
|
||||
for block_idx, block in enumerate(self.transformer_blocks):
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
freqs_cis=freqs_cis,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
timestep=timestep,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
skip_layer_mask= None if skip_layer_mask is None else skip_layer_mask[block_idx],
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
)
|
||||
if ltxv_model._interrupt:
|
||||
return [None]
|
||||
|
||||
else:
|
||||
for block_idx, block in enumerate(self.transformer_blocks):
|
||||
for i, (one_hidden_states, one_encoder_hidden_states, one_encoder_attention_mask,one_timestep) in enumerate(zip(hidden_states, encoder_hidden_states,encoder_attention_mask,timestep)):
|
||||
hidden_states[i][...] = block(
|
||||
one_hidden_states.unsqueeze(0),
|
||||
freqs_cis=freqs_cis,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=one_encoder_hidden_states.unsqueeze(0),
|
||||
encoder_attention_mask=one_encoder_attention_mask.unsqueeze(0),
|
||||
timestep=one_timestep.unsqueeze(0),
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
skip_layer_mask= None if skip_layer_mask is None else skip_layer_mask[block_idx, i],
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
)
|
||||
if ltxv_model._interrupt:
|
||||
return [None]
|
||||
|
||||
# 3. Output
|
||||
scale_shift_values = (
|
||||
self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
|
||||
)
|
||||
shift, scale = scale_shift_values[:, :, 0].unsqueeze(-2), scale_shift_values[:, :, 1].unsqueeze(-2)
|
||||
hidden_states = self.norm_out(hidden_states)
|
||||
# Modulation
|
||||
|
||||
|
||||
hidden_states = reshape_hidden_states(hidden_states, scale.shape[1])
|
||||
# hidden_states = hidden_states * (1 + scale)
|
||||
hidden_states *= 1 + scale
|
||||
hidden_states += shift
|
||||
hidden_states = restore_hidden_states_shape(hidden_states)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
if not return_dict:
|
||||
return (hidden_states,)
|
||||
|
||||
return Transformer3DModelOutput(sample=hidden_states)
|
||||
Reference in New Issue
Block a user