1198 lines
46 KiB
Python
Executable File
1198 lines
46 KiB
Python
Executable File
# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/vae.py
|
|
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
|
from typing import Tuple, Union
|
|
|
|
import torch
|
|
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, cache_x=None):
|
|
padding = list(self._padding)
|
|
if cache_x is not None and self._padding[4] > 0:
|
|
cache_x = cache_x.to(x.device)
|
|
x = torch.cat([cache_x, x], dim=2)
|
|
padding[4] -= cache_x.shape[2]
|
|
x = F.pad(x, padding)
|
|
|
|
return super().forward(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.
|
|
|
|
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., 2.), mode='nearest-exact'),
|
|
nn.Conv2d(dim, dim // 2, 3, padding=1))
|
|
elif mode == 'upsample3d':
|
|
self.resample = nn.Sequential(
|
|
Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
|
|
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, feat_cache=None, feat_idx=[0]):
|
|
b, c, t, h, w = x.size()
|
|
if self.mode == 'upsample3d':
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
if feat_cache[idx] is None:
|
|
feat_cache[idx] = 'Rep'
|
|
feat_idx[0] += 1
|
|
else:
|
|
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[
|
|
idx] is not None and feat_cache[idx] != 'Rep':
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
if cache_x.shape[2] < 2 and feat_cache[
|
|
idx] is not None and feat_cache[idx] == 'Rep':
|
|
cache_x = torch.cat([
|
|
torch.zeros_like(cache_x).to(cache_x.device),
|
|
cache_x
|
|
],
|
|
dim=2)
|
|
if feat_cache[idx] == 'Rep':
|
|
x = self.time_conv(x)
|
|
else:
|
|
x = self.time_conv(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
|
|
x = x.reshape(b, 2, c, t, h, w)
|
|
x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
|
|
3)
|
|
x = x.reshape(b, c, t * 2, h, w)
|
|
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':
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
if feat_cache[idx] is None:
|
|
feat_cache[idx] = x.clone()
|
|
feat_idx[0] += 1
|
|
else:
|
|
|
|
cache_x = x[:, :, -1:, :, :].clone()
|
|
# if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
|
|
# # cache last frame of last two chunk
|
|
# cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
|
|
|
|
x = self.time_conv(
|
|
torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
return x
|
|
|
|
def init_weight(self, conv):
|
|
conv_weight = conv.weight
|
|
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,1,1] = init_matrix * 0.5
|
|
conv_weight.data[:, :, 1, 0, 0] = init_matrix #* 0.5
|
|
conv.weight.data.copy_(conv_weight)
|
|
nn.init.zeros_(conv.bias.data)
|
|
|
|
def init_weight2(self, conv):
|
|
conv_weight = conv.weight.data
|
|
nn.init.zeros_(conv_weight)
|
|
c1, c2, t, h, w = conv_weight.size()
|
|
init_matrix = torch.eye(c1 // 2, c2)
|
|
#init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
|
|
conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
|
|
conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
|
|
conv.weight.data.copy_(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, feat_cache=None, feat_idx=[0]):
|
|
h = self.shortcut(x)
|
|
for layer in self.residual:
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
x = layer(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = layer(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
|
|
|
|
|
|
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(3, dims[0], 3, padding=1)
|
|
|
|
# downsample blocks
|
|
downsamples = []
|
|
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
|
# residual (+attention) blocks
|
|
for _ in range(num_res_blocks):
|
|
downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
|
if scale in attn_scales:
|
|
downsamples.append(AttentionBlock(out_dim))
|
|
in_dim = out_dim
|
|
|
|
# downsample block
|
|
if i != len(dim_mult) - 1:
|
|
mode = 'downsample3d' if temperal_downsample[
|
|
i] else 'downsample2d'
|
|
downsamples.append(Resample(out_dim, mode=mode))
|
|
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, feat_cache=None, feat_idx=[0]):
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
x = self.conv1(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv1(x)
|
|
|
|
## downsamples
|
|
for layer in self.downsamples:
|
|
if feat_cache is not None:
|
|
x = layer(x, feat_cache, feat_idx)
|
|
else:
|
|
x = layer(x)
|
|
|
|
## middle
|
|
for layer in self.middle:
|
|
if isinstance(layer, ResidualBlock) and feat_cache is not None:
|
|
x = layer(x, feat_cache, feat_idx)
|
|
else:
|
|
x = layer(x)
|
|
|
|
## head
|
|
for layer in self.head:
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
x = layer(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = layer(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:])):
|
|
# residual (+attention) blocks
|
|
if i == 1 or i == 2 or i == 3:
|
|
in_dim = in_dim // 2
|
|
for _ in range(num_res_blocks + 1):
|
|
upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
|
if scale in attn_scales:
|
|
upsamples.append(AttentionBlock(out_dim))
|
|
in_dim = out_dim
|
|
|
|
# upsample block
|
|
if i != len(dim_mult) - 1:
|
|
mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
|
|
upsamples.append(Resample(out_dim, mode=mode))
|
|
scale *= 2.0
|
|
self.upsamples = nn.Sequential(*upsamples)
|
|
|
|
# output blocks
|
|
self.head = nn.Sequential(
|
|
RMS_norm(out_dim, images=False), nn.SiLU(),
|
|
CausalConv3d(out_dim, 3, 3, padding=1))
|
|
|
|
def forward(self, x, feat_cache=None, feat_idx=[0]):
|
|
## conv1
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
x = self.conv1(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv1(x)
|
|
|
|
## middle
|
|
for layer in self.middle:
|
|
if isinstance(layer, ResidualBlock) and feat_cache is not None:
|
|
x = layer(x, feat_cache, feat_idx)
|
|
else:
|
|
x = layer(x)
|
|
|
|
## upsamples
|
|
for layer in self.upsamples:
|
|
if feat_cache is not None:
|
|
x = layer(x, feat_cache, feat_idx)
|
|
else:
|
|
x = layer(x)
|
|
|
|
## head
|
|
for layer in self.head:
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([
|
|
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
|
|
cache_x.device), cache_x
|
|
],
|
|
dim=2)
|
|
x = layer(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
def count_conv3d(model):
|
|
count = 0
|
|
for m in model.modules():
|
|
if isinstance(m, CausalConv3d):
|
|
count += 1
|
|
return count
|
|
|
|
|
|
class AutoencoderKLWan_(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
|
|
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(dim, z_dim, dim_mult, num_res_blocks,
|
|
attn_scales, self.temperal_upsample, dropout)
|
|
|
|
def forward(self, x):
|
|
mu, log_var = self.encode(x)
|
|
z = self.reparameterize(mu, log_var)
|
|
x_recon = self.decode(z)
|
|
return x_recon, mu, log_var
|
|
|
|
def encode(self, x, scale=None):
|
|
self.clear_cache()
|
|
## cache
|
|
t = x.shape[2]
|
|
iter_ = 1 + (t - 1) // 4
|
|
if scale != None:
|
|
scale = [item.to(x.device, x.dtype) for item in scale]
|
|
## 对encode输入的x，按时间拆分为1、4、4、4....
|
|
for i in range(iter_):
|
|
self._enc_conv_idx = [0]
|
|
if i == 0:
|
|
out = self.encoder(
|
|
x[:, :, :1, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx)
|
|
else:
|
|
out_ = self.encoder(
|
|
x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx)
|
|
out = torch.cat([out, out_], 2)
|
|
mu, log_var = self.conv1(out).chunk(2, dim=1)
|
|
if scale != None:
|
|
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)
|
|
self.clear_cache()
|
|
return x
|
|
|
|
def encode_stream(self, x, scale=None, is_first_chunk=True):
|
|
# Streaming variant of `encode`. Unlike `encode`, it does NOT call
|
|
# `clear_cache()` at the start/end, so the causal temporal feature
|
|
# cache (`self._enc_feat_map`, mutated in place by the encoder's
|
|
# CausalConv3d layers) is threaded across successive pixel chunks. The
|
|
# caller is responsible for calling `clear_cache()` exactly once before
|
|
# the first chunk and once after the last chunk. Encoding the pixel
|
|
# sequence chunk-by-chunk this way is mathematically equivalent to a
|
|
# single `encode` call over the whole sequence (no boundary seams).
|
|
#
|
|
# Temporal alignment (mirrors the VAE's 1,4,4,... compression):
|
|
# - is_first_chunk=True : chunk contains the very first pixel frame;
|
|
# `t` must be `1 + 4*k` (k>=0) -> yields `1 + k` latent frames.
|
|
# - is_first_chunk=False: continuation chunk; `t` must be a multiple
|
|
# of 4 -> yields `t // 4` latent frames.
|
|
# x: [b,c,t,h,w]
|
|
t = x.shape[2]
|
|
if scale is not None:
|
|
scale = [item.to(x.device, x.dtype) for item in scale]
|
|
out = None
|
|
if is_first_chunk:
|
|
assert (t - 1) % 4 == 0, \
|
|
f"first streaming encode chunk needs t=1+4*k frames, got t={t}"
|
|
iter_ = 1 + (t - 1) // 4
|
|
for i in range(iter_):
|
|
self._enc_conv_idx = [0]
|
|
if i == 0:
|
|
cur = self.encoder(
|
|
x[:, :, :1, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx)
|
|
else:
|
|
cur = self.encoder(
|
|
x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx)
|
|
out = cur if out is None else torch.cat([out, cur], 2)
|
|
else:
|
|
assert t % 4 == 0, \
|
|
f"continuation streaming encode chunk needs t=4*k frames, got t={t}"
|
|
iter_ = t // 4
|
|
for i in range(iter_):
|
|
self._enc_conv_idx = [0]
|
|
cur = self.encoder(
|
|
x[:, :, 4 * i:4 * (i + 1), :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx)
|
|
out = cur if out is None else torch.cat([out, cur], 2)
|
|
# conv1 has temporal kernel size 1, so applying it per chunk is
|
|
# identical to applying it once over the concatenated sequence.
|
|
mu, log_var = self.conv1(out).chunk(2, dim=1)
|
|
if scale is not None:
|
|
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_frame(self, frame, feat_map_in):
|
|
# Decode a single latent frame with a functional (non-mutating) cache
|
|
# so it can be wrapped by gradient checkpointing safely. The input
|
|
# cache list is copied and never mutated in place, and the updated
|
|
# cache is returned explicitly to be threaded to the next frame.
|
|
conv_idx = [0]
|
|
feat_map = list(feat_map_in)
|
|
out = self.decoder(frame, feat_cache=feat_map, feat_idx=conv_idx)
|
|
return out, feat_map
|
|
|
|
def decode(self, z, scale=None):
|
|
self.clear_cache()
|
|
# z: [b,c,t,h,w]
|
|
if scale != None:
|
|
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]
|
|
iter_ = z.shape[2]
|
|
x = self.conv2(z)
|
|
# Per-frame gradient checkpointing: each latent frame is an independent
|
|
# recompute unit, so backward only materializes one frame's decoder
|
|
# activations at a time instead of all frames at once.
|
|
use_ckpt = getattr(self, "gradient_checkpointing", False) \
|
|
and torch.is_grad_enabled()
|
|
feat_map = self._feat_map
|
|
outs = []
|
|
for i in range(iter_):
|
|
frame = x[:, :, i:i + 1, :, :]
|
|
if use_ckpt:
|
|
out_, feat_map = torch_checkpoint(
|
|
self._decode_frame, frame, feat_map,
|
|
use_reentrant=False)
|
|
else:
|
|
out_, feat_map = self._decode_frame(frame, feat_map)
|
|
outs.append(out_)
|
|
out = torch.cat(outs, 2)
|
|
self.clear_cache()
|
|
return out
|
|
|
|
def decode_stream(self, z, scale=None):
|
|
# Streaming variant of `decode`. Unlike `decode`, it does NOT call
|
|
# `clear_cache()` at the start/end, so the causal temporal feature
|
|
# cache (`self._feat_map`) is threaded across successive chunks. The
|
|
# caller is responsible for calling `clear_cache()` exactly once before
|
|
# the first chunk and once after the last chunk. Decoding the latent
|
|
# sequence chunk-by-chunk this way is mathematically equivalent to a
|
|
# single `decode` call over the whole sequence (no boundary seams),
|
|
# while enabling "generate a chunk, decode a chunk" streaming.
|
|
# z: [b,c,t,h,w]
|
|
if scale is not None:
|
|
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]
|
|
iter_ = z.shape[2]
|
|
x = self.conv2(z)
|
|
use_ckpt = getattr(self, "gradient_checkpointing", False) \
|
|
and torch.is_grad_enabled()
|
|
feat_map = self._feat_map
|
|
outs = []
|
|
for i in range(iter_):
|
|
frame = x[:, :, i:i + 1, :, :]
|
|
if use_ckpt:
|
|
out_, feat_map = torch_checkpoint(
|
|
self._decode_frame, frame, feat_map,
|
|
use_reentrant=False)
|
|
else:
|
|
out_, feat_map = self._decode_frame(frame, feat_map)
|
|
outs.append(out_)
|
|
# Persist the updated cache so the next chunk continues seamlessly.
|
|
self._feat_map = feat_map
|
|
out = torch.cat(outs, 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 clear_cache(self):
|
|
self._conv_num = count_conv3d(self.decoder)
|
|
self._conv_idx = [0]
|
|
self._feat_map = [None] * self._conv_num
|
|
#cache encode
|
|
self._enc_conv_num = count_conv3d(self.encoder)
|
|
self._enc_conv_idx = [0]
|
|
self._enc_feat_map = [None] * self._enc_conv_num
|
|
|
|
|
|
def _video_vae(z_dim=None, **kwargs):
|
|
"""
|
|
Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
|
|
"""
|
|
# params
|
|
cfg = dict(
|
|
dim=96,
|
|
z_dim=z_dim,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_downsample=[False, True, True],
|
|
dropout=0.0)
|
|
cfg.update(**kwargs)
|
|
|
|
# init model
|
|
model = AutoencoderKLWan_(**cfg)
|
|
|
|
return model
|
|
|
|
|
|
class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
|
_supports_gradient_checkpointing = True
|
|
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
latent_channels=16,
|
|
temporal_compression_ratio=4,
|
|
spatial_compression_ratio=8
|
|
):
|
|
super().__init__()
|
|
mean = [
|
|
-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
|
|
0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
|
|
]
|
|
std = [
|
|
2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
|
|
3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
|
|
]
|
|
self.mean = torch.tensor(mean, dtype=torch.float32)
|
|
self.std = torch.tensor(std, dtype=torch.float32)
|
|
self.scale = [self.mean, 1.0 / self.std]
|
|
|
|
# init model
|
|
self.model = _video_vae(
|
|
z_dim=latent_channels,
|
|
)
|
|
|
|
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 _encode_stream(self, x, is_first_chunk=True):
|
|
# Streaming encode of one pixel chunk over the full batch at once. The
|
|
# whole batch must be encoded in a single call so that the persistent
|
|
# `self.model._enc_feat_map` stays consistent across chunks (a per-sample
|
|
# loop would clobber the cache between samples/chunks).
|
|
return self.model.encode_stream(x, self.scale, is_first_chunk=is_first_chunk)
|
|
|
|
@apply_forward_hook
|
|
def encode_stream(
|
|
self, x: torch.Tensor, is_first_chunk: bool = True, return_dict: bool = True
|
|
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
|
|
h = self._encode_stream(x, is_first_chunk=is_first_chunk)
|
|
|
|
posterior = DiagonalGaussianDistribution(h)
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
def _decode(self, zs):
|
|
# Gradient checkpointing is applied per-frame inside self.model.decode
|
|
# (see AutoencoderKLWan_.decode), which lowers the backward peak to a
|
|
# single latent frame's decoder activations.
|
|
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)
|
|
|
|
def clear_cache(self):
|
|
# Reset the causal temporal feature cache. Must be called once before
|
|
# the first streaming chunk and once after the last one.
|
|
self.model.clear_cache()
|
|
|
|
def _decode_stream(self, zs):
|
|
# Streaming decode of one chunk over the full batch at once. The whole
|
|
# batch must be decoded in a single call so that the persistent
|
|
# `self.model._feat_map` stays consistent across chunks (a per-sample
|
|
# loop would clobber the cache between samples/chunks).
|
|
self.model.gradient_checkpointing = self.gradient_checkpointing
|
|
dec = self.model.decode_stream(zs, self.scale).clamp_(-1, 1)
|
|
return DecoderOutput(sample=dec)
|
|
|
|
@apply_forward_hook
|
|
def decode_stream(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
|
|
decoded = self._decode_stream(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
|
|
|
|
def convert_state_dict_keys(state_dict):
|
|
"""
|
|
Convert old checkpoint keys to new format for compatibility.
|
|
"""
|
|
import re
|
|
new_state_dict = {}
|
|
|
|
for old_key in state_dict:
|
|
new_key = old_key
|
|
|
|
# Map quant_conv -> conv1, post_quant_conv -> conv2
|
|
if '.quant_conv.' in old_key and 'post_quant_conv' not in old_key:
|
|
new_key = old_key.replace('.quant_conv.', '.conv1.')
|
|
elif '.post_quant_conv.' in old_key:
|
|
new_key = old_key.replace('.post_quant_conv.', '.conv2.')
|
|
|
|
# Map mid_block.attentions -> middle (attention is at index 1)
|
|
if '.encoder.mid_block.attentions.' in old_key:
|
|
match = re.match(r'(.*)\.encoder\.mid_block\.attentions\.(\d+)\.(norm|to_qkv|proj)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
new_key = f"{prefix}.encoder.middle.1.{layer_name}.{param_name}"
|
|
else:
|
|
print(f"Warning: unrecognized key pattern '{old_key}', keeping original key.")
|
|
|
|
if '.decoder.mid_block.attentions.' in old_key:
|
|
match = re.match(r'(.*)\.decoder\.mid_block\.attentions\.(\d+)\.(norm|to_qkv|proj)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
new_key = f"{prefix}.decoder.middle.1.{layer_name}.{param_name}"
|
|
else:
|
|
print(f"Warning: unrecognized key pattern '{old_key}', keeping original key.")
|
|
|
|
# 1. Map encoder.down_blocks -> encoder.downsamples
|
|
if '.encoder.down_blocks.' in old_key:
|
|
match = re.match(r'(.*)\.encoder\.down_blocks\.(\d+)\.(conv1|conv2|norm1|norm2|conv_shortcut|resample\.\d+|time_conv)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
block_idx = match.group(2)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
|
|
layer_map = {
|
|
'conv1': 'residual.2',
|
|
'conv2': 'residual.6',
|
|
'norm1': 'residual.0',
|
|
'norm2': 'residual.3',
|
|
'conv_shortcut': 'shortcut'
|
|
}
|
|
new_layer = layer_map.get(layer_name, layer_name)
|
|
new_key = f"{prefix}.encoder.downsamples.{block_idx}.{new_layer}.{param_name}"
|
|
|
|
# 2. Map decoder.up_blocks.X.resnets.Y -> decoder.upsamples
|
|
elif '.decoder.up_blocks.' in old_key and '.resnets.' in old_key:
|
|
match = re.match(r'(.*)\.decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.(conv1|conv2|norm1|norm2|conv_shortcut)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
block_idx = match.group(2)
|
|
resnet_idx = match.group(3)
|
|
layer_name = match.group(4)
|
|
param_name = match.group(5)
|
|
|
|
# Calculate upsample index based on actual structure
|
|
# up_blocks.0: resnets 0,1,2 -> upsamples 0,1,2
|
|
# up_blocks.1: resnets 0,1,2 -> upsamples 4,5,6 (3 is upsampler)
|
|
# up_blocks.2: resnets 0,1,2 -> upsamples 8,9,10 (7 is upsampler)
|
|
# up_blocks.3: resnets 0,1,2 -> upsamples 12,13,14 (11 is upsampler)
|
|
block_start = {0: 0, 1: 4, 2: 8, 3: 12}
|
|
upsample_idx = block_start.get(int(block_idx), 0) + int(resnet_idx)
|
|
|
|
layer_map = {
|
|
'conv1': 'residual.2',
|
|
'conv2': 'residual.6',
|
|
'norm1': 'residual.0',
|
|
'norm2': 'residual.3',
|
|
'conv_shortcut': 'shortcut'
|
|
}
|
|
new_layer = layer_map.get(layer_name, layer_name)
|
|
new_key = f"{prefix}.decoder.upsamples.{upsample_idx}.{new_layer}.{param_name}"
|
|
|
|
# 3. Map decoder.up_blocks.X.upsamplers -> decoder.upsamples
|
|
elif '.decoder.up_blocks.' in old_key and '.upsamplers.' in old_key:
|
|
match = re.match(r'(.*)\.decoder\.up_blocks\.(\d+)\.upsamplers\.0\.(resample\.\d+|time_conv)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
block_idx = match.group(2)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
|
|
# upsamplers are at indices 3, 7, 11
|
|
upsampler_idx = {0: 3, 1: 7, 2: 11, 3: None}
|
|
idx = upsampler_idx.get(int(block_idx))
|
|
if idx is not None:
|
|
new_key = f"{prefix}.decoder.upsamples.{idx}.{layer_name}.{param_name}"
|
|
else:
|
|
continue # Skip if no upsampler for this block
|
|
|
|
# 4. Map encoder.mid_block.resnets -> encoder.middle
|
|
elif '.encoder.mid_block.resnets.' in old_key:
|
|
match = re.match(r'(.*)\.encoder\.mid_block\.resnets\.(\d+)\.(conv1|conv2|norm1|norm2)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
resnet_idx = match.group(2)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
|
|
# middle structure: [ResidualBlock(0), AttentionBlock(1), ResidualBlock(2)]
|
|
middle_idx = int(resnet_idx) * 2
|
|
|
|
layer_map = {
|
|
'conv1': 'residual.2',
|
|
'conv2': 'residual.6',
|
|
'norm1': 'residual.0',
|
|
'norm2': 'residual.3'
|
|
}
|
|
new_layer = layer_map.get(layer_name, layer_name)
|
|
new_key = f"{prefix}.encoder.middle.{middle_idx}.{new_layer}.{param_name}"
|
|
|
|
# 5. Map decoder.mid_block.resnets -> decoder.middle
|
|
elif '.decoder.mid_block.resnets.' in old_key:
|
|
match = re.match(r'(.*)\.decoder\.mid_block\.resnets\.(\d+)\.(conv1|conv2|norm1|norm2)\.(.*)', old_key)
|
|
if match:
|
|
prefix = match.group(1)
|
|
resnet_idx = match.group(2)
|
|
layer_name = match.group(3)
|
|
param_name = match.group(4)
|
|
|
|
middle_idx = int(resnet_idx) * 2
|
|
|
|
layer_map = {
|
|
'conv1': 'residual.2',
|
|
'conv2': 'residual.6',
|
|
'norm1': 'residual.0',
|
|
'norm2': 'residual.3'
|
|
}
|
|
new_layer = layer_map.get(layer_name, layer_name)
|
|
new_key = f"{prefix}.decoder.middle.{middle_idx}.{new_layer}.{param_name}"
|
|
|
|
# 6. Map encoder.norm_out -> encoder.head
|
|
elif '.encoder.norm_out.' in old_key:
|
|
new_key = old_key.replace('.encoder.norm_out.', '.encoder.head.0.')
|
|
|
|
# 7. Map decoder.norm_out -> decoder.head
|
|
elif '.decoder.norm_out.' in old_key:
|
|
new_key = old_key.replace('.decoder.norm_out.', '.decoder.head.0.')
|
|
|
|
# 9. Map conv_in -> conv1
|
|
if '.encoder.conv_in.' in new_key:
|
|
new_key = new_key.replace('.encoder.conv_in.', '.encoder.conv1.')
|
|
if '.decoder.conv_in.' in new_key:
|
|
new_key = new_key.replace('.decoder.conv_in.', '.decoder.conv1.')
|
|
|
|
# 10. Map conv_out -> head.2
|
|
if '.encoder.conv_out.' in new_key:
|
|
new_key = new_key.replace('.encoder.conv_out.', '.encoder.head.2.')
|
|
if '.decoder.conv_out.' in new_key:
|
|
new_key = new_key.replace('.decoder.conv_out.', '.decoder.head.2.')
|
|
|
|
new_state_dict[new_key] = state_dict[old_key]
|
|
|
|
return new_state_dict
|
|
|
|
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)
|
|
|
|
# Add model. prefix
|
|
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
|
|
# Convert keys for old checkpoint compatibility
|
|
state_dict = convert_state_dict_keys(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
|
|
|
|
|
|
class AutoencoderKLWanCompileQwenImage(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
attn_scales = [],
|
|
base_dim = 96,
|
|
dim_mult = [
|
|
1,
|
|
2,
|
|
4,
|
|
4
|
|
],
|
|
dropout = 0.0,
|
|
num_res_blocks = 2,
|
|
temperal_downsample = [
|
|
False,
|
|
True,
|
|
True
|
|
],
|
|
z_dim = 16,
|
|
latents_mean = [
|
|
-0.7571,
|
|
-0.7089,
|
|
-0.9113,
|
|
0.1075,
|
|
-0.1745,
|
|
0.9653,
|
|
-0.1517,
|
|
1.5508,
|
|
0.4134,
|
|
-0.0715,
|
|
0.5517,
|
|
-0.3632,
|
|
-0.1922,
|
|
-0.9497,
|
|
0.2503,
|
|
-0.2921
|
|
],
|
|
latents_std = [
|
|
2.8184,
|
|
1.4541,
|
|
2.3275,
|
|
2.6558,
|
|
1.2196,
|
|
1.7708,
|
|
2.6052,
|
|
2.0743,
|
|
3.2687,
|
|
2.1526,
|
|
2.8652,
|
|
1.5579,
|
|
1.6382,
|
|
1.1253,
|
|
2.8251,
|
|
1.916
|
|
],
|
|
temporal_compression_ratio=4,
|
|
spatial_compression_ratio=8
|
|
):
|
|
super().__init__()
|
|
cfg = dict(
|
|
dim=base_dim,
|
|
z_dim=z_dim,
|
|
dim_mult=dim_mult,
|
|
num_res_blocks=num_res_blocks,
|
|
attn_scales=attn_scales,
|
|
temperal_downsample=temperal_downsample,
|
|
dropout=dropout)
|
|
|
|
# init model
|
|
self.model = AutoencoderKLWan_(**cfg)
|
|
|
|
self.dim = base_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]
|
|
self.temporal_compression_ratio = temporal_compression_ratio
|
|
self.spatial_compression_ratio = spatial_compression_ratio
|
|
|
|
def _encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = [
|
|
self.model.encode(u.unsqueeze(0)).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):
|
|
dec = [
|
|
self.model.decode(u.unsqueeze(0)).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 |