Skip to content

Commit e711aaf

Browse files
authored
Lower VAE loading requirements:Create a new branch for GPU memory calculations in qwen-image vae (Comfy-Org#11199)
1 parent 57ddb7f commit e711aaf

1 file changed

Lines changed: 4 additions & 2 deletions

File tree

comfy/sd.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -549,8 +549,10 @@ def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None)
549549
ddconfig = {"dim": dim, "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}
550550
self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
551551
self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
552-
self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
553-
self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
552+
self.memory_used_encode = lambda shape, dtype: (1500 if shape[2]<=4 else 6000) * shape[3] * shape[4] * model_management.dtype_size(dtype)
553+
self.memory_used_decode = lambda shape, dtype: (2200 if shape[2]<=4 else 7000) * shape[3] * shape[4] * (8*8) * model_management.dtype_size(dtype)
554+
555+
554556
# Hunyuan 3d v2 2.0 & 2.1
555557
elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
556558

0 commit comments

Comments
 (0)