wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-15 08:40:50 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Support Wan 2.1 lightVAE

Browse Source
This commit is contained in:
kijai 2025-11-26 17:14:56 +02:00
parent 7ef46f66ff
commit 8e93a15857
2 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
comfy/ldm/wan/vae.py
View File

@ -231,7 +231,8 @@ class Encoder3d(nn.Module):
num_res_blocks=2,
attn_scales=[],
temperal_downsample=[True, True, False],
dropout=0.0):
dropout=0.0,
pruning_rate=0.0):
super().__init__()
self.dim = dim
self.z_dim = z_dim
@ -242,6 +243,7 @@ class Encoder3d(nn.Module):
# dimensions
dims = [dim * u for u in [1] + dim_mult]
dims = [int(d * (1 - pruning_rate)) for d in dims]
scale = 1.0
# init block
@ -335,7 +337,8 @@ class Decoder3d(nn.Module):
num_res_blocks=2,
attn_scales=[],
temperal_upsample=[False, True, True],
dropout=0.0):
dropout=0.0,
pruning_rate=0.0):
super().__init__()
self.dim = dim
self.z_dim = z_dim
@ -346,6 +349,7 @@ class Decoder3d(nn.Module):
# dimensions
dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
dims = [int(d * (1 - pruning_rate)) for d in dims]
scale = 1.0 / 2**(len(dim_mult) - 2)
# init block
@ -449,7 +453,8 @@ class WanVAE(nn.Module):
num_res_blocks=2,
attn_scales=[],
temperal_downsample=[True, True, False],
dropout=0.0):
dropout=0.0,
pruning_rate=0.0):
super().__init__()
self.dim = dim
self.z_dim = z_dim
@ -461,11 +466,11 @@ class WanVAE(nn.Module):
# modules
self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
attn_scales, self.temperal_downsample, dropout)
attn_scales, self.temperal_downsample, dropout, pruning_rate)
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)
attn_scales, self.temperal_upsample, dropout, pruning_rate)
def encode(self, x):
conv_idx = [0]

5
comfy/sd.py
View File

@ -528,13 +528,16 @@ class VAE:
self.memory_used_encode = lambda shape, dtype: 3300 * shape[3] * shape[4] * model_management.dtype_size(dtype)
self.memory_used_decode = lambda shape, dtype: 8000 * shape[3] * shape[4] * (16 * 16) * model_management.dtype_size(dtype)
else: # Wan 2.1 VAE
pruning_rate = 0.0
if sd["decoder.middle.0.residual.0.gamma"].shape[0] == 96: # lightx2v lightvae
pruning_rate = 0.75
self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
self.upscale_index_formula = (4, 8, 8)
self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
self.downscale_index_formula = (4, 8, 8)
self.latent_dim = 3
self.latent_channels = 16
ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0, "pruning_rate": pruning_rate}
self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)