wehub-resource-sync
424 lines
15 KiB
Python
424 lines
15 KiB
Python
import logging
|
|
import os
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from wan_5b.modules.vae2_2 import (
|
|
CausalConv3d,
|
|
Decoder3d,
|
|
Encoder3d,
|
|
count_conv3d,
|
|
patchify,
|
|
unpatchify,
|
|
)
|
|
|
|
|
|
def _extract_checkpoint_state_dict(raw):
|
|
state = raw
|
|
if isinstance(state, dict) and "state_dict" in state:
|
|
state = state["state_dict"]
|
|
if isinstance(state, dict) and "gen_model" in state:
|
|
state = state["gen_model"]
|
|
if isinstance(state, dict) and "generator" in state:
|
|
state = state["generator"]
|
|
if not isinstance(state, dict):
|
|
raise ValueError("Unsupported checkpoint format: expected a dict-like state_dict.")
|
|
return state
|
|
|
|
|
|
def _map_lightvae_key_to_wanvae(key):
|
|
def _map_resnet_tail(tail):
|
|
if tail.startswith("norm1."):
|
|
return "residual.0." + tail[len("norm1."):]
|
|
if tail.startswith("conv1."):
|
|
return "residual.2." + tail[len("conv1."):]
|
|
if tail.startswith("norm2."):
|
|
return "residual.3." + tail[len("norm2."):]
|
|
if tail.startswith("conv2."):
|
|
return "residual.6." + tail[len("conv2."):]
|
|
if tail.startswith("conv_shortcut."):
|
|
return "shortcut." + tail[len("conv_shortcut."):]
|
|
return tail
|
|
|
|
if key.startswith("dynamic_feature_projection_heads."):
|
|
return None
|
|
|
|
if key.startswith("quant_conv."):
|
|
return key.replace("quant_conv.", "conv1.", 1)
|
|
if key.startswith("post_quant_conv."):
|
|
return key.replace("post_quant_conv.", "conv2.", 1)
|
|
|
|
if key.startswith("encoder.conv_in."):
|
|
return key.replace("encoder.conv_in.", "encoder.conv1.", 1)
|
|
if key.startswith("encoder.mid_block.resnets.0."):
|
|
tail = key[len("encoder.mid_block.resnets.0."):]
|
|
return "encoder.middle.0." + _map_resnet_tail(tail)
|
|
if key.startswith("encoder.mid_block.attentions.0."):
|
|
return key.replace("encoder.mid_block.attentions.0.", "encoder.middle.1.", 1)
|
|
if key.startswith("encoder.mid_block.resnets.1."):
|
|
tail = key[len("encoder.mid_block.resnets.1."):]
|
|
return "encoder.middle.2." + _map_resnet_tail(tail)
|
|
if key.startswith("encoder.norm_out."):
|
|
return key.replace("encoder.norm_out.", "encoder.head.0.", 1)
|
|
if key.startswith("encoder.conv_out."):
|
|
return key.replace("encoder.conv_out.", "encoder.head.2.", 1)
|
|
|
|
if key.startswith("encoder.down_blocks."):
|
|
parts = key.split(".")
|
|
if len(parts) >= 6 and parts[3] == "resnets":
|
|
tail = ".".join(parts[5:])
|
|
return f"encoder.downsamples.{parts[2]}.downsamples.{parts[4]}." + _map_resnet_tail(tail)
|
|
if len(parts) >= 7 and parts[3] == "downsampler" and parts[4] == "resample":
|
|
return f"encoder.downsamples.{parts[2]}.downsamples.2.resample.{parts[5]}." + ".".join(parts[6:])
|
|
if len(parts) >= 6 and parts[3] == "downsampler" and parts[4] == "time_conv":
|
|
return f"encoder.downsamples.{parts[2]}.downsamples.2.time_conv." + ".".join(parts[5:])
|
|
|
|
if key.startswith("decoder.conv_in."):
|
|
return key.replace("decoder.conv_in.", "decoder.conv1.", 1)
|
|
if key.startswith("decoder.mid_block.resnets.0."):
|
|
tail = key[len("decoder.mid_block.resnets.0."):]
|
|
return "decoder.middle.0." + _map_resnet_tail(tail)
|
|
if key.startswith("decoder.mid_block.attentions.0."):
|
|
return key.replace("decoder.mid_block.attentions.0.", "decoder.middle.1.", 1)
|
|
if key.startswith("decoder.mid_block.resnets.1."):
|
|
tail = key[len("decoder.mid_block.resnets.1."):]
|
|
return "decoder.middle.2." + _map_resnet_tail(tail)
|
|
if key.startswith("decoder.norm_out."):
|
|
return key.replace("decoder.norm_out.", "decoder.head.0.", 1)
|
|
if key.startswith("decoder.conv_out."):
|
|
return key.replace("decoder.conv_out.", "decoder.head.2.", 1)
|
|
|
|
if key.startswith("decoder.up_blocks."):
|
|
parts = key.split(".")
|
|
if len(parts) >= 6 and parts[3] == "resnets":
|
|
tail = ".".join(parts[5:])
|
|
return f"decoder.upsamples.{parts[2]}.upsamples.{parts[4]}." + _map_resnet_tail(tail)
|
|
if len(parts) >= 7 and parts[3] == "upsampler" and parts[4] == "resample":
|
|
return f"decoder.upsamples.{parts[2]}.upsamples.3.resample.{parts[5]}." + ".".join(parts[6:])
|
|
if len(parts) >= 6 and parts[3] == "upsampler" and parts[4] == "time_conv":
|
|
return f"decoder.upsamples.{parts[2]}.upsamples.3.time_conv." + ".".join(parts[5:])
|
|
|
|
return key
|
|
|
|
|
|
def _normalize_vae_state_dict(raw_state):
|
|
state = _extract_checkpoint_state_dict(raw_state)
|
|
normalized = {}
|
|
for key, value in state.items():
|
|
mapped_key = _map_lightvae_key_to_wanvae(key)
|
|
if mapped_key is None:
|
|
continue
|
|
normalized[mapped_key] = value
|
|
return normalized
|
|
|
|
|
|
def infer_lightvae_pruning_rate_from_ckpt(vae_path, full_decoder_conv1_out=1024):
|
|
if vae_path is None or not os.path.exists(vae_path):
|
|
return None
|
|
try:
|
|
raw_state = torch.load(vae_path, map_location="cpu")
|
|
state = _extract_checkpoint_state_dict(raw_state)
|
|
except Exception as exc:
|
|
logging.warning("Failed to load checkpoint for pruning-rate inference: %s", exc)
|
|
return None
|
|
|
|
weight = None
|
|
if isinstance(state, dict):
|
|
if "decoder.conv_in.weight" in state:
|
|
weight = state["decoder.conv_in.weight"]
|
|
elif "decoder.conv1.weight" in state:
|
|
weight = state["decoder.conv1.weight"]
|
|
|
|
if weight is None:
|
|
try:
|
|
normalized_state = _normalize_vae_state_dict(state)
|
|
weight = normalized_state.get("decoder.conv1.weight", None)
|
|
except Exception:
|
|
weight = None
|
|
|
|
if weight is None or not hasattr(weight, "shape") or len(weight.shape) < 1:
|
|
return None
|
|
|
|
student_out = int(weight.shape[0])
|
|
if full_decoder_conv1_out <= 0:
|
|
return None
|
|
|
|
pruning_rate = 1.0 - (float(student_out) / float(full_decoder_conv1_out))
|
|
pruning_rate = max(0.0, min(0.99, pruning_rate))
|
|
return round(pruning_rate, 6)
|
|
|
|
|
|
def convert_to_channels_last_3d(module):
|
|
for child in module.children():
|
|
if isinstance(child, nn.Conv3d):
|
|
child.weight.data = child.weight.data.to(memory_format=torch.channels_last_3d)
|
|
else:
|
|
convert_to_channels_last_3d(child)
|
|
|
|
|
|
class PrunableWanVAE(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim=160,
|
|
dec_dim=256,
|
|
z_dim=48,
|
|
dim_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
attn_scales=[],
|
|
temperal_downsample=[False, True, True],
|
|
dropout=0.0,
|
|
pruning_rate=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]
|
|
|
|
dim = max(1, int(round(dim * (1.0 - pruning_rate))))
|
|
dec_dim = max(1, int(round(dec_dim * (1.0 - pruning_rate))))
|
|
|
|
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,
|
|
)
|
|
|
|
def encode(self, x, scale):
|
|
self.clear_cache()
|
|
x = patchify(x, patch_size=2)
|
|
total_steps = 1 + (x.shape[2] - 1) // 4
|
|
for step in range(total_steps):
|
|
self._enc_conv_idx = [0]
|
|
if step == 0:
|
|
out = self.encoder(
|
|
x[:, :, :1, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx,
|
|
)
|
|
else:
|
|
out_chunk = self.encoder(
|
|
x[:, :, 1 + 4 * (step - 1):1 + 4 * step, :, :],
|
|
feat_cache=self._enc_feat_map,
|
|
feat_idx=self._enc_conv_idx,
|
|
)
|
|
out = torch.cat([out, out_chunk], 2)
|
|
mu, _ = 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]
|
|
self.clear_cache()
|
|
return mu
|
|
|
|
def decode(self, z, scale):
|
|
self.clear_cache()
|
|
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]
|
|
total_steps = z.shape[2]
|
|
x = self.conv2(z)
|
|
for step in range(total_steps):
|
|
self._conv_idx = [0]
|
|
if step == 0:
|
|
out = self.decoder(
|
|
x[:, :, step:step + 1, :, :],
|
|
feat_cache=self._feat_map,
|
|
feat_idx=self._conv_idx,
|
|
first_chunk=True,
|
|
)
|
|
else:
|
|
out_chunk = self.decoder(
|
|
x[:, :, step:step + 1, :, :],
|
|
feat_cache=self._feat_map,
|
|
feat_idx=self._conv_idx,
|
|
)
|
|
out = torch.cat([out, out_chunk], 2)
|
|
out = unpatchify(out, patch_size=2)
|
|
self.clear_cache()
|
|
return out
|
|
|
|
def cached_decode(self, z, 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]
|
|
total_steps = z.shape[2]
|
|
x = self.conv2(z)
|
|
is_first = self._feat_map[0] is None
|
|
for step in range(total_steps):
|
|
self._conv_idx = [0]
|
|
if step == 0 and is_first:
|
|
out = self.decoder(
|
|
x[:, :, step:step + 1, :, :],
|
|
feat_cache=self._feat_map,
|
|
feat_idx=self._conv_idx,
|
|
first_chunk=True,
|
|
)
|
|
elif step == 0:
|
|
out = self.decoder(
|
|
x[:, :, step:step + 1, :, :],
|
|
feat_cache=self._feat_map,
|
|
feat_idx=self._conv_idx,
|
|
)
|
|
else:
|
|
out_chunk = self.decoder(
|
|
x[:, :, step:step + 1, :, :],
|
|
feat_cache=self._feat_map,
|
|
feat_idx=self._conv_idx,
|
|
)
|
|
out = torch.cat([out, out_chunk], 2)
|
|
return unpatchify(out, patch_size=2)
|
|
|
|
def clear_cache(self):
|
|
self._conv_num = count_conv3d(self.decoder)
|
|
self._conv_idx = [0]
|
|
self._feat_map = [None] * self._conv_num
|
|
self._enc_conv_num = count_conv3d(self.encoder)
|
|
self._enc_conv_idx = [0]
|
|
self._enc_feat_map = [None] * self._enc_conv_num
|
|
|
|
|
|
def _load_lightvae_model(pretrained_path=None, z_dim=48, dim=160, device="cpu", **kwargs):
|
|
cfg = dict(
|
|
dim=dim,
|
|
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)
|
|
|
|
with torch.device("meta"):
|
|
model = PrunableWanVAE(**cfg)
|
|
|
|
if pretrained_path is None or not os.path.exists(pretrained_path):
|
|
raise FileNotFoundError(f"VAE checkpoint not found at {pretrained_path}")
|
|
|
|
logging.info("loading %s", pretrained_path)
|
|
raw_state = torch.load(pretrained_path, map_location="cpu")
|
|
state_dict = _normalize_vae_state_dict(raw_state)
|
|
missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
|
|
logging.info(
|
|
"LightVAE checkpoint loaded with strict=False (missing=%d, unexpected=%d)",
|
|
len(missing),
|
|
len(unexpected),
|
|
)
|
|
|
|
convert_to_channels_last_3d(model)
|
|
return model
|
|
|
|
|
|
class LightVAE5BWrapper(nn.Module):
|
|
def __init__(
|
|
self,
|
|
vae_path: str,
|
|
pruning_rate: Optional[float] = None,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
device: Optional[torch.device] = None,
|
|
):
|
|
super().__init__()
|
|
if device is None:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
if pruning_rate is None:
|
|
pruning_rate = infer_lightvae_pruning_rate_from_ckpt(vae_path)
|
|
if pruning_rate is None:
|
|
pruning_rate = 0.75
|
|
logging.warning(
|
|
"Unable to infer LightVAE pruning rate from checkpoint; fallback to 0.75."
|
|
)
|
|
|
|
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,
|
|
]
|
|
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,
|
|
]
|
|
self.mean = torch.tensor(mean, dtype=torch.float32)
|
|
self.std = torch.tensor(std, dtype=torch.float32)
|
|
self.vae_path = os.path.abspath(vae_path)
|
|
self.pruning_rate = pruning_rate
|
|
self.device = torch.device(device)
|
|
self.dtype = dtype
|
|
|
|
self.model = _load_lightvae_model(
|
|
pretrained_path=self.vae_path,
|
|
pruning_rate=self.pruning_rate,
|
|
).eval().requires_grad_(False)
|
|
self.to(device=self.device, dtype=self.dtype)
|
|
|
|
def to(self, device=None, dtype=None):
|
|
device = self.device if device is None else torch.device(device)
|
|
dtype = self.dtype if dtype is None else dtype
|
|
self.model.to(device=device, dtype=dtype)
|
|
self.mean = self.mean.to(device=device, dtype=dtype)
|
|
self.std = self.std.to(device=device, dtype=dtype)
|
|
self.device = device
|
|
self.dtype = dtype
|
|
return self
|
|
|
|
def eval(self):
|
|
super().eval()
|
|
self.model.eval()
|
|
return self
|
|
|
|
@torch.no_grad()
|
|
def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
|
|
zs = latent.permute(0, 2, 1, 3, 4)
|
|
if use_cache:
|
|
assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
|
|
|
|
scale = [self.mean, 1.0 / self.std]
|
|
decode_fn = self.model.cached_decode if use_cache else self.model.decode
|
|
|
|
output = []
|
|
for item in zs:
|
|
output.append(
|
|
decode_fn(item.unsqueeze(0).to(device=self.device, dtype=self.dtype), scale)
|
|
.float()
|
|
.clamp_(-1, 1)
|
|
.squeeze(0)
|
|
)
|
|
output = torch.stack(output, dim=0)
|
|
return output.permute(0, 2, 1, 3, 4)
|