797 lines
31 KiB
Python
797 lines
31 KiB
Python
# Ported from https://github.com/madebyollin/taehv
|
|
# Tiny AutoEncoder for Wan 2.1 (taew2_1, 16ch latent) and Wan 2.2 5B (taew2_2, 48ch latent).
|
|
# TAE shares the exact latent space of the full-size Wan VAEs, so it can be used
|
|
# as a cheap drop-in encoder/decoder (e.g. for live previewing or low-memory decoding)
|
|
# at the cost of slightly lower reconstruction quality.
|
|
from typing import Tuple, Union
|
|
|
|
import os
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from collections import namedtuple
|
|
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 tqdm.auto import tqdm
|
|
|
|
# Deterministic-logvar used to wrap the (non-stochastic) TAE encoder output
|
|
# into a DiagonalGaussianDistribution compatible with the full Wan VAE interface.
|
|
DETERMINISTIC_LOGVAR = -30.0
|
|
|
|
TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
|
|
|
|
|
|
def _load_raw_state_dict(path):
|
|
"""Load a flat tensor state_dict from a .safetensors/.pth/.bin checkpoint.
|
|
|
|
Mirrors the normalization done in AutoencoderTinyWan.from_pretrained so that a
|
|
companion encoder checkpoint can be borrowed the same way.
|
|
"""
|
|
path = os.fspath(path)
|
|
if path.endswith(".safetensors"):
|
|
from safetensors.torch import load_file
|
|
|
|
sd = load_file(path)
|
|
else:
|
|
sd = torch.load(path, map_location="cpu", weights_only=True)
|
|
if isinstance(sd, dict):
|
|
if isinstance(sd.get("state_dict"), dict):
|
|
sd = sd["state_dict"]
|
|
elif isinstance(sd.get("decoder"), dict):
|
|
sd = sd["decoder"]
|
|
if isinstance(sd, dict) and all(k.startswith("model.") for k in sd):
|
|
sd = {k[len("model.") :]: v for k, v in sd.items()}
|
|
return sd
|
|
|
|
|
|
def conv(n_in, n_out, **kwargs):
|
|
return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
|
|
|
|
|
|
class Clamp(nn.Module):
|
|
def forward(self, x):
|
|
return torch.tanh(x / 3) * 3
|
|
|
|
|
|
class MemBlock(nn.Module):
|
|
def __init__(self, n_in, n_out):
|
|
super().__init__()
|
|
self.conv = nn.Sequential(
|
|
conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)
|
|
)
|
|
self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
|
|
self.act = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x, past):
|
|
return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
|
|
|
|
|
|
class SuperMemBlock(nn.Module):
|
|
"""MemBlock variant used by the Super decoder (ConvNeXt-style: 7x7 depthwise conv + inverted bottleneck)."""
|
|
|
|
def __init__(self, n_f):
|
|
super().__init__()
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(n_f * 2, n_f * 2, 7, padding=3, groups=n_f * 2, bias=False),
|
|
nn.Conv2d(n_f * 2, n_f * 4, 1),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv2d(n_f * 4, n_f, 1, bias=False),
|
|
)
|
|
|
|
def forward(self, x, past):
|
|
return self.conv(torch.cat([x, past], 1)) + x
|
|
|
|
|
|
class TPool(nn.Module):
|
|
def __init__(self, n_f, stride):
|
|
super().__init__()
|
|
self.stride = stride
|
|
self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False)
|
|
|
|
def forward(self, x):
|
|
_NT, C, H, W = x.shape
|
|
return self.conv(x.reshape(-1, self.stride * C, H, W))
|
|
|
|
|
|
class TGrow(nn.Module):
|
|
def __init__(self, n_f, stride):
|
|
super().__init__()
|
|
self.stride = stride
|
|
self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False)
|
|
|
|
def forward(self, x):
|
|
_NT, C, H, W = x.shape
|
|
x = self.conv(x)
|
|
return x.reshape(-1, C, H, W)
|
|
|
|
|
|
def apply_model_with_memblocks_parallel(model, x, show_progress_bar):
|
|
"""Apply a sequential model with memblocks over the time axis in parallel.
|
|
|
|
Args:
|
|
- model: nn.Sequential of blocks to apply
|
|
- x: input data, of dimensions NTCHW
|
|
- show_progress_bar: if True, enables tqdm progressbar display
|
|
|
|
Returns NTCHW tensor of output data.
|
|
"""
|
|
assert x.ndim == 5, f"TAE operates on NTCHW tensors, but got {x.ndim}-dim tensor"
|
|
N, T, C, H, W = x.shape
|
|
x = x.reshape(N * T, C, H, W)
|
|
|
|
# parallel over input timesteps, iterate over blocks
|
|
for b in tqdm(model, disable=not show_progress_bar):
|
|
if isinstance(b, (MemBlock, SuperMemBlock)):
|
|
NT, C, H, W = x.shape
|
|
T = NT // N
|
|
_x = x.reshape(N, T, C, H, W)
|
|
# pad with zeros along time axis (i.e. empty memory), slice
|
|
block_memory = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :T].reshape(x.shape)
|
|
x = b(x, block_memory)
|
|
else:
|
|
x = b(x)
|
|
NT, C, H, W = x.shape
|
|
T = NT // N
|
|
return x.view(N, T, C, H, W)
|
|
|
|
|
|
def apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar=None):
|
|
"""Process the work queue (a graph traversal over blocks and timesteps)
|
|
until an output frame is produced or the queue is empty.
|
|
Mutates memory and work_queue in place.
|
|
|
|
Returns N1CHW output tensor, or None if the queue needs more input.
|
|
"""
|
|
while work_queue:
|
|
xt, i = work_queue.pop(0)
|
|
if progress_bar is not None and i == 0:
|
|
progress_bar.update(1)
|
|
if i == len(model):
|
|
return xt.unsqueeze(1)
|
|
b = model[i]
|
|
if isinstance(b, (MemBlock, SuperMemBlock)):
|
|
# mem blocks are simple since we're visiting the graph in causal order
|
|
if memory[i] is None:
|
|
xt_new = b(xt, xt * 0)
|
|
else:
|
|
xt_new = b(xt, memory[i])
|
|
memory[i] = xt
|
|
work_queue.insert(0, TWorkItem(xt_new, i + 1))
|
|
elif isinstance(b, TPool):
|
|
# pool blocks accumulate inputs until they have enough to pool
|
|
if memory[i] is None:
|
|
memory[i] = []
|
|
memory[i].append(xt)
|
|
if len(memory[i]) > b.stride:
|
|
raise ValueError(f"TPool memory overflow: {len(memory[i])} items for stride {b.stride}")
|
|
elif len(memory[i]) == b.stride:
|
|
N, C, H, W = xt.shape
|
|
xt = b(torch.cat(memory[i], 1).view(N * b.stride, C, H, W))
|
|
memory[i] = []
|
|
work_queue.insert(0, TWorkItem(xt, i + 1))
|
|
elif isinstance(b, TGrow):
|
|
xt = b(xt)
|
|
NT, C, H, W = xt.shape
|
|
for xt_next in reversed(xt.view(NT // b.stride, b.stride * C, H, W).chunk(b.stride, 1)):
|
|
work_queue.insert(0, TWorkItem(xt_next, i + 1))
|
|
else:
|
|
xt = b(xt)
|
|
work_queue.insert(0, TWorkItem(xt, i + 1))
|
|
return None
|
|
|
|
|
|
def apply_model_with_memblocks_sequential(model, x, show_progress_bar):
|
|
"""Apply a sequential model with memblocks, iterating over timesteps as
|
|
well as blocks (slow but uses O(1) memory w.r.t. sequence length).
|
|
|
|
Args:
|
|
- model: nn.Sequential of blocks to apply
|
|
- x: input data, of dimensions NTCHW
|
|
- show_progress_bar: if True, enables tqdm progressbar display
|
|
|
|
Returns NTCHW tensor of output data.
|
|
"""
|
|
assert x.ndim == 5, f"TAE operates on NTCHW tensors, but got {x.ndim}-dim tensor"
|
|
work_queue = [TWorkItem(xt, 0) for xt in x.unbind(1)]
|
|
memory = [None] * len(model)
|
|
progress_bar = tqdm(range(len(work_queue)), disable=not show_progress_bar)
|
|
out = []
|
|
while work_queue:
|
|
xt = apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar)
|
|
if xt is not None:
|
|
out.append(xt)
|
|
progress_bar.close()
|
|
return torch.cat(out, 1)
|
|
|
|
|
|
def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
|
|
"""Apply a sequential model with memblocks to the given input.
|
|
|
|
Args:
|
|
- model: nn.Sequential of blocks to apply
|
|
- x: input data, of dimensions NTCHW
|
|
- parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
|
|
if False, each timestep will be processed sequentially (slow but uses O(1) memory)
|
|
- show_progress_bar: if True, enables tqdm progressbar display
|
|
|
|
Returns NTCHW tensor of output data.
|
|
"""
|
|
if parallel:
|
|
return apply_model_with_memblocks_parallel(model, x, show_progress_bar)
|
|
else:
|
|
return apply_model_with_memblocks_sequential(model, x, show_progress_bar)
|
|
|
|
|
|
class TAEWan(nn.Module):
|
|
"""Tiny AutoEncoder core for the Wan family (Wan 2.1 16ch / Wan 2.2 48ch latents).
|
|
|
|
Operates on NTCHW tensors with pixel values in [0, 1] and unnormalized
|
|
(~Gaussian) latents, matching the reference taehv implementation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
encoder_time_downscale=(True, True, False),
|
|
decoder_time_upscale=(False, True, True),
|
|
decoder_space_upscale=(True, True, True),
|
|
patch_size=1,
|
|
latent_channels=16,
|
|
arch_variant=None,
|
|
):
|
|
"""
|
|
Args:
|
|
encoder_time_downscale: whether temporal downsampling is enabled for each block.
|
|
decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
|
|
decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
|
|
patch_size: input/output pixelshuffle patch-size for this model.
|
|
latent_channels: number of latent channels (z dim) for this model.
|
|
arch_variant: decoder architecture variant. None (base) or "super" (higher-quality, ~2x decoder params).
|
|
"""
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.latent_channels = latent_channels
|
|
self.image_channels = 3
|
|
if len(decoder_time_upscale) == 2:
|
|
decoder_time_upscale = (False, *decoder_time_upscale)
|
|
assert arch_variant in (None, "super"), f"unrecognized arch_variant {arch_variant!r}"
|
|
self.encoder = nn.Sequential(
|
|
conv(self.image_channels * self.patch_size ** 2, 64),
|
|
nn.ReLU(inplace=True),
|
|
TPool(64, 2 if encoder_time_downscale[0] else 1),
|
|
conv(64, 64, stride=2, bias=False),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
TPool(64, 2 if encoder_time_downscale[1] else 1),
|
|
conv(64, 64, stride=2, bias=False),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
TPool(64, 2 if encoder_time_downscale[2] else 1),
|
|
conv(64, 64, stride=2, bias=False),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
MemBlock(64, 64),
|
|
conv(64, self.latent_channels),
|
|
)
|
|
if arch_variant == "super":
|
|
n_f = [512, 256, 128, 64]
|
|
self.decoder = nn.Sequential(
|
|
nn.Conv2d(self.latent_channels, n_f[0], 1, bias=False),
|
|
SuperMemBlock(n_f[0]),
|
|
SuperMemBlock(n_f[0]),
|
|
SuperMemBlock(n_f[0]),
|
|
conv(n_f[0], n_f[1] * (2 if decoder_space_upscale[0] else 1) ** 2),
|
|
nn.ReLU(inplace=True),
|
|
nn.PixelShuffle(2 if decoder_space_upscale[0] else 1),
|
|
TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
|
|
SuperMemBlock(n_f[1]),
|
|
SuperMemBlock(n_f[1]),
|
|
SuperMemBlock(n_f[1]),
|
|
conv(n_f[1], n_f[2] * (2 if decoder_space_upscale[1] else 1) ** 2),
|
|
nn.ReLU(inplace=True),
|
|
nn.PixelShuffle(2 if decoder_space_upscale[1] else 1),
|
|
TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
|
|
SuperMemBlock(n_f[2]),
|
|
SuperMemBlock(n_f[2]),
|
|
SuperMemBlock(n_f[2]),
|
|
conv(n_f[2], n_f[3] * (2 if decoder_space_upscale[2] else 1) ** 2),
|
|
nn.ReLU(inplace=True),
|
|
nn.PixelShuffle(2 if decoder_space_upscale[2] else 1),
|
|
TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1),
|
|
conv(n_f[3], self.image_channels * self.patch_size ** 2),
|
|
)
|
|
else:
|
|
n_f = [256, 128, 64, 64]
|
|
self.decoder = nn.Sequential(
|
|
Clamp(),
|
|
conv(self.latent_channels, n_f[0]),
|
|
nn.ReLU(inplace=True),
|
|
MemBlock(n_f[0], n_f[0]),
|
|
MemBlock(n_f[0], n_f[0]),
|
|
MemBlock(n_f[0], n_f[0]),
|
|
nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
|
|
TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1),
|
|
conv(n_f[0], n_f[1], bias=False),
|
|
MemBlock(n_f[1], n_f[1]),
|
|
MemBlock(n_f[1], n_f[1]),
|
|
MemBlock(n_f[1], n_f[1]),
|
|
nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
|
|
TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1),
|
|
conv(n_f[1], n_f[2], bias=False),
|
|
MemBlock(n_f[2], n_f[2]),
|
|
MemBlock(n_f[2], n_f[2]),
|
|
MemBlock(n_f[2], n_f[2]),
|
|
nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
|
|
TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1),
|
|
conv(n_f[2], n_f[3], bias=False),
|
|
nn.ReLU(inplace=True),
|
|
conv(n_f[3], self.image_channels * self.patch_size ** 2),
|
|
)
|
|
# computed properties
|
|
self.t_downscale = 2 ** sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool))
|
|
self.t_upscale = 2 ** sum(t.stride == 2 for t in self.decoder if isinstance(t, TGrow))
|
|
self.frames_to_trim = self.t_upscale - 1
|
|
|
|
def patch_tgrow_layers(self, sd):
|
|
"""Patch TGrow layers to use a smaller kernel if temporal upscaling
|
|
was partially disabled relative to the checkpoint.
|
|
|
|
Args:
|
|
sd: state dict to patch
|
|
"""
|
|
new_sd = self.state_dict()
|
|
for i, layer in enumerate(self.decoder):
|
|
if isinstance(layer, TGrow):
|
|
key = f"decoder.{i}.conv.weight"
|
|
if key in sd and sd[key].shape[0] > new_sd[key].shape[0]:
|
|
# take the last-timestep output channels
|
|
sd[key] = sd[key][-new_sd[key].shape[0] :]
|
|
return sd
|
|
|
|
def preprocess_input_frames(self, x):
|
|
"""Preprocess RGB input frames prior to the main encoder sequence."""
|
|
if self.patch_size > 1:
|
|
x = F.pixel_unshuffle(x, self.patch_size)
|
|
return x
|
|
|
|
def postprocess_output_frames(self, x):
|
|
"""Postprocess RGB frames after the main decoder sequence."""
|
|
if self.patch_size > 1:
|
|
x = F.pixel_shuffle(x, self.patch_size)
|
|
return x.clamp_(0, 1)
|
|
|
|
def encode_video(self, x, parallel=True, show_progress_bar=False):
|
|
"""Encode a sequence of frames.
|
|
|
|
Args:
|
|
x: input NTCHW RGB (C=3) tensor with values in [0, 1].
|
|
parallel: if True, all frames will be processed at once.
|
|
(this is faster but may require more memory).
|
|
if False, frames will be processed sequentially.
|
|
Returns NTCHW latent tensor with ~Gaussian values.
|
|
"""
|
|
if isinstance(self.encoder, nn.Identity):
|
|
raise RuntimeError("仅 Decoder 的 TAE 权重不支持 encode；请加载包含 Encoder 的完整权重。")
|
|
x = self.preprocess_input_frames(x)
|
|
if x.shape[1] % self.t_downscale != 0:
|
|
# pad at end to multiple of self.t_downscale
|
|
n_pad = self.t_downscale - x.shape[1] % self.t_downscale
|
|
padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
|
|
x = torch.cat([x, padding], 1)
|
|
return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
|
|
|
|
def decode_video(self, x, parallel=True, show_progress_bar=False):
|
|
"""Decode a sequence of frames.
|
|
|
|
Args:
|
|
x: input NTCHW latent (C=self.latent_channels) tensor with ~Gaussian values.
|
|
parallel: if True, all frames will be processed at once.
|
|
(this is faster but may require more memory).
|
|
if False, frames will be processed sequentially.
|
|
Returns NTCHW RGB tensor with ~[0, 1] values.
|
|
"""
|
|
x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
|
|
x = self.postprocess_output_frames(x)
|
|
return x[:, self.frames_to_trim :]
|
|
|
|
|
|
class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
|
"""VideoX-Fun-compatible wrapper around TAEWan.
|
|
|
|
Implements the same encode/decode protocol as AutoencoderKLWan /
|
|
AutoencoderKLWan3_8: encode/decode take BCTHW tensors in [-1, 1] and
|
|
latents are normalized with the per-channel mean/std of the corresponding
|
|
full-size Wan VAE, so TAE latents are interchangeable with the real ones.
|
|
|
|
Weight files (from https://github.com/madebyollin/taehv):
|
|
- taew2_1*.pth / .safetensors: Wan 2.1 & Wan 2.2 14B (16ch latent)
|
|
- taew2_2*.pth / .safetensors: Wan 2.2 5B (48ch latent, patch_size=2)
|
|
"""
|
|
|
|
# The TAE has no checkpointed block path, but the flag is declared so that
|
|
# `enable_gradient_checkpointing()` (a shared training-script call) is a
|
|
# no-op instead of raising, mirroring AutoencoderKLWan.
|
|
_supports_gradient_checkpointing = True
|
|
|
|
# per-channel statistics of the full-size Wan VAE latent spaces
|
|
WAN2_1_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,
|
|
]
|
|
WAN2_1_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,
|
|
]
|
|
WAN2_2_MEAN = [
|
|
-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,
|
|
]
|
|
WAN2_2_STD = [
|
|
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,
|
|
]
|
|
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
latent_channels=16,
|
|
patch_size=1,
|
|
arch_variant=None,
|
|
encoder_time_downscale=(True, True, False),
|
|
decoder_time_upscale=(False, True, True),
|
|
decoder_space_upscale=(True, True, True),
|
|
temporal_compression_ratio=4,
|
|
spatial_compression_ratio=None,
|
|
parallel=True,
|
|
decoder_only=False,
|
|
):
|
|
super().__init__()
|
|
if latent_channels == 16:
|
|
mean, std = self.WAN2_1_MEAN, self.WAN2_1_STD
|
|
elif latent_channels == 48:
|
|
mean, std = self.WAN2_2_MEAN, self.WAN2_2_STD
|
|
else:
|
|
raise ValueError(
|
|
f"AutoencoderTinyWan supports 16ch (taew2_1) or 48ch (taew2_2) latents, got {latent_channels}"
|
|
)
|
|
assert len(mean) == latent_channels and len(std) == latent_channels
|
|
self.mean = torch.tensor(mean, dtype=torch.float32)
|
|
self.std = torch.tensor(std, dtype=torch.float32)
|
|
self.scale = [self.mean, 1.0 / self.std]
|
|
|
|
# spatial compression is 8 from convs times the pixel-shuffle patch size
|
|
if spatial_compression_ratio is None:
|
|
spatial_compression_ratio = 8 * patch_size
|
|
self.spatial_compression_ratio = spatial_compression_ratio
|
|
# Persist the derived value into the frozen config as well: some pipelines
|
|
# read `vae.config.spatial_compression_ratio` (not the attribute), and the
|
|
# register_to_config decorator would otherwise keep the `None` default.
|
|
self.register_to_config(spatial_compression_ratio=spatial_compression_ratio)
|
|
|
|
self.model = TAEWan(
|
|
encoder_time_downscale=tuple(encoder_time_downscale),
|
|
decoder_time_upscale=tuple(decoder_time_upscale),
|
|
decoder_space_upscale=tuple(decoder_space_upscale),
|
|
patch_size=patch_size,
|
|
latent_channels=latent_channels,
|
|
arch_variant=arch_variant,
|
|
)
|
|
if decoder_only:
|
|
self.model.encoder = nn.Identity()
|
|
self.parallel = parallel
|
|
self.gradient_checkpointing = False
|
|
|
|
def _set_gradient_checkpointing(self, *args, **kwargs):
|
|
# TAE has no gradient checkpointing path; accept the flag for API parity.
|
|
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")
|
|
|
|
def _encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
# x: BCTHW in [-1, 1] -> normalized latents BCLHW.
|
|
# NOTE: the official TAE weights were distilled against the *normalized*
|
|
# Wan latent space (see the official diffusers wrapper, which sets
|
|
# latents_mean=zeros/latents_std=ones and passes pipeline latents
|
|
# through unchanged), so no mean/std scaling is applied here: the
|
|
# encoder output is already interchangeable with full-VAE latents.
|
|
x = (x.add(1).div_(2)).clamp_(0, 1) # [-1, 1] -> [0, 1]
|
|
x = x.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
|
|
zs = []
|
|
for u in x:
|
|
z = self.model.encode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0)
|
|
zs.append(z.permute(1, 0, 2, 3)) # TCHW -> CTHW
|
|
return torch.stack(zs)
|
|
|
|
@apply_forward_hook
|
|
def encode(
|
|
self, x: torch.Tensor, return_dict: bool = True
|
|
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
|
|
h = self._encode(x)
|
|
# TAE encoder is deterministic; wrap the output in a degenerate
|
|
# gaussian (near-zero variance) for interface compatibility.
|
|
logvar = torch.full_like(h, DETERMINISTIC_LOGVAR)
|
|
posterior = DiagonalGaussianDistribution(torch.cat([h, logvar], dim=1))
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
def _decode(self, zs):
|
|
# zs: BCLHW normalized latents -> BCTHW in [-1, 1].
|
|
# NOTE: no denormalization here, for the same reason as in _encode:
|
|
# the official TAE decoder consumes normalized latents directly.
|
|
zs = zs.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
|
|
dec = []
|
|
for u in zs:
|
|
d = self.model.decode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0)
|
|
d = d.mul_(2).sub_(1).clamp_(-1, 1) # [0, 1] -> [-1, 1]
|
|
dec.append(d)
|
|
dec = torch.stack(dec).permute(0, 2, 1, 3, 4) # NTCHW -> BCTHW
|
|
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 forward(self, z: torch.Tensor, return_dict: bool = True):
|
|
# Use forward so decoder training runs through DDP hooks.
|
|
return self.decode(z, return_dict=return_dict)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_path, additional_kwargs=None, **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
|
|
|
|
pretrained_model_path = os.fspath(pretrained_model_path)
|
|
subfolder = kwargs.pop("subfolder", None)
|
|
if subfolder:
|
|
pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
|
|
model_kwargs = filter_kwargs(cls, additional_kwargs or {})
|
|
model_kwargs.update(filter_kwargs(cls, kwargs))
|
|
if os.path.isdir(pretrained_model_path):
|
|
if os.path.exists(os.path.join(pretrained_model_path, "config.json")):
|
|
loader_kwargs = {k: v for k, v in kwargs.items() if k not in model_kwargs}
|
|
return_loading_info = loader_kwargs.pop("output_loading_info", False)
|
|
model, info = super().from_pretrained(
|
|
pretrained_model_path, **model_kwargs, **loader_kwargs, output_loading_info=True
|
|
)
|
|
if any(info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
|
|
raise RuntimeError(f"TAE 配置与权重不匹配或权重不完整：{info}")
|
|
return (model, info) if return_loading_info else model
|
|
candidates = sorted(
|
|
f for f in os.listdir(pretrained_model_path) if f.endswith((".pt", ".pth", ".bin", ".safetensors"))
|
|
)
|
|
if len(candidates) != 1:
|
|
raise ValueError(f"请指定唯一的 TAE 权重文件，目录内找到：{candidates}")
|
|
pretrained_model_path = os.path.join(pretrained_model_path, candidates[0])
|
|
|
|
if pretrained_model_path.endswith(".safetensors"):
|
|
from safetensors.torch import load_file
|
|
|
|
state_dict = load_file(pretrained_model_path)
|
|
else:
|
|
state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True)
|
|
if isinstance(state_dict, dict):
|
|
if isinstance(state_dict.get("state_dict"), dict):
|
|
state_dict = state_dict["state_dict"]
|
|
elif isinstance(state_dict.get("decoder"), dict):
|
|
state_dict = state_dict["decoder"]
|
|
if (
|
|
not isinstance(state_dict, dict)
|
|
or not state_dict
|
|
or not all(isinstance(k, str) and isinstance(v, torch.Tensor) for k, v in state_dict.items())
|
|
):
|
|
raise ValueError("无效的 TAE 权重：需要 tensor state_dict 或 checkpoint['decoder']。")
|
|
if all(k.startswith("model.") for k in state_dict):
|
|
state_dict = {k[len("model.") :]: v for k, v in state_dict.items()}
|
|
|
|
# Borrow an encoder when the checkpoint is decoder-only. Some taew2_2_super
|
|
# releases ship only the Super decoder; without an encoder, encode() would
|
|
# raise. We merge encoder.* weights from a companion base checkpoint so the
|
|
# model becomes a full encoder(base)+decoder(super) TAE.
|
|
# Priority: explicit `encoder_path` (via additional_kwargs / vae_kwargs or
|
|
# kwargs) > auto sibling found by stripping "_super" from the filename.
|
|
encoder_path = None
|
|
if isinstance(additional_kwargs, dict):
|
|
encoder_path = additional_kwargs.get("encoder_path")
|
|
if encoder_path is None:
|
|
encoder_path = kwargs.get("encoder_path")
|
|
if not any(k.startswith("encoder.") for k in state_dict):
|
|
if encoder_path is None:
|
|
root, ext = os.path.splitext(pretrained_model_path)
|
|
sibling = root.replace("_super", "") + ext
|
|
if sibling != pretrained_model_path and os.path.exists(sibling):
|
|
encoder_path = sibling
|
|
if encoder_path is not None and os.path.exists(encoder_path):
|
|
enc_sd = {k: v for k, v in _load_raw_state_dict(encoder_path).items() if k.startswith("encoder.")}
|
|
if enc_sd:
|
|
state_dict = {**enc_sd, **state_dict}
|
|
print(f"[TAE] decoder-only 权重：已从 {encoder_path} 借用 {len(enc_sd)} 个 encoder 键以支持 encode。")
|
|
if not any(k.startswith("encoder.") for k in state_dict):
|
|
print("[TAE][warn] 未找到 encoder 权重，将以 decoder_only 加载，encode() 不可用。"
|
|
"可通过 vae_kwargs.encoder_path=<base taew2_2 权重路径> 指定，"
|
|
"或把 base 权重放在与 super 同名（去掉 _super）的路径下。")
|
|
|
|
if "decoder.0.weight" in state_dict and "decoder.1.weight" not in state_dict:
|
|
arch_variant, first_key, block_key = "super", "decoder.0.weight", "decoder.1.conv.0.weight"
|
|
first_channels, first_kernel, block_shape = 512, (1, 1), (1024, 1, 7, 7)
|
|
elif "decoder.1.weight" in state_dict and "decoder.0.weight" not in state_dict:
|
|
arch_variant, first_key, block_key = None, "decoder.1.weight", "decoder.3.conv.0.weight"
|
|
first_channels, first_kernel, block_shape = 256, (3, 3), (256, 512, 3, 3)
|
|
else:
|
|
raise ValueError("无法识别 TAE 架构：Decoder 首层权重缺失或存在冲突。")
|
|
first = state_dict[first_key]
|
|
block = state_dict.get(block_key)
|
|
last = state_dict.get("decoder.22.weight")
|
|
if (
|
|
first.ndim != 4
|
|
or first.shape[0] != first_channels
|
|
or tuple(first.shape[2:]) != first_kernel
|
|
or block is None
|
|
or tuple(block.shape) != block_shape
|
|
or last is None
|
|
or last.ndim != 4
|
|
):
|
|
raise ValueError("无法识别 TAE 架构：首层、记忆块或输出层的形状不符合普通/Super TAE。")
|
|
latent_channels = first.shape[1]
|
|
patch_size = {16: 1, 48: 2}.get(latent_channels)
|
|
if patch_size is None or tuple(last.shape) != (3 * patch_size ** 2, 64, 3, 3):
|
|
raise ValueError("TAE latent 通道或输出 patch_size 不符合 Wan 2.1/2.2 架构。")
|
|
inferred = {
|
|
"arch_variant": arch_variant,
|
|
"latent_channels": latent_channels,
|
|
"patch_size": patch_size,
|
|
"decoder_only": not any(k.startswith("encoder.") for k in state_dict),
|
|
}
|
|
for key, value in inferred.items():
|
|
if key in model_kwargs and model_kwargs[key] != value:
|
|
raise ValueError(f"TAE 参数 {key}={model_kwargs[key]!r} 与权重识别值 {value!r} 冲突。")
|
|
model_kwargs[key] = value
|
|
model = cls(**model_kwargs)
|
|
state_dict = model.model.patch_tgrow_layers(dict(state_dict))
|
|
model.model.load_state_dict(state_dict, strict=True)
|
|
torch_dtype = kwargs.get("torch_dtype")
|
|
if torch_dtype == "auto":
|
|
torch_dtype = first.dtype
|
|
if torch_dtype is not None:
|
|
model.to(dtype=torch_dtype)
|
|
return model
|