Conversion script for ncsnpp models (#98) · python273/diffusers@f794432 · GitHub
Skip to content

Commit f794432

Browse files
ArthurZuckerArthurZucker
authored
Conversion script for ncsnpp models (huggingface#98)
* added kwargs for easier intialisation of random model * initial commit for conversion script * current debug script * update * Update * done * add updated debug conversion script * style * clean conversion script
1 parent 182b164 commit f794432

3 files changed

Lines changed: 252 additions & 64 deletions

File tree

‎debug_conversion.py‎

Lines changed: 87 additions & 63 deletions

‎scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
# coding=utf-8
2+
# Copyright 2022 The HuggingFace Inc. team.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
""" Conversion script for the NCSNPP checkpoints. """
16+
17+
import argparse
18+
import json
19+
import torch
20+
from diffusers import UNetUnconditionalModel
21+
22+
23+
24+
def convert_ncsnpp_checkpoint(checkpoint, config):
25+
"""
26+
Takes a state dict and the path to
27+
"""
28+
new_model_architecture = UNetUnconditionalModel(**config)
29+
new_model_architecture.time_steps.W.data= checkpoint['all_modules.0.W'].data
30+
new_model_architecture.time_steps.weight.data = checkpoint['all_modules.0.W'].data
31+
new_model_architecture.time_embedding.linear_1.weight.data = checkpoint['all_modules.1.weight'].data
32+
new_model_architecture.time_embedding.linear_1.bias.data = checkpoint['all_modules.1.bias'].data
33+
34+
new_model_architecture.time_embedding.linear_2.weight.data = checkpoint['all_modules.2.weight'].data
35+
new_model_architecture.time_embedding.linear_2.bias.data= checkpoint['all_modules.2.bias'].data
36+
37+
new_model_architecture.conv_in.weight.data = checkpoint['all_modules.3.weight'].data
38+
new_model_architecture.conv_in.bias.data = checkpoint['all_modules.3.bias'].data
39+
40+
new_model_architecture.conv_norm_out.weight.data = checkpoint[list(checkpoint.keys())[-4]].data
41+
new_model_architecture.conv_norm_out.bias.data = checkpoint[list(checkpoint.keys())[-3]].data
42+
new_model_architecture.conv_out.weight.data = checkpoint[list(checkpoint.keys())[-2]].data
43+
new_model_architecture.conv_out.bias.data = checkpoint[list(checkpoint.keys())[-1]].data
44+
45+
module_index = 4
46+
47+
48+
def set_attention_weights(new_layer,old_checkpoint,index):
49+
new_layer.query.weight.data = old_checkpoint[f"all_modules.{index}.NIN_0.W"].data.T
50+
new_layer.key.weight.data = old_checkpoint[f"all_modules.{index}.NIN_1.W"].data.T
51+
new_layer.value.weight.data = old_checkpoint[f"all_modules.{index}.NIN_2.W"].data.T
52+
53+
new_layer.query.bias.data = old_checkpoint[f"all_modules.{index}.NIN_0.b"].data
54+
new_layer.key.bias.data = old_checkpoint[f"all_modules.{index}.NIN_1.b"].data
55+
new_layer.value.bias.data = old_checkpoint[f"all_modules.{index}.NIN_2.b"].data
56+
57+
new_layer.proj_attn.weight.data = old_checkpoint[f"all_modules.{index}.NIN_3.W"].data.T
58+
new_layer.proj_attn.bias.data = old_checkpoint[f"all_modules.{index}.NIN_3.b"].data
59+
60+
new_layer.group_norm.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data
61+
new_layer.group_norm.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.bias"].data
62+
63+
def set_resnet_weights(new_layer,old_checkpoint,index):
64+
new_layer.conv1.weight.data = old_checkpoint[f"all_modules.{index}.Conv_0.weight"].data
65+
new_layer.conv1.bias.data = old_checkpoint[f"all_modules.{index}.Conv_0.bias"].data
66+
new_layer.norm1.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data
67+
new_layer.norm1.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.bias"].data
68+
69+
new_layer.conv2.weight.data = old_checkpoint[f"all_modules.{index}.Conv_1.weight"].data
70+
new_layer.conv2.bias.data = old_checkpoint[f"all_modules.{index}.Conv_1.bias"].data
71+
new_layer.norm2.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_1.weight"].data
72+
new_layer.norm2.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_1.bias"].data
73+
74+
new_layer.time_emb_proj.weight.data = old_checkpoint[f"all_modules.{index}.Dense_0.weight"].data
75+
new_layer.time_emb_proj.bias.data = old_checkpoint[f"all_modules.{index}.Dense_0.bias"].data
76+
77+
if new_layer.in_channels != new_layer.out_channels or new_layer.up or new_layer.down:
78+
new_layer.conv_shortcut.weight.data = old_checkpoint[f"all_modules.{index}.Conv_2.weight"].data
79+
new_layer.conv_shortcut.bias.data = old_checkpoint[f"all_modules.{index}.Conv_2.bias"].data
80+
81+
for i, block in enumerate(new_model_architecture.downsample_blocks):
82+
has_attentions = hasattr(block, "attentions")
83+
for j in range(len(block.resnets)):
84+
set_resnet_weights(block.resnets[j],checkpoint, module_index)
85+
module_index += 1
86+
if has_attentions:
87+
set_attention_weights(block.attentions[j],checkpoint, module_index)
88+
module_index += 1
89+
90+
if hasattr(block, "downsamplers") and block.downsamplers is not None:
91+
set_resnet_weights(block.resnet_down,checkpoint, module_index)
92+
module_index += 1
93+
block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.Conv_0.weight"].data
94+
block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.Conv_0.bias"].data
95+
module_index += 1
96+
97+
98+
99+
set_resnet_weights(new_model_architecture.mid.resnets[0],checkpoint,module_index)
100+
module_index += 1
101+
set_attention_weights(new_model_architecture.mid.attentions[0],checkpoint, module_index)
102+
module_index += 1
103+
set_resnet_weights(new_model_architecture.mid.resnets[1],checkpoint,module_index)
104+
module_index += 1
105+
106+
for i, block in enumerate(new_model_architecture.upsample_blocks):
107+
has_attentions = hasattr(block, "attentions")
108+
for j in range(len(block.resnets)):
109+
set_resnet_weights(block.resnets[j],checkpoint, module_index)
110+
module_index += 1
111+
if has_attentions:
112+
set_attention_weights(block.attentions[0],checkpoint, module_index) # why can there only be a single attention layer for up?
113+
module_index += 1
114+
115+
if hasattr(block, "resnet_up") and block.resnet_up is not None:
116+
block.skip_norm.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
117+
block.skip_norm.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data
118+
module_index += 1
119+
block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
120+
block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data
121+
module_index += 1
122+
set_resnet_weights(block.resnet_up,checkpoint, module_index)
123+
module_index += 1
124+
125+
new_model_architecture.conv_norm_out.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
126+
new_model_architecture.conv_norm_out.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data
127+
module_index += 1
128+
new_model_architecture.conv_out.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
129+
new_model_architecture.conv_out.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data
130+
131+
return new_model_architecture.state_dict()
132+
133+
if __name__ == "__main__":
134+
parser = argparse.ArgumentParser()
135+
136+
parser.add_argument(
137+
"--checkpoint_path", default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model.pt", type=str, required=False, help="Path to the checkpoint to convert."
138+
)
139+
140+
parser.add_argument(
141+
"--config_file",
142+
default="/Users/arthurzucker/Work/diffusers/ArthurZ/config.json",
143+
type=str,
144+
required=False,
145+
help="The config json file corresponding to the architecture.",
146+
)
147+
148+
parser.add_argument(
149+
"--dump_path", default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt", type=str, required=False, help="Path to the output model."
150+
)
151+
152+
args = parser.parse_args()
153+
154+
155+
156+
157+
checkpoint = torch.load(args.checkpoint_path, map_location="cpu")
158+
159+
with open(args.config_file) as f:
160+
config = json.loads(f.read())
161+
162+
163+
converted_checkpoint = convert_ncsnpp_checkpoint(checkpoint, config,)
164+
torch.save(converted_checkpoint, args.dump_path)

‎src/diffusers/models/unet_unconditional.py‎

Lines changed: 1 addition & 1 deletion

0 commit comments

Comments
 (0)