mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-08-21 12:06:51 +08:00
719 lines
24 KiB
Python
719 lines
24 KiB
Python
|
# vae_mlx_final.py
|
||
|
|
import logging
|
||
|
|
from typing import Optional, List, Tuple
|
||
|
|
|
||
|
|
import mlx.core as mx
|
||
|
|
import mlx.nn as nn
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from mlx.utils import tree_unflatten
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
'WanVAE',
|
||
|
|
]
|
||
|
|
|
||
|
|
CACHE_T = 2
|
||
|
|
|
||
|
|
debug_line = 0
|
||
|
|
|
||
|
|
|
||
|
|
class CausalConv3d(nn.Conv3d):
|
||
|
|
"""
|
||
|
|
Causal 3d convolution for MLX.
|
||
|
|
Expects input in BTHWC format (batch, time, height, width, channels).
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
# Padding order: (W, W, H, H, T, 0)
|
||
|
|
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 __call__(self, x, cache_x=None):
|
||
|
|
padding = list(self._padding)
|
||
|
|
if cache_x is not None and self._padding[4] > 0:
|
||
|
|
x = mx.concatenate([cache_x, x], axis=1) # Concat along time axis
|
||
|
|
padding[4] -= cache_x.shape[1]
|
||
|
|
|
||
|
|
# Pad in BTHWC format
|
||
|
|
pad_width = [(0, 0), (padding[4], padding[5]), (padding[2], padding[3]),
|
||
|
|
(padding[0], padding[1]), (0, 0)]
|
||
|
|
x = mx.pad(x, pad_width)
|
||
|
|
|
||
|
|
result = super().__call__(x)
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
class RMS_norm(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, dim, channel_first=False, images=True, bias=False):
|
||
|
|
super().__init__()
|
||
|
|
self.channel_first = channel_first
|
||
|
|
self.images = images
|
||
|
|
self.scale = dim**0.5
|
||
|
|
|
||
|
|
# Just keep as 1D - let broadcasting do its magic
|
||
|
|
self.gamma = mx.ones((dim,))
|
||
|
|
self.bias = mx.zeros((dim,)) if bias else 0.
|
||
|
|
|
||
|
|
def __call__(self, x):
|
||
|
|
# F.normalize in PyTorch does L2 normalization, not RMS!
|
||
|
|
# For NHWC/BTHWC format, normalize along the last axis
|
||
|
|
# L2 norm: sqrt(sum(x^2))
|
||
|
|
norm = mx.sqrt(mx.sum(x * x, axis=-1, keepdims=True) + 1e-6)
|
||
|
|
x = x / norm
|
||
|
|
return x * self.scale * self.gamma + self.bias
|
||
|
|
|
||
|
|
|
||
|
|
class Upsample(nn.Module):
|
||
|
|
"""
|
||
|
|
Upsampling layer that matches PyTorch's behavior.
|
||
|
|
"""
|
||
|
|
def __init__(self, scale_factor, mode='nearest-exact'):
|
||
|
|
super().__init__()
|
||
|
|
self.scale_factor = scale_factor
|
||
|
|
self.mode = mode # mode is now unused, but kept for signature consistency
|
||
|
|
|
||
|
|
def __call__(self, x):
|
||
|
|
# For NHWC format (n, h, w, c)
|
||
|
|
|
||
|
|
# NOTE: For an integer scale_factor like 2.0, PyTorch's 'nearest-exact'
|
||
|
|
# is equivalent to a simple repeat operation. The previous coordinate-based
|
||
|
|
# sampling was not correct for this model and caused the divergence.
|
||
|
|
|
||
|
|
scale_h, scale_w = self.scale_factor
|
||
|
|
|
||
|
|
out = mx.repeat(x, int(scale_h), axis=1) # Repeat along H dimension
|
||
|
|
out = mx.repeat(out, int(scale_w), axis=2) # Repeat along W dimension
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
class AsymmetricPad(nn.Module):
|
||
|
|
"""A module to apply asymmetric padding, compatible with nn.Sequential."""
|
||
|
|
def __init__(self, pad_width: tuple):
|
||
|
|
super().__init__()
|
||
|
|
self.pad_width = pad_width
|
||
|
|
|
||
|
|
def __call__(self, x):
|
||
|
|
return mx.pad(x, self.pad_width)
|
||
|
|
|
||
|
|
# Update your Resample class to use 'nearest-exact'
|
||
|
|
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))
|
||
|
|
|
||
|
|
# --- CORRECTED PADDING LOGIC ---
|
||
|
|
elif mode == 'downsample2d':
|
||
|
|
# Replicate PyTorch's ZeroPad2d((0, 1, 0, 1)) + Conv2d(stride=2)
|
||
|
|
# Use the new AsymmetricPad module.
|
||
|
|
# Pad width for NHWC format is ((N), (H), (W), (C))
|
||
|
|
# Pad H with (top, bottom) and W with (left, right)
|
||
|
|
pad_layer = AsymmetricPad(pad_width=((0, 0), (0, 1), (0, 1), (0, 0)))
|
||
|
|
conv_layer = nn.Conv2d(dim, dim, 3, stride=(2, 2), padding=0)
|
||
|
|
self.resample = nn.Sequential(pad_layer, conv_layer)
|
||
|
|
|
||
|
|
elif mode == 'downsample3d':
|
||
|
|
# The spatial downsampling part uses the same logic
|
||
|
|
pad_layer = AsymmetricPad(pad_width=((0, 0), (0, 1), (0, 1), (0, 0)))
|
||
|
|
conv_layer = nn.Conv2d(dim, dim, 3, stride=(2, 2), padding=0)
|
||
|
|
self.resample = nn.Sequential(pad_layer, conv_layer)
|
||
|
|
|
||
|
|
self.time_conv = CausalConv3d(
|
||
|
|
dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
|
||
|
|
|
||
|
|
else:
|
||
|
|
self.resample = nn.Identity()
|
||
|
|
|
||
|
|
def __call__(self, x, feat_cache=None, feat_idx=[0]):
|
||
|
|
# The __call__ method logic remains unchanged from your original code
|
||
|
|
b, t, h, w, c = x.shape
|
||
|
|
|
||
|
|
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:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None and feat_cache[idx] != 'Rep':
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None and feat_cache[idx] == 'Rep':
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
mx.zeros_like(cache_x), cache_x
|
||
|
|
], axis=1)
|
||
|
|
|
||
|
|
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, t, h, w, 2, c)
|
||
|
|
x = mx.stack([x[:, :, :, :, 0, :], x[:, :, :, :, 1, :]], axis=2)
|
||
|
|
x = x.reshape(b, t * 2, h, w, c)
|
||
|
|
|
||
|
|
t = x.shape[1]
|
||
|
|
x = x.reshape(b * t, h, w, c)
|
||
|
|
|
||
|
|
x = self.resample(x)
|
||
|
|
|
||
|
|
_, h_new, w_new, c_new = x.shape
|
||
|
|
x = x.reshape(b, t, h_new, w_new, c_new)
|
||
|
|
|
||
|
|
if self.mode == 'downsample3d':
|
||
|
|
if feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
if feat_cache[idx] is None:
|
||
|
|
feat_cache[idx] = x
|
||
|
|
feat_idx[0] += 1
|
||
|
|
else:
|
||
|
|
cache_x = x[:, -1:, :, :, :]
|
||
|
|
x = self.time_conv(
|
||
|
|
mx.concatenate([feat_cache[idx][:, -1:, :, :, :], x], axis=1))
|
||
|
|
feat_cache[idx] = cache_x
|
||
|
|
feat_idx[0] += 1
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
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) if dropout > 0 else nn.Identity(),
|
||
|
|
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 __call__(self, x, feat_cache=None, feat_idx=[0]):
|
||
|
|
h = self.shortcut(x)
|
||
|
|
|
||
|
|
for i, layer in enumerate(self.residual.layers):
|
||
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
cache_x = x[:, -CACHE_T:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None:
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
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
|
||
|
|
self.proj.weight = mx.zeros_like(self.proj.weight)
|
||
|
|
|
||
|
|
def __call__(self, x):
|
||
|
|
# x is in BTHWC format
|
||
|
|
identity = x
|
||
|
|
b, t, h, w, c = x.shape
|
||
|
|
x = x.reshape(b * t, h, w, c) # Combine batch and time
|
||
|
|
x = self.norm(x)
|
||
|
|
# compute query, key, value
|
||
|
|
qkv = self.to_qkv(x) # Output: (b*t, h, w, 3*c)
|
||
|
|
qkv = qkv.reshape(b * t, h * w, 3 * c)
|
||
|
|
q, k, v = mx.split(qkv, 3, axis=-1)
|
||
|
|
|
||
|
|
# Reshape for attention
|
||
|
|
q = q.reshape(b * t, h * w, c)
|
||
|
|
k = k.reshape(b * t, h * w, c)
|
||
|
|
v = v.reshape(b * t, h * w, c)
|
||
|
|
|
||
|
|
# Scaled dot product attention
|
||
|
|
scale = 1.0 / mx.sqrt(mx.array(c, dtype=q.dtype))
|
||
|
|
scores = (q @ k.transpose(0, 2, 1)) * scale
|
||
|
|
weights = mx.softmax(scores, axis=-1)
|
||
|
|
x = weights @ v
|
||
|
|
x = x.reshape(b * t, h, w, c)
|
||
|
|
|
||
|
|
# output
|
||
|
|
x = self.proj(x)
|
||
|
|
x = x.reshape(b, t, h, w, c)
|
||
|
|
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(dims[-1], dims[-1], dropout),
|
||
|
|
AttentionBlock(dims[-1]),
|
||
|
|
ResidualBlock(dims[-1], dims[-1], dropout)
|
||
|
|
)
|
||
|
|
|
||
|
|
# output blocks
|
||
|
|
self.head = nn.Sequential(
|
||
|
|
RMS_norm(dims[-1], images=False),
|
||
|
|
nn.SiLU(),
|
||
|
|
CausalConv3d(dims[-1], z_dim, 3, padding=1)
|
||
|
|
)
|
||
|
|
|
||
|
|
def __call__(self, x, feat_cache=None, feat_idx=[0]):
|
||
|
|
if feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
cache_x = x[:, -CACHE_T:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None:
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
x = self.conv1(x, feat_cache[idx])
|
||
|
|
feat_cache[idx] = cache_x
|
||
|
|
feat_idx[0] += 1
|
||
|
|
else:
|
||
|
|
x = self.conv1(x)
|
||
|
|
|
||
|
|
## downsamples
|
||
|
|
for i, layer in enumerate(self.downsamples.layers):
|
||
|
|
if feat_cache is not None:
|
||
|
|
x = layer(x, feat_cache, feat_idx)
|
||
|
|
else:
|
||
|
|
x = layer(x)
|
||
|
|
|
||
|
|
## middle
|
||
|
|
for layer in self.middle.layers:
|
||
|
|
if isinstance(layer, ResidualBlock) and feat_cache is not None:
|
||
|
|
x = layer(x, feat_cache, feat_idx)
|
||
|
|
else:
|
||
|
|
x = layer(x)
|
||
|
|
|
||
|
|
## head
|
||
|
|
for i, layer in enumerate(self.head.layers):
|
||
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
cache_x = x[:, -CACHE_T:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None:
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
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(dims[-1], images=False),
|
||
|
|
nn.SiLU(),
|
||
|
|
CausalConv3d(dims[-1], 3, 3, padding=1)
|
||
|
|
)
|
||
|
|
|
||
|
|
def __call__(self, x, feat_cache=None, feat_idx=[0]):
|
||
|
|
## conv1
|
||
|
|
if feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
cache_x = x[:, -CACHE_T:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None:
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
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.layers:
|
||
|
|
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.layers:
|
||
|
|
if feat_cache is not None:
|
||
|
|
x = layer(x, feat_cache, feat_idx)
|
||
|
|
else:
|
||
|
|
x = layer(x)
|
||
|
|
|
||
|
|
## head
|
||
|
|
for i, layer in enumerate(self.head.layers):
|
||
|
|
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||
|
|
idx = feat_idx[0]
|
||
|
|
cache_x = x[:, -CACHE_T:, :, :, :]
|
||
|
|
if cache_x.shape[1] < 2 and feat_cache[idx] is not None:
|
||
|
|
cache_x = mx.concatenate([
|
||
|
|
feat_cache[idx][:, -1:, :, :, :], cache_x
|
||
|
|
], axis=1)
|
||
|
|
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 name, module in model.named_modules():
|
||
|
|
if isinstance(module, CausalConv3d):
|
||
|
|
count += 1
|
||
|
|
return count
|
||
|
|
|
||
|
|
|
||
|
|
class WanVAE_(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 encode(self, x, scale):
|
||
|
|
# x is in BTHWC format
|
||
|
|
self.clear_cache()
|
||
|
|
## cache
|
||
|
|
t = x.shape[1]
|
||
|
|
iter_ = 1 + (t - 1) // 4
|
||
|
|
## Split encode input x by time into 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 = mx.concatenate([out, out_], axis=1)
|
||
|
|
|
||
|
|
z = self.conv1(out)
|
||
|
|
mu, log_var = mx.split(z, 2, axis=-1) # Split along channel dimension
|
||
|
|
|
||
|
|
if isinstance(scale[0], mx.array):
|
||
|
|
# Reshape scale for broadcasting in BTHWC format
|
||
|
|
scale_mean = scale[0].reshape(1, 1, 1, 1, self.z_dim)
|
||
|
|
scale_std = scale[1].reshape(1, 1, 1, 1, self.z_dim)
|
||
|
|
mu = (mu - scale_mean) * scale_std
|
||
|
|
else:
|
||
|
|
mu = (mu - scale[0]) * scale[1]
|
||
|
|
self.clear_cache()
|
||
|
|
|
||
|
|
return mu, log_var
|
||
|
|
|
||
|
|
def decode(self, z, scale):
|
||
|
|
# z is in BTHWC format
|
||
|
|
self.clear_cache()
|
||
|
|
if isinstance(scale[0], mx.array):
|
||
|
|
scale_mean = scale[0].reshape(1, 1, 1, 1, self.z_dim)
|
||
|
|
scale_std = scale[1].reshape(1, 1, 1, 1, self.z_dim)
|
||
|
|
z = z / scale_std + scale_mean
|
||
|
|
else:
|
||
|
|
z = z / scale[1] + scale[0]
|
||
|
|
iter_ = z.shape[1]
|
||
|
|
x = self.conv2(z)
|
||
|
|
for i in range(iter_):
|
||
|
|
self._conv_idx = [0]
|
||
|
|
if i == 0:
|
||
|
|
out = self.decoder(
|
||
|
|
x[:, i:i + 1, :, :, :],
|
||
|
|
feat_cache=self._feat_map,
|
||
|
|
feat_idx=self._conv_idx)
|
||
|
|
else:
|
||
|
|
out_ = self.decoder(
|
||
|
|
x[:, i:i + 1, :, :, :],
|
||
|
|
feat_cache=self._feat_map,
|
||
|
|
feat_idx=self._conv_idx)
|
||
|
|
out = mx.concatenate([out, out_], axis=1)
|
||
|
|
self.clear_cache()
|
||
|
|
return out
|
||
|
|
|
||
|
|
def reparameterize(self, mu, log_var):
|
||
|
|
std = mx.exp(0.5 * log_var)
|
||
|
|
eps = mx.random.normal(std.shape)
|
||
|
|
return eps * std + mu
|
||
|
|
|
||
|
|
def __call__(self, x):
|
||
|
|
mu, log_var = self.encode(x, self.scale)
|
||
|
|
z = self.reparameterize(mu, log_var)
|
||
|
|
x_recon = self.decode(z, self.scale)
|
||
|
|
return x_recon, mu, log_var
|
||
|
|
|
||
|
|
def sample(self, imgs, deterministic=False):
|
||
|
|
mu, log_var = self.encode(imgs, self.scale)
|
||
|
|
if deterministic:
|
||
|
|
return mu
|
||
|
|
std = mx.exp(0.5 * mx.clip(log_var, -30.0, 20.0))
|
||
|
|
return mu + std * mx.random.normal(std.shape)
|
||
|
|
|
||
|
|
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(pretrained_path=None, 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 = WanVAE_(**cfg)
|
||
|
|
|
||
|
|
# load checkpoint
|
||
|
|
if pretrained_path:
|
||
|
|
logging.info(f'loading {pretrained_path}')
|
||
|
|
weights = mx.load(pretrained_path)
|
||
|
|
model.update(tree_unflatten(list(weights.items())))
|
||
|
|
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
class WanVAE:
|
||
|
|
|
||
|
|
def __init__(self,
|
||
|
|
z_dim=16,
|
||
|
|
vae_pth='cache/vae_step_411000.pth',
|
||
|
|
dtype=mx.float32):
|
||
|
|
self.dtype = dtype
|
||
|
|
|
||
|
|
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 = mx.array(mean, dtype=dtype)
|
||
|
|
self.std = mx.array(std, dtype=dtype)
|
||
|
|
self.scale = [self.mean, 1.0 / self.std]
|
||
|
|
|
||
|
|
# init model
|
||
|
|
self.model = _video_vae(
|
||
|
|
pretrained_path=vae_pth,
|
||
|
|
z_dim=z_dim,
|
||
|
|
)
|
||
|
|
|
||
|
|
def encode(self, videos):
|
||
|
|
"""
|
||
|
|
videos: A list of videos each with shape [C, T, H, W].
|
||
|
|
Returns: List of encoded videos in [C, T, H, W] format.
|
||
|
|
"""
|
||
|
|
encoded = []
|
||
|
|
for video in videos:
|
||
|
|
# Convert CTHW -> BTHWC
|
||
|
|
x = mx.expand_dims(video, axis=0) # Add batch dimension
|
||
|
|
x = x.transpose(0, 2, 3, 4, 1) # BCTHW -> BTHWC
|
||
|
|
|
||
|
|
# Encode
|
||
|
|
z = self.model.encode(x, self.scale)[0] # Get mu only
|
||
|
|
|
||
|
|
# Convert back BTHWC -> CTHW and remove batch dimension
|
||
|
|
z = z.transpose(0, 4, 1, 2, 3) # BTHWC -> BCTHW
|
||
|
|
z = z.squeeze(0) # Remove batch dimension -> CTHW
|
||
|
|
|
||
|
|
encoded.append(z.astype(mx.float32))
|
||
|
|
|
||
|
|
return encoded
|
||
|
|
|
||
|
|
def decode(self, zs):
|
||
|
|
"""
|
||
|
|
zs: A list of latent codes each with shape [C, T, H, W].
|
||
|
|
Returns: List of decoded videos in [C, T, H, W] format.
|
||
|
|
"""
|
||
|
|
decoded = []
|
||
|
|
for z in zs:
|
||
|
|
# Convert CTHW -> BTHWC
|
||
|
|
x = mx.expand_dims(z, axis=0) # Add batch dimension
|
||
|
|
x = x.transpose(0, 2, 3, 4, 1) # BCTHW -> BTHWC
|
||
|
|
|
||
|
|
# Decode
|
||
|
|
x = self.model.decode(x, self.scale)
|
||
|
|
|
||
|
|
# Convert back BTHWC -> CTHW and remove batch dimension
|
||
|
|
x = x.transpose(0, 4, 1, 2, 3) # BTHWC -> BCTHW
|
||
|
|
x = x.squeeze(0) # Remove batch dimension -> CTHW
|
||
|
|
|
||
|
|
# Clamp values
|
||
|
|
x = mx.clip(x, -1, 1)
|
||
|
|
|
||
|
|
decoded.append(x.astype(mx.float32))
|
||
|
|
|
||
|
|
return decoded
|