Support saving some more modelspec types.

This commit is contained in:
comfyanonymous 2024-05-13 21:54:11 -04:00
parent cf6e1efb69
commit 2de3b69b30
1 changed files with 6 additions and 1 deletions

View File

@ -175,9 +175,14 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi
enable_modelspec = True
if isinstance(model.model, comfy.model_base.SDXL):
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
if isinstance(model.model, comfy.model_base.SDXL_instructpix2pix):
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-edit"
else:
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
elif isinstance(model.model, comfy.model_base.SDXLRefiner):
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
elif isinstance(model.model, comfy.model_base.SVD_img2vid):
metadata["modelspec.architecture"] = "stable-video-diffusion-img2vid-v1"
else:
enable_modelspec = False