From be9e8a64c9f0fd375cc02f491721b1753748da26 Mon Sep 17 00:00:00 2001 From: bash-j Date: Fri, 16 Aug 2024 10:57:05 +0930 Subject: [PATCH] save_checkpoint VAE optional Make VAE optional in save_checkpoint. Function argument has vae=None, but you call a method on the vae object without checking first if it is None. --- comfy/sd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index edd0b51d8..b24498bd8 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -668,7 +668,8 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m model_management.load_models_gpu(load_models, force_patch_weights=True) clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None - sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), clip_vision_sd) + vae = vae.get_sd() if vae is not None else None + sd = model.model.state_dict_for_saving(clip_sd, vae, clip_vision_sd) for k in extra_keys: sd[k] = extra_keys[k]