Refactor diffusers model convert code to be able to reuse it. · code/app-python-comfyui@0fc483d · GitHub
Skip to content

Commit 0fc483d

Browse files
comfyanonymouscomfyanonymous
committed
Refactor diffusers model convert code to be able to reuse it.
1 parent f3ac938 commit 0fc483d

3 files changed

Lines changed: 113 additions & 109 deletions

File tree

‎comfy/diffusers_convert.py‎

Lines changed: 0 additions & 107 deletions

‎comfy/diffusers_load.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
import json
2+
import os
3+
import yaml
4+
5+
import folder_paths
6+
from comfy.ldm.util import instantiate_from_config
7+
from comfy.sd import ModelPatcher, load_model_weights, CLIP, VAE
8+
import os.path as osp
9+
import re
10+
import torch
11+
from safetensors.torch import load_file, save_file
12+
import diffusers_convert
13+
14+
def load_diffusers(model_path, fp16=True, output_vae=True, output_clip=True, embedding_directory=None):
15+
diffusers_unet_conf = json.load(open(osp.join(model_path, "unet/config.json")))
16+
diffusers_scheduler_conf = json.load(open(osp.join(model_path, "scheduler/scheduler_config.json")))
17+
18+
# magic
19+
v2 = diffusers_unet_conf["sample_size"] == 96
20+
if 'prediction_type' in diffusers_scheduler_conf:
21+
v_pred = diffusers_scheduler_conf['prediction_type'] == 'v_prediction'
22+
23+
if v2:
24+
if v_pred:
25+
config_path = folder_paths.get_full_path("configs", 'v2-inference-v.yaml')
26+
else:
27+
config_path = folder_paths.get_full_path("configs", 'v2-inference.yaml')
28+
else:
29+
config_path = folder_paths.get_full_path("configs", 'v1-inference.yaml')
30+
31+
with open(config_path, 'r') as stream:
32+
config = yaml.safe_load(stream)
33+
34+
model_config_params = config['model']['params']
35+
clip_config = model_config_params['cond_stage_config']
36+
scale_factor = model_config_params['scale_factor']
37+
vae_config = model_config_params['first_stage_config']
38+
vae_config['scale_factor'] = scale_factor
39+
model_config_params["unet_config"]["params"]["use_fp16"] = fp16
40+
41+
unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
42+
vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
43+
text_enc_path = osp.join(model_path, "text_encoder", "model.safetensors")
44+
45+
# Load models from safetensors if it exists, if it doesn't pytorch
46+
if osp.exists(unet_path):
47+
unet_state_dict = load_file(unet_path, device="cpu")
48+
else:
49+
unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
50+
unet_state_dict = torch.load(unet_path, map_location="cpu")
51+
52+
if osp.exists(vae_path):
53+
vae_state_dict = load_file(vae_path, device="cpu")
54+
else:
55+
vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
56+
vae_state_dict = torch.load(vae_path, map_location="cpu")
57+
58+
if osp.exists(text_enc_path):
59+
text_enc_dict = load_file(text_enc_path, device="cpu")
60+
else:
61+
text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
62+
text_enc_dict = torch.load(text_enc_path, map_location="cpu")
63+
64+
# Convert the UNet model
65+
unet_state_dict = diffusers_convert.convert_unet_state_dict(unet_state_dict)
66+
unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
67+
68+
# Convert the VAE model
69+
vae_state_dict = diffusers_convert.convert_vae_state_dict(vae_state_dict)
70+
vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
71+
72+
# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
73+
is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict
74+
75+
if is_v20_model:
76+
# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
77+
text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()}
78+
text_enc_dict = diffusers_convert.convert_text_enc_state_dict_v20(text_enc_dict)
79+
text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
80+
else:
81+
text_enc_dict = diffusers_convert.convert_text_enc_state_dict(text_enc_dict)
82+
text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
83+
84+
# Put together new checkpoint
85+
sd = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
86+
87+
clip = None
88+
vae = None
89+
90+
class WeightsLoader(torch.nn.Module):
91+
pass
92+
93+
w = WeightsLoader()
94+
load_state_dict_to = []
95+
if output_vae:
96+
vae = VAE(scale_factor=scale_factor, config=vae_config)
97+
w.first_stage_model = vae.first_stage_model
98+
load_state_dict_to = [w]
99+
100+
if output_clip:
101+
clip = CLIP(config=clip_config, embedding_directory=embedding_directory)
102+
w.cond_stage_model = clip.cond_stage_model
103+
load_state_dict_to = [w]
104+
105+
model = instantiate_from_config(config["model"])
106+
model = load_model_weights(model, sd, verbose=False, load_state_dict_to=load_state_dict_to)
107+
108+
if fp16:
109+
model = model.half()
110+
111+
return ModelPatcher(model), clip, vae

‎nodes.py‎

Lines changed: 2 additions & 2 deletions

0 commit comments

Comments
 (0)