xinyun/ComfyUI-MochiWrapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use dtype for vea.

Browse Source 
Somehow, download node is not using dtype for vae.
This commit is contained in:
Yoshimasa Niwa 2024-11-05 13:46:45 +09:00
parent 99285ca1e7
commit 24a834be79
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
nodes.py
View File

@ -177,14 +177,15 @@ class DownloadAndLoadMochiModel:
nonlinearity="silu",
output_nonlinearity="silu",
causal=True,
dtype=dtype,
)
vae_sd = load_torch_file(vae_path)
if is_accelerate_available:
for key in vae_sd:
set_module_tensor_to_device(vae, key, dtype=torch.bfloat16, device=offload_device, value=vae_sd[key])
set_module_tensor_to_device(vae, key, dtype=dtype, device=offload_device, value=vae_sd[key])
else:
vae.load_state_dict(vae_sd, strict=True)
vae.eval().to(torch.bfloat16).to("cpu")
vae.eval().to(dtype).to("cpu")
del vae_sd
return (model, vae,)