945 lines
28 KiB
Python
945 lines
28 KiB
Python
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
|
import logging
|
|
from typing import Tuple, Union
|
|
|
|
import torch
|
|
import torch.cuda.amp as amp
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
|
from diffusers.loaders.single_file_model import FromOriginalModelMixin
|
|
from diffusers.models.autoencoders.vae import (DecoderOutput,
|
|
DiagonalGaussianDistribution)
|
|
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
|
from diffusers.models.modeling_utils import ModelMixin
|
|
from diffusers.utils.accelerate_utils import apply_forward_hook
|
|
from einops import rearrange
|
|
from torch.utils.checkpoint import checkpoint as torch_checkpoint
|
|
|
|
CACHE_T = 2
|
|
|
|
|
|
class CausalConv3d(nn.Conv3d):
|
|
"""
|
|
Causal 3d convolusion.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._padding = (
|
|
self.padding[2],
|
|
self.padding[2],
|
|
self.padding[1],
|
|
self.padding[1],
|
|
2 * self.padding[0],
|
|
0,
|
|
)
|
|
self.padding = (0, 0, 0)
|
|
|
|
def forward(self, x):
|
|
# Single-pass causal conv: the full temporal sequence is available at
|
|
# once, so the past is provided purely by the front temporal pad in
|
|
# ``self._padding`` (no feat_cache threading).
|
|
x = F.pad(x, self._padding)
|
|
x = super().forward(x)
|
|
return x
|
|
|
|
|
|
class RMS_norm(nn.Module):
|
|
|
|
def __init__(self, dim, channel_first=True, images=True, bias=False):
|
|
super().__init__()
|
|
broadcastable_dims = (1, 1, 1) if not images else (1, 1)
|
|
shape = (dim, *broadcastable_dims) if channel_first else (dim,)
|
|
|
|
self.channel_first = channel_first
|
|
self.scale = dim**0.5
|
|
self.gamma = nn.Parameter(torch.ones(shape))
|
|
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
|
|
|
|
def forward(self, x):
|
|
return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
|
|
self.scale * self.gamma + self.bias)
|
|
|
|
|
|
class Upsample(nn.Upsample):
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Fix bfloat16 support for nearest neighbor interpolation.
|
|
"""
|
|
return super().forward(x.float()).type_as(x)
|
|
|
|
|
|
class Resample(nn.Module):
|
|
|
|
def __init__(self, dim, mode):
|
|
assert mode in (
|
|
"none",
|
|
"upsample2d",
|
|
"upsample3d",
|
|
"downsample2d",
|
|
"downsample3d",
|
|
)
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.mode = mode
|
|
|
|
# layers
|
|
if mode == "upsample2d":
|
|
self.resample = nn.Sequential(
|
|
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
|
nn.Conv2d(dim, dim, 3, padding=1),
|
|
)
|
|
elif mode == "upsample3d":
|
|
self.resample = nn.Sequential(
|
|
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
|
nn.Conv2d(dim, dim, 3, padding=1),
|
|
# nn.Conv2d(dim, dim//2, 3, padding=1)
|
|
)
|
|
self.time_conv = CausalConv3d(
|
|
dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
|
|
elif mode == "downsample2d":
|
|
self.resample = nn.Sequential(
|
|
nn.ZeroPad2d((0, 1, 0, 1)),
|
|
nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
|
elif mode == "downsample3d":
|
|
self.resample = nn.Sequential(
|
|
nn.ZeroPad2d((0, 1, 0, 1)),
|
|
nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
|
self.time_conv = CausalConv3d(
|
|
dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
|
|
else:
|
|
self.resample = nn.Identity()
|
|
|
|
def forward(self, x):
|
|
b, c, t, h, w = x.size()
|
|
|
|
if self.mode == 'upsample3d':
|
|
# Causal temporal upsample without cache: the first frame is kept
|
|
# as-is and only the remaining frames are temporally doubled. This
|
|
# reproduces the chunked path (frame 0 emitted with zero left
|
|
# context, later frames with 2 frames of context).
|
|
first_frame = x[:, :, :1, :, :]
|
|
rest_frames = x[:, :, 1:, :, :] if t > 1 else None
|
|
if rest_frames is not None:
|
|
rest_frames = self.time_conv(rest_frames)
|
|
rest_frames = rest_frames.reshape(b, 2, c, t - 1, h, w)
|
|
rest_frames = torch.stack((rest_frames[:, 0], rest_frames[:, 1]),
|
|
3)
|
|
rest_frames = rest_frames.reshape(b, c, (t - 1) * 2, h, w)
|
|
x = torch.cat([first_frame, rest_frames], dim=2)
|
|
else:
|
|
x = first_frame
|
|
|
|
t = x.shape[2]
|
|
x = rearrange(x, 'b c t h w -> (b t) c h w')
|
|
x = self.resample(x)
|
|
x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
|
|
|
|
if self.mode == 'downsample3d':
|
|
# Causal temporal downsample without cache: first frame passes
|
|
# through unchanged and the stride-2 causal conv runs over the full
|
|
# sequence (its own windows supply the previous-frame left context),
|
|
# exactly reproducing the chunked cache-threaded path.
|
|
if x.shape[2] > 1:
|
|
first_frame = x[:, :, :1, :, :]
|
|
rest_frames = self.time_conv(x)
|
|
x = torch.cat([first_frame, rest_frames], dim=2)
|
|
|
|
return x
|
|
|
|
def init_weight(self, conv):
|
|
conv_weight = conv.weight.detach().clone()
|
|
nn.init.zeros_(conv_weight)
|
|
c1, c2, t, h, w = conv_weight.size()
|
|
one_matrix = torch.eye(c1, c2)
|
|
init_matrix = one_matrix
|
|
nn.init.zeros_(conv_weight)
|
|
conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5
|
|
conv.weight = nn.Parameter(conv_weight)
|
|
nn.init.zeros_(conv.bias.data)
|
|
|
|
def init_weight2(self, conv):
|
|
conv_weight = conv.weight.data.detach().clone()
|
|
nn.init.zeros_(conv_weight)
|
|
c1, c2, t, h, w = conv_weight.size()
|
|
init_matrix = torch.eye(c1 // 2, c2)
|
|
conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
|
|
conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
|
|
conv.weight = nn.Parameter(conv_weight)
|
|
nn.init.zeros_(conv.bias.data)
|
|
|
|
|
|
class ResidualBlock(nn.Module):
|
|
|
|
def __init__(self, in_dim, out_dim, dropout=0.0):
|
|
super().__init__()
|
|
self.in_dim = in_dim
|
|
self.out_dim = out_dim
|
|
|
|
# layers
|
|
self.residual = nn.Sequential(
|
|
RMS_norm(in_dim, images=False),
|
|
nn.SiLU(),
|
|
CausalConv3d(in_dim, out_dim, 3, padding=1),
|
|
RMS_norm(out_dim, images=False),
|
|
nn.SiLU(),
|
|
nn.Dropout(dropout),
|
|
CausalConv3d(out_dim, out_dim, 3, padding=1),
|
|
)
|
|
self.shortcut = (
|
|
CausalConv3d(in_dim, out_dim, 1)
|
|
if in_dim != out_dim else nn.Identity())
|
|
|
|
def forward(self, x):
|
|
h = self.shortcut(x)
|
|
x = self.residual(x)
|
|
return x + h
|
|
|
|
|
|
class AttentionBlock(nn.Module):
|
|
"""
|
|
Causal self-attention with a single head.
|
|
"""
|
|
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
# layers
|
|
self.norm = RMS_norm(dim)
|
|
self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
|
|
self.proj = nn.Conv2d(dim, dim, 1)
|
|
|
|
# zero out the last layer params
|
|
nn.init.zeros_(self.proj.weight)
|
|
|
|
def forward(self, x):
|
|
identity = x
|
|
b, c, t, h, w = x.size()
|
|
x = rearrange(x, "b c t h w -> (b t) c h w")
|
|
x = self.norm(x)
|
|
# compute query, key, value
|
|
q, k, v = (
|
|
self.to_qkv(x).reshape(b * t, 1, c * 3,
|
|
-1).permute(0, 1, 3,
|
|
2).contiguous().chunk(3, dim=-1))
|
|
|
|
# apply attention
|
|
x = F.scaled_dot_product_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
)
|
|
x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
|
|
|
|
# output
|
|
x = self.proj(x)
|
|
x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
|
|
return x + identity
|
|
|
|
|
|
def patchify(x, patch_size):
|
|
if patch_size == 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, r=patch_size)
|
|
elif x.dim() == 5:
|
|
x = rearrange(
|
|
x,
|
|
"b c f (h q) (w r) -> b (c r q) f h w",
|
|
q=patch_size,
|
|
r=patch_size,
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid input shape: {x.shape}")
|
|
|
|
return x
|
|
|
|
|
|
def unpatchify(x, patch_size):
|
|
if patch_size == 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, r=patch_size)
|
|
elif x.dim() == 5:
|
|
x = rearrange(
|
|
x,
|
|
"b (c r q) f h w -> b c f (h q) (w r)",
|
|
q=patch_size,
|
|
r=patch_size,
|
|
)
|
|
return x
|
|
|
|
|
|
class AvgDown3D(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
factor_t,
|
|
factor_s=1,
|
|
):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.factor_t = factor_t
|
|
self.factor_s = factor_s
|
|
self.factor = self.factor_t * self.factor_s * self.factor_s
|
|
|
|
assert in_channels * self.factor % out_channels == 0
|
|
self.group_size = in_channels * self.factor // out_channels
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
|
|
pad = (0, 0, 0, 0, pad_t, 0)
|
|
x = F.pad(x, pad)
|
|
B, C, T, H, W = x.shape
|
|
x = x.view(
|
|
B,
|
|
C,
|
|
T // self.factor_t,
|
|
self.factor_t,
|
|
H // self.factor_s,
|
|
self.factor_s,
|
|
W // self.factor_s,
|
|
self.factor_s,
|
|
)
|
|
x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
|
|
x = x.view(
|
|
B,
|
|
C * self.factor,
|
|
T // self.factor_t,
|
|
H // self.factor_s,
|
|
W // self.factor_s,
|
|
)
|
|
x = x.view(
|
|
B,
|
|
self.out_channels,
|
|
self.group_size,
|
|
T // self.factor_t,
|
|
H // self.factor_s,
|
|
W // self.factor_s,
|
|
)
|
|
x = x.mean(dim=2)
|
|
return x
|
|
|
|
|
|
class DupUp3D(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
factor_t,
|
|
factor_s=1,
|
|
):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
|
|
self.factor_t = factor_t
|
|
self.factor_s = factor_s
|
|
self.factor = self.factor_t * self.factor_s * self.factor_s
|
|
|
|
assert out_channels * self.factor % in_channels == 0
|
|
self.repeats = out_channels * self.factor // in_channels
|
|
|
|
def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
|
|
x = x.repeat_interleave(self.repeats, dim=1)
|
|
x = x.view(
|
|
x.size(0),
|
|
self.out_channels,
|
|
self.factor_t,
|
|
self.factor_s,
|
|
self.factor_s,
|
|
x.size(2),
|
|
x.size(3),
|
|
x.size(4),
|
|
)
|
|
x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
|
|
x = x.view(
|
|
x.size(0),
|
|
self.out_channels,
|
|
x.size(2) * self.factor_t,
|
|
x.size(4) * self.factor_s,
|
|
x.size(6) * self.factor_s,
|
|
)
|
|
if first_chunk:
|
|
x = x[:, :, self.factor_t - 1:, :, :]
|
|
return x
|
|
|
|
|
|
class Down_ResidualBlock(nn.Module):
|
|
|
|
def __init__(self,
|
|
in_dim,
|
|
out_dim,
|
|
dropout,
|
|
mult,
|
|
temperal_downsample=False,
|
|
down_flag=False):
|
|
super().__init__()
|
|
|
|
# Shortcut path with downsample
|
|
self.avg_shortcut = AvgDown3D(
|
|
in_dim,
|
|
out_dim,
|
|
factor_t=2 if temperal_downsample else 1,
|
|
factor_s=2 if down_flag else 1,
|
|
)
|
|
|
|
# Main path with residual blocks and downsample
|
|
downsamples = []
|
|
for _ in range(mult):
|
|
downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
|
in_dim = out_dim
|
|
|
|
# Add the final downsample block
|
|
if down_flag:
|
|
mode = "downsample3d" if temperal_downsample else "downsample2d"
|
|
downsamples.append(Resample(out_dim, mode=mode))
|
|
|
|
self.downsamples = nn.Sequential(*downsamples)
|
|
|
|
def forward(self, x):
|
|
x_copy = x
|
|
for module in self.downsamples:
|
|
x = module(x)
|
|
|
|
return x + self.avg_shortcut(x_copy)
|
|
|
|
|
|
class Up_ResidualBlock(nn.Module):
|
|
|
|
def __init__(self,
|
|
in_dim,
|
|
out_dim,
|
|
dropout,
|
|
mult,
|
|
temperal_upsample=False,
|
|
up_flag=False):
|
|
super().__init__()
|
|
# Shortcut path with upsample
|
|
if up_flag:
|
|
self.avg_shortcut = DupUp3D(
|
|
in_dim,
|
|
out_dim,
|
|
factor_t=2 if temperal_upsample else 1,
|
|
factor_s=2 if up_flag else 1,
|
|
)
|
|
else:
|
|
self.avg_shortcut = None
|
|
|
|
# Main path with residual blocks and upsample
|
|
upsamples = []
|
|
for _ in range(mult):
|
|
upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
|
in_dim = out_dim
|
|
|
|
# Add the final upsample block
|
|
if up_flag:
|
|
mode = "upsample3d" if temperal_upsample else "upsample2d"
|
|
upsamples.append(Resample(out_dim, mode=mode))
|
|
|
|
self.upsamples = nn.Sequential(*upsamples)
|
|
|
|
def forward(self, x):
|
|
# Single-pass: the whole latent sequence starts at frame 0, so the
|
|
# DupUp3D shortcut always applies the ``first_chunk`` front trim.
|
|
x_main = x
|
|
for module in self.upsamples:
|
|
x_main = module(x_main)
|
|
if self.avg_shortcut is not None:
|
|
x_shortcut = self.avg_shortcut(x, first_chunk=True)
|
|
return x_main + x_shortcut
|
|
else:
|
|
return x_main
|
|
|
|
|
|
class Encoder3d(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
dim=128,
|
|
z_dim=4,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_downsample=[True, True, False],
|
|
dropout=0.0,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.z_dim = z_dim
|
|
self.dim_mult = dim_mult
|
|
self.num_res_blocks = num_res_blocks
|
|
self.attn_scales = attn_scales
|
|
self.temperal_downsample = temperal_downsample
|
|
|
|
# dimensions
|
|
dims = [dim * u for u in [1] + dim_mult]
|
|
scale = 1.0
|
|
|
|
# init block
|
|
self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
|
|
|
|
# downsample blocks
|
|
downsamples = []
|
|
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
|
t_down_flag = (
|
|
temperal_downsample[i]
|
|
if i < len(temperal_downsample) else False)
|
|
downsamples.append(
|
|
Down_ResidualBlock(
|
|
in_dim=in_dim,
|
|
out_dim=out_dim,
|
|
dropout=dropout,
|
|
mult=num_res_blocks,
|
|
temperal_downsample=t_down_flag,
|
|
down_flag=i != len(dim_mult) - 1,
|
|
))
|
|
scale /= 2.0
|
|
self.downsamples = nn.Sequential(*downsamples)
|
|
|
|
# middle blocks
|
|
self.middle = nn.Sequential(
|
|
ResidualBlock(out_dim, out_dim, dropout),
|
|
AttentionBlock(out_dim),
|
|
ResidualBlock(out_dim, out_dim, dropout),
|
|
)
|
|
|
|
# # output blocks
|
|
self.head = nn.Sequential(
|
|
RMS_norm(out_dim, images=False),
|
|
nn.SiLU(),
|
|
CausalConv3d(out_dim, z_dim, 3, padding=1),
|
|
)
|
|
|
|
def forward(self, x):
|
|
# Single-pass encode: the whole (patchified) video goes through the
|
|
# encoder in one forward, no feat_cache / chunk loop.
|
|
x = self.conv1(x)
|
|
x = self.downsamples(x)
|
|
x = self.middle(x)
|
|
x = self.head(x)
|
|
return x
|
|
|
|
|
|
class Decoder3d(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
dim=128,
|
|
z_dim=4,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_upsample=[False, True, True],
|
|
dropout=0.0,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.z_dim = z_dim
|
|
self.dim_mult = dim_mult
|
|
self.num_res_blocks = num_res_blocks
|
|
self.attn_scales = attn_scales
|
|
self.temperal_upsample = temperal_upsample
|
|
|
|
# dimensions
|
|
dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
|
|
scale = 1.0 / 2**(len(dim_mult) - 2)
|
|
# init block
|
|
self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
|
|
|
|
# middle blocks
|
|
self.middle = nn.Sequential(
|
|
ResidualBlock(dims[0], dims[0], dropout),
|
|
AttentionBlock(dims[0]),
|
|
ResidualBlock(dims[0], dims[0], dropout),
|
|
)
|
|
|
|
# upsample blocks
|
|
upsamples = []
|
|
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
|
t_up_flag = temperal_upsample[i] if i < len(
|
|
temperal_upsample) else False
|
|
upsamples.append(
|
|
Up_ResidualBlock(
|
|
in_dim=in_dim,
|
|
out_dim=out_dim,
|
|
dropout=dropout,
|
|
mult=num_res_blocks + 1,
|
|
temperal_upsample=t_up_flag,
|
|
up_flag=i != len(dim_mult) - 1,
|
|
))
|
|
self.upsamples = nn.Sequential(*upsamples)
|
|
|
|
# output blocks
|
|
self.head = nn.Sequential(
|
|
RMS_norm(out_dim, images=False),
|
|
nn.SiLU(),
|
|
CausalConv3d(out_dim, 12, 3, padding=1),
|
|
)
|
|
|
|
# Toggle for single-pass per-block gradient checkpointing, propagated
|
|
# from AutoencoderKLWan2_2_.decode.
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, x):
|
|
# Single-pass decode. When gradient checkpointing is enabled each
|
|
# sub-block is recomputed during backward so only one block's
|
|
# activations stay alive at a time (vs. the whole upsampled video).
|
|
def _run(module, t):
|
|
if self.gradient_checkpointing and torch.is_grad_enabled():
|
|
return torch_checkpoint(module, t, use_reentrant=False)
|
|
return module(t)
|
|
|
|
x = self.conv1(x)
|
|
for layer in self.middle:
|
|
x = _run(layer, x)
|
|
for layer in self.upsamples:
|
|
x = _run(layer, x)
|
|
for layer in self.head:
|
|
x = _run(layer, x)
|
|
return x
|
|
|
|
|
|
def count_conv3d(model):
|
|
count = 0
|
|
for m in model.modules():
|
|
if isinstance(m, CausalConv3d):
|
|
count += 1
|
|
return count
|
|
|
|
|
|
class AutoencoderKLWan2_2_(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
dim=160,
|
|
dec_dim=256,
|
|
z_dim=16,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_downsample=[True, True, False],
|
|
dropout=0.0,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.z_dim = z_dim
|
|
self.dim_mult = dim_mult
|
|
self.num_res_blocks = num_res_blocks
|
|
self.attn_scales = attn_scales
|
|
self.temperal_downsample = temperal_downsample
|
|
self.temperal_upsample = temperal_downsample[::-1]
|
|
|
|
# modules
|
|
self.encoder = Encoder3d(
|
|
dim,
|
|
z_dim * 2,
|
|
dim_mult,
|
|
num_res_blocks,
|
|
attn_scales,
|
|
self.temperal_downsample,
|
|
dropout,
|
|
)
|
|
self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
|
|
self.conv2 = CausalConv3d(z_dim, z_dim, 1)
|
|
self.decoder = Decoder3d(
|
|
dec_dim,
|
|
z_dim,
|
|
dim_mult,
|
|
num_res_blocks,
|
|
attn_scales,
|
|
self.temperal_upsample,
|
|
dropout,
|
|
)
|
|
|
|
# When True the decoder runs with per-block gradient checkpointing.
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, x, scale=[0, 1]):
|
|
mu = self.encode(x, scale)
|
|
x_recon = self.decode(mu, scale)
|
|
return x_recon, mu
|
|
|
|
def encode(self, x, scale):
|
|
# Single-pass encode: patchify the full video and run the encoder once
|
|
# (no chunk loop, no feat_cache).
|
|
scale = [item.to(x.device, x.dtype) for item in scale]
|
|
x = patchify(x, patch_size=2)
|
|
out = self.encoder(x)
|
|
mu, log_var = self.conv1(out).chunk(2, dim=1)
|
|
if isinstance(scale[0], torch.Tensor):
|
|
mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
|
|
1, self.z_dim, 1, 1, 1)
|
|
else:
|
|
mu = (mu - scale[0]) * scale[1]
|
|
x = torch.cat([mu, log_var], dim=1)
|
|
return x
|
|
|
|
def decode(self, z, scale):
|
|
# Single-pass decode: run the whole latent sequence through the decoder
|
|
# in one forward. Because there is no feat_cache mutation, gradient
|
|
# checkpointing can be applied per decoder block (see Decoder3d.forward)
|
|
# instead of the previous per-frame cache-threading workaround.
|
|
scale = [item.to(z.device, z.dtype) for item in scale]
|
|
if isinstance(scale[0], torch.Tensor):
|
|
z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
|
|
1, self.z_dim, 1, 1, 1)
|
|
else:
|
|
z = z / scale[1] + scale[0]
|
|
self.decoder.gradient_checkpointing = self.gradient_checkpointing
|
|
x = self.conv2(z)
|
|
out = self.decoder(x)
|
|
out = unpatchify(out, patch_size=2)
|
|
return out
|
|
|
|
def reparameterize(self, mu, log_var):
|
|
std = torch.exp(0.5 * log_var)
|
|
eps = torch.randn_like(std)
|
|
return eps * std + mu
|
|
|
|
def sample(self, imgs, deterministic=False):
|
|
mu, log_var = self.encode(imgs)
|
|
if deterministic:
|
|
return mu
|
|
std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
|
|
return mu + std * torch.randn_like(std)
|
|
|
|
|
|
def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
|
|
# params
|
|
cfg = dict(
|
|
dim=dim,
|
|
z_dim=z_dim,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_downsample=[True, True, True],
|
|
dropout=0.0,
|
|
)
|
|
cfg.update(**kwargs)
|
|
|
|
# init model
|
|
model = AutoencoderKLWan2_2_(**cfg)
|
|
|
|
return model
|
|
|
|
|
|
class AutoencoderKLWan3_8_1Steps(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
|
_supports_gradient_checkpointing = True
|
|
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
latent_channels=48,
|
|
c_dim=160,
|
|
vae_pth=None,
|
|
dim_mult=[1, 2, 4, 4],
|
|
temperal_downsample=[False, True, True],
|
|
temporal_compression_ratio=4,
|
|
spatial_compression_ratio=16
|
|
):
|
|
super().__init__()
|
|
mean = torch.tensor(
|
|
[
|
|
-0.2289,
|
|
-0.0052,
|
|
-0.1323,
|
|
-0.2339,
|
|
-0.2799,
|
|
0.0174,
|
|
0.1838,
|
|
0.1557,
|
|
-0.1382,
|
|
0.0542,
|
|
0.2813,
|
|
0.0891,
|
|
0.1570,
|
|
-0.0098,
|
|
0.0375,
|
|
-0.1825,
|
|
-0.2246,
|
|
-0.1207,
|
|
-0.0698,
|
|
0.5109,
|
|
0.2665,
|
|
-0.2108,
|
|
-0.2158,
|
|
0.2502,
|
|
-0.2055,
|
|
-0.0322,
|
|
0.1109,
|
|
0.1567,
|
|
-0.0729,
|
|
0.0899,
|
|
-0.2799,
|
|
-0.1230,
|
|
-0.0313,
|
|
-0.1649,
|
|
0.0117,
|
|
0.0723,
|
|
-0.2839,
|
|
-0.2083,
|
|
-0.0520,
|
|
0.3748,
|
|
0.0152,
|
|
0.1957,
|
|
0.1433,
|
|
-0.2944,
|
|
0.3573,
|
|
-0.0548,
|
|
-0.1681,
|
|
-0.0667,
|
|
], dtype=torch.float32
|
|
)
|
|
std = torch.tensor(
|
|
[
|
|
0.4765,
|
|
1.0364,
|
|
0.4514,
|
|
1.1677,
|
|
0.5313,
|
|
0.4990,
|
|
0.4818,
|
|
0.5013,
|
|
0.8158,
|
|
1.0344,
|
|
0.5894,
|
|
1.0901,
|
|
0.6885,
|
|
0.6165,
|
|
0.8454,
|
|
0.4978,
|
|
0.5759,
|
|
0.3523,
|
|
0.7135,
|
|
0.6804,
|
|
0.5833,
|
|
1.4146,
|
|
0.8986,
|
|
0.5659,
|
|
0.7069,
|
|
0.5338,
|
|
0.4889,
|
|
0.4917,
|
|
0.4069,
|
|
0.4999,
|
|
0.6866,
|
|
0.4093,
|
|
0.5709,
|
|
0.6065,
|
|
0.6415,
|
|
0.4944,
|
|
0.5726,
|
|
1.2042,
|
|
0.5458,
|
|
1.6887,
|
|
0.3971,
|
|
1.0600,
|
|
0.3943,
|
|
0.5537,
|
|
0.5444,
|
|
0.4089,
|
|
0.7468,
|
|
0.7744,
|
|
], dtype=torch.float32
|
|
)
|
|
self.scale = [mean, 1.0 / std]
|
|
|
|
# init model
|
|
self.model = _video_vae(
|
|
pretrained_path=vae_pth,
|
|
z_dim=latent_channels,
|
|
dim=c_dim,
|
|
dim_mult=dim_mult,
|
|
temperal_downsample=temperal_downsample,
|
|
).eval().requires_grad_(False)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def _set_gradient_checkpointing(self, *args, **kwargs):
|
|
if "value" in kwargs:
|
|
self.gradient_checkpointing = kwargs["value"]
|
|
elif "enable" in kwargs:
|
|
self.gradient_checkpointing = kwargs["enable"]
|
|
else:
|
|
raise ValueError("Invalid set gradient checkpointing")
|
|
self.model.gradient_checkpointing = self.gradient_checkpointing
|
|
|
|
def _encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = [
|
|
self.model.encode(u.unsqueeze(0), self.scale).squeeze(0)
|
|
for u in x
|
|
]
|
|
x = torch.stack(x)
|
|
return x
|
|
|
|
@apply_forward_hook
|
|
def encode(
|
|
self, x: torch.Tensor, return_dict: bool = True
|
|
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
|
|
h = self._encode(x)
|
|
|
|
posterior = DiagonalGaussianDistribution(h)
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
def _decode(self, zs):
|
|
# Single-pass decode with per-block gradient checkpointing inside the
|
|
# decoder (see Decoder3d.forward), which lowers the backward peak to a
|
|
# single decoder block's activations while keeping the whole video in
|
|
# one forward (no feat_cache chunk loop).
|
|
self.model.gradient_checkpointing = self.gradient_checkpointing
|
|
dec = [
|
|
self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0)
|
|
for u in zs
|
|
]
|
|
dec = torch.stack(dec)
|
|
|
|
return DecoderOutput(sample=dec)
|
|
|
|
@apply_forward_hook
|
|
def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
|
|
decoded = self._decode(z).sample
|
|
|
|
if not return_dict:
|
|
return (decoded,)
|
|
return DecoderOutput(sample=decoded)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_path, additional_kwargs={}):
|
|
def filter_kwargs(cls, kwargs):
|
|
import inspect
|
|
sig = inspect.signature(cls.__init__)
|
|
valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
|
|
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
|
|
return filtered_kwargs
|
|
|
|
model = cls(**filter_kwargs(cls, additional_kwargs))
|
|
if pretrained_model_path.endswith(".safetensors"):
|
|
from safetensors.torch import load_file, safe_open
|
|
state_dict = load_file(pretrained_model_path)
|
|
else:
|
|
state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True)
|
|
tmp_state_dict = {}
|
|
for key in state_dict:
|
|
if not key.startswith("model."):
|
|
tmp_state_dict["model." + key] = state_dict[key]
|
|
else:
|
|
tmp_state_dict[key] = state_dict[key]
|
|
state_dict = tmp_state_dict
|
|
m, u = model.load_state_dict(state_dict, strict=False)
|
|
print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
|
|
print(m, u)
|
|
return model |