wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-10 05:22:34 +08:00
Code Issues Actions 14 Packages Projects Releases Wiki Activity

Add key to indicate checkpoint is v_prediction when saving.

Browse Source
This commit is contained in:
comfyanonymous 2023-07-18 00:25:53 -04:00
parent 22abe3af9f
commit cea5c2adfb
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
comfy/model_base.py
View File

@ -99,6 +99,10 @@ class BaseModel(torch.nn.Module):
if self.get_dtype() == torch.float16:
clip_state_dict = utils.convert_sd_to(clip_state_dict, torch.float16)
vae_state_dict = utils.convert_sd_to(vae_state_dict, torch.float16)
if self.model_type == ModelType.V_PREDICTION:
unet_state_dict["v_pred"] = torch.tensor([])
return {**unet_state_dict, **vae_state_dict, **clip_state_dict}