mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-11 12:06:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support properly saving CosXL checkpoints.

Browse Source
This commit is contained in:
comfyanonymous
2024-04-08 00:36:22 -04:00
parent d644b6bcd8
commit 30abc324c2
2 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/sd.py
View File

@@ -600,7 +600,7 @@ def load_unet(unet_path):
raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
return model
def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None):
def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}):
clip_sd = None
load_models = [model]
if clip is not None:
@@ -610,4 +610,7 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
model_management.load_models_gpu(load_models)
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)
for k in extra_keys:
sd[k] = extra_keys[k]
comfy.utils.save_torch_file(sd, output_path, metadata=metadata)