Fix loading SD1.5 diffusers checkpoint.

This commit is contained in:
comfyanonymous 2023-04-07 01:28:15 -04:00
parent 8b9ac8fedb
commit 58ed0f2da4
2 changed files with 5 additions and 3 deletions

View File

@ -272,7 +272,8 @@ def load_diffusers(model_path, fp16=True, output_vae=True, output_clip=True, emb
# magic
v2 = diffusers_unet_conf["sample_size"] == 96
v_pred = diffusers_scheduler_conf['prediction_type'] == 'v_prediction'
if 'prediction_type' in diffusers_scheduler_conf:
v_pred = diffusers_scheduler_conf['prediction_type'] == 'v_prediction'
if v2:
if v_pred:
@ -290,6 +291,7 @@ def load_diffusers(model_path, fp16=True, output_vae=True, output_clip=True, emb
scale_factor = model_config_params['scale_factor']
vae_config = model_config_params['first_stage_config']
vae_config['scale_factor'] = scale_factor
model_config_params["unet_config"]["params"]["use_fp16"] = fp16
unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.safetensors")

View File

@ -231,7 +231,7 @@ class DiffusersLoader:
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
FUNCTION = "load_checkpoint"
CATEGORY = "loaders"
CATEGORY = "advanced/loaders"
def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
for search_path in folder_paths.get_folder_paths("diffusers"):
@ -240,7 +240,7 @@ class DiffusersLoader:
if model_path in paths:
model_path = os.path.join(search_path, model_path)
break
search_paths = folder_paths.get_folder_paths("diffusers")
return load_diffusers(model_path, fp16=model_management.should_use_fp16(), output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))