John6666 commited on
Commit
0dcf861
1 Parent(s): a846f5a

Upload 5 files

Browse files
README.md CHANGED
@@ -1,12 +1,12 @@
1
- ---
2
- title: Convert Repo To Safetensors Sd
3
- emoji: 😻
4
- colorFrom: green
5
- colorTo: yellow
6
- sdk: gradio
7
- sdk_version: 4.38.1
8
- app_file: app.py
9
- pinned: false
10
- ---
11
-
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
1
+ ---
2
+ title: Convert diffusers SD1.5 repo to single Safetensors
3
+ emoji: 🐶
4
+ colorFrom: yellow
5
+ colorTo: red
6
+ sdk: gradio
7
+ sdk_version: 4.38.1
8
+ app_file: app.py
9
+ pinned: false
10
+ ---
11
+
12
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ from convert_repo_to_safetensors_sd_gr import convert_repo_to_safetensors_multi_sd
4
+ os.environ['HF_OUTPUT_REPO'] = 'John6666/safetensors_converting_test'
5
+
6
+ css = """"""
7
+
8
+ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
9
+ with gr.Column():
10
+ repo_id = gr.Textbox(label="Repo ID", placeholder="author/model", value="", lines=1)
11
+ is_half = gr.Checkbox(label="Half precision", value=True)
12
+ is_upload = gr.Checkbox(label="Upload safetensors to HF Repo", info="Fast download, but files will be public.", value=False)
13
+ uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None)
14
+ run_button = gr.Button(value="Convert")
15
+ st_file = gr.Files(label="Output", interactive=False)
16
+ st_md = gr.Markdown()
17
+
18
+ gr.on(
19
+ triggers=[repo_id.submit, run_button.click],
20
+ fn=convert_repo_to_safetensors_multi_sd,
21
+ inputs=[repo_id, st_file, is_upload, uploaded_urls, is_half],
22
+ outputs=[st_file, uploaded_urls, st_md],
23
+ )
24
+
25
+ demo.queue()
26
+ demo.launch()
convert_repo_to_safetensors_sd.py ADDED
@@ -0,0 +1,263 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
2
+ # *Only* converts the UNet, VAE, and Text Encoder.
3
+ # Does not convert optimizer state or any other thing.
4
+ # Written by jachiam
5
+
6
+ import argparse
7
+ import os.path as osp
8
+
9
+ import torch
10
+
11
+
12
+ # =================#
13
+ # UNet Conversion #
14
+ # =================#
15
+
16
+ unet_conversion_map = [
17
+ # (stable-diffusion, HF Diffusers)
18
+ ("time_embed.0.weight", "time_embedding.linear_1.weight"),
19
+ ("time_embed.0.bias", "time_embedding.linear_1.bias"),
20
+ ("time_embed.2.weight", "time_embedding.linear_2.weight"),
21
+ ("time_embed.2.bias", "time_embedding.linear_2.bias"),
22
+ ("input_blocks.0.0.weight", "conv_in.weight"),
23
+ ("input_blocks.0.0.bias", "conv_in.bias"),
24
+ ("out.0.weight", "conv_norm_out.weight"),
25
+ ("out.0.bias", "conv_norm_out.bias"),
26
+ ("out.2.weight", "conv_out.weight"),
27
+ ("out.2.bias", "conv_out.bias"),
28
+ ]
29
+
30
+ unet_conversion_map_resnet = [
31
+ # (stable-diffusion, HF Diffusers)
32
+ ("in_layers.0", "norm1"),
33
+ ("in_layers.2", "conv1"),
34
+ ("out_layers.0", "norm2"),
35
+ ("out_layers.3", "conv2"),
36
+ ("emb_layers.1", "time_emb_proj"),
37
+ ("skip_connection", "conv_shortcut"),
38
+ ]
39
+
40
+ unet_conversion_map_layer = []
41
+ # hardcoded number of downblocks and resnets/attentions...
42
+ # would need smarter logic for other networks.
43
+ for i in range(4):
44
+ # loop over downblocks/upblocks
45
+
46
+ for j in range(2):
47
+ # loop over resnets/attentions for downblocks
48
+ hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
49
+ sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
50
+ unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
51
+
52
+ if i < 3:
53
+ # no attention layers in down_blocks.3
54
+ hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
55
+ sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
56
+ unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
57
+
58
+ for j in range(3):
59
+ # loop over resnets/attentions for upblocks
60
+ hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
61
+ sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
62
+ unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
63
+
64
+ if i > 0:
65
+ # no attention layers in up_blocks.0
66
+ hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
67
+ sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
68
+ unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
69
+
70
+ if i < 3:
71
+ # no downsample in down_blocks.3
72
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
73
+ sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
74
+ unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
75
+
76
+ # no upsample in up_blocks.3
77
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
78
+ sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
79
+ unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
80
+
81
+ hf_mid_atn_prefix = "mid_block.attentions.0."
82
+ sd_mid_atn_prefix = "middle_block.1."
83
+ unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
84
+
85
+ for j in range(2):
86
+ hf_mid_res_prefix = f"mid_block.resnets.{j}."
87
+ sd_mid_res_prefix = f"middle_block.{2*j}."
88
+ unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
89
+
90
+
91
+ def convert_unet_state_dict(unet_state_dict):
92
+ # buyer beware: this is a *brittle* function,
93
+ # and correct output requires that all of these pieces interact in
94
+ # the exact order in which I have arranged them.
95
+ mapping = {k: k for k in unet_state_dict.keys()}
96
+ for sd_name, hf_name in unet_conversion_map:
97
+ mapping[hf_name] = sd_name
98
+ for k, v in mapping.items():
99
+ if "resnets" in k:
100
+ for sd_part, hf_part in unet_conversion_map_resnet:
101
+ v = v.replace(hf_part, sd_part)
102
+ mapping[k] = v
103
+ for k, v in mapping.items():
104
+ for sd_part, hf_part in unet_conversion_map_layer:
105
+ v = v.replace(hf_part, sd_part)
106
+ mapping[k] = v
107
+ new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
108
+ return new_state_dict
109
+
110
+
111
+ # ================#
112
+ # VAE Conversion #
113
+ # ================#
114
+
115
+ vae_conversion_map = [
116
+ # (stable-diffusion, HF Diffusers)
117
+ ("nin_shortcut", "conv_shortcut"),
118
+ ("norm_out", "conv_norm_out"),
119
+ ("mid.attn_1.", "mid_block.attentions.0."),
120
+ ]
121
+
122
+ for i in range(4):
123
+ # down_blocks have two resnets
124
+ for j in range(2):
125
+ hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
126
+ sd_down_prefix = f"encoder.down.{i}.block.{j}."
127
+ vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
128
+
129
+ if i < 3:
130
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
131
+ sd_downsample_prefix = f"down.{i}.downsample."
132
+ vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
133
+
134
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
135
+ sd_upsample_prefix = f"up.{3-i}.upsample."
136
+ vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
137
+
138
+ # up_blocks have three resnets
139
+ # also, up blocks in hf are numbered in reverse from sd
140
+ for j in range(3):
141
+ hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
142
+ sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
143
+ vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
144
+
145
+ # this part accounts for mid blocks in both the encoder and the decoder
146
+ for i in range(2):
147
+ hf_mid_res_prefix = f"mid_block.resnets.{i}."
148
+ sd_mid_res_prefix = f"mid.block_{i+1}."
149
+ vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
150
+
151
+
152
+ vae_conversion_map_attn = [
153
+ # (stable-diffusion, HF Diffusers)
154
+ ("norm.", "group_norm."),
155
+ ("q.", "query."),
156
+ ("k.", "key."),
157
+ ("v.", "value."),
158
+ ("proj_out.", "proj_attn."),
159
+ ]
160
+
161
+
162
+ def reshape_weight_for_sd(w):
163
+ # convert HF linear weights to SD conv2d weights
164
+ return w.reshape(*w.shape, 1, 1)
165
+
166
+
167
+ def convert_vae_state_dict(vae_state_dict):
168
+ mapping = {k: k for k in vae_state_dict.keys()}
169
+ for k, v in mapping.items():
170
+ for sd_part, hf_part in vae_conversion_map:
171
+ v = v.replace(hf_part, sd_part)
172
+ mapping[k] = v
173
+ for k, v in mapping.items():
174
+ if "attentions" in k:
175
+ for sd_part, hf_part in vae_conversion_map_attn:
176
+ v = v.replace(hf_part, sd_part)
177
+ mapping[k] = v
178
+ new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
179
+ weights_to_convert = ["q", "k", "v", "proj_out"]
180
+ for k, v in new_state_dict.items():
181
+ for weight_name in weights_to_convert:
182
+ if f"mid.attn_1.{weight_name}.weight" in k:
183
+ print(f"Reshaping {k} for SD format")
184
+ new_state_dict[k] = reshape_weight_for_sd(v)
185
+ return new_state_dict
186
+
187
+
188
+ # =========================#
189
+ # Text Encoder Conversion #
190
+ # =========================#
191
+ # pretty much a no-op
192
+
193
+
194
+ def convert_text_enc_state_dict(text_enc_dict):
195
+ return text_enc_dict
196
+
197
+
198
+ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True):
199
+ from safetensors.torch import load_file
200
+ input_safetensors = False
201
+ unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
202
+ if not osp.exists(unet_path):
203
+ unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
204
+ input_safetensors = True
205
+ vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
206
+ if not osp.exists(vae_path):
207
+ vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
208
+ input_safetensors = True
209
+ text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
210
+ if not osp.exists(text_enc_path):
211
+ text_enc_path = osp.join(model_path, "text_encoder", "model.safetensors")
212
+ input_safetensors = True
213
+
214
+ # Convert the UNet model
215
+ unet_state_dict = torch.load(unet_path, map_location='cpu') if not input_safetensors else load_file(unet_path)
216
+ unet_state_dict = convert_unet_state_dict(unet_state_dict)
217
+ unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
218
+
219
+ # Convert the VAE model
220
+ vae_state_dict = torch.load(vae_path, map_location='cpu') if not input_safetensors else load_file(vae_path)
221
+ vae_state_dict = convert_vae_state_dict(vae_state_dict)
222
+ vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
223
+
224
+ # Convert the text encoder model
225
+ text_enc_dict = torch.load(text_enc_path, map_location='cpu') if not input_safetensors else load_file(text_enc_path)
226
+ text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
227
+ text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
228
+
229
+ # Put together new checkpoint
230
+ state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
231
+ if half:
232
+ state_dict = {k:v.half() for k,v in state_dict.items()}
233
+ state_dict = {"state_dict": state_dict}
234
+ torch.save(state_dict, checkpoint_path)
235
+
236
+
237
+ def download_repo(repo_id, dir_path):
238
+ from huggingface_hub import snapshot_download
239
+ try:
240
+ snapshot_download(repo_id=repo_id, local_dir=dir_path)
241
+ except Exception as e:
242
+ print(f"Error: Failed to download {repo_id}. ")
243
+ return
244
+
245
+
246
+ def convert_repo_to_safetensors(repo_id, half = True):
247
+ download_dir = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}"
248
+ output_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}.safetensors"
249
+ download_repo(repo_id, download_dir)
250
+ convert_diffusers_to_safetensors(download_dir, output_filename, half)
251
+ return output_filename
252
+
253
+
254
+ if __name__ == "__main__":
255
+ parser = argparse.ArgumentParser()
256
+
257
+ parser.add_argument("--repo_id", default=None, type=str, required=True, help="HF Repo ID of the model to convert.")
258
+ parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
259
+
260
+ args = parser.parse_args()
261
+ assert args.repo_id is not None, "Must provide a Repo ID!"
262
+
263
+ convert_repo_to_safetensors(args.repo_id, args.half)
convert_repo_to_safetensors_sd_gr.py ADDED
@@ -0,0 +1,300 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
2
+ # *Only* converts the UNet, VAE, and Text Encoder.
3
+ # Does not convert optimizer state or any other thing.
4
+ # Written by jachiam
5
+
6
+ import argparse
7
+ import os.path as osp
8
+
9
+ import torch
10
+ import gradio as gr
11
+
12
+ # =================#
13
+ # UNet Conversion #
14
+ # =================#
15
+
16
+ unet_conversion_map = [
17
+ # (stable-diffusion, HF Diffusers)
18
+ ("time_embed.0.weight", "time_embedding.linear_1.weight"),
19
+ ("time_embed.0.bias", "time_embedding.linear_1.bias"),
20
+ ("time_embed.2.weight", "time_embedding.linear_2.weight"),
21
+ ("time_embed.2.bias", "time_embedding.linear_2.bias"),
22
+ ("input_blocks.0.0.weight", "conv_in.weight"),
23
+ ("input_blocks.0.0.bias", "conv_in.bias"),
24
+ ("out.0.weight", "conv_norm_out.weight"),
25
+ ("out.0.bias", "conv_norm_out.bias"),
26
+ ("out.2.weight", "conv_out.weight"),
27
+ ("out.2.bias", "conv_out.bias"),
28
+ ]
29
+
30
+ unet_conversion_map_resnet = [
31
+ # (stable-diffusion, HF Diffusers)
32
+ ("in_layers.0", "norm1"),
33
+ ("in_layers.2", "conv1"),
34
+ ("out_layers.0", "norm2"),
35
+ ("out_layers.3", "conv2"),
36
+ ("emb_layers.1", "time_emb_proj"),
37
+ ("skip_connection", "conv_shortcut"),
38
+ ]
39
+
40
+ unet_conversion_map_layer = []
41
+ # hardcoded number of downblocks and resnets/attentions...
42
+ # would need smarter logic for other networks.
43
+ for i in range(4):
44
+ # loop over downblocks/upblocks
45
+
46
+ for j in range(2):
47
+ # loop over resnets/attentions for downblocks
48
+ hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
49
+ sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
50
+ unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
51
+
52
+ if i < 3:
53
+ # no attention layers in down_blocks.3
54
+ hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
55
+ sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
56
+ unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
57
+
58
+ for j in range(3):
59
+ # loop over resnets/attentions for upblocks
60
+ hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
61
+ sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
62
+ unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
63
+
64
+ if i > 0:
65
+ # no attention layers in up_blocks.0
66
+ hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
67
+ sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
68
+ unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
69
+
70
+ if i < 3:
71
+ # no downsample in down_blocks.3
72
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
73
+ sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
74
+ unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
75
+
76
+ # no upsample in up_blocks.3
77
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
78
+ sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
79
+ unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
80
+
81
+ hf_mid_atn_prefix = "mid_block.attentions.0."
82
+ sd_mid_atn_prefix = "middle_block.1."
83
+ unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
84
+
85
+ for j in range(2):
86
+ hf_mid_res_prefix = f"mid_block.resnets.{j}."
87
+ sd_mid_res_prefix = f"middle_block.{2*j}."
88
+ unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
89
+
90
+
91
+ def convert_unet_state_dict(unet_state_dict):
92
+ # buyer beware: this is a *brittle* function,
93
+ # and correct output requires that all of these pieces interact in
94
+ # the exact order in which I have arranged them.
95
+ mapping = {k: k for k in unet_state_dict.keys()}
96
+ for sd_name, hf_name in unet_conversion_map:
97
+ mapping[hf_name] = sd_name
98
+ for k, v in mapping.items():
99
+ if "resnets" in k:
100
+ for sd_part, hf_part in unet_conversion_map_resnet:
101
+ v = v.replace(hf_part, sd_part)
102
+ mapping[k] = v
103
+ for k, v in mapping.items():
104
+ for sd_part, hf_part in unet_conversion_map_layer:
105
+ v = v.replace(hf_part, sd_part)
106
+ mapping[k] = v
107
+ new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
108
+ return new_state_dict
109
+
110
+
111
+ # ================#
112
+ # VAE Conversion #
113
+ # ================#
114
+
115
+ vae_conversion_map = [
116
+ # (stable-diffusion, HF Diffusers)
117
+ ("nin_shortcut", "conv_shortcut"),
118
+ ("norm_out", "conv_norm_out"),
119
+ ("mid.attn_1.", "mid_block.attentions.0."),
120
+ ]
121
+
122
+ for i in range(4):
123
+ # down_blocks have two resnets
124
+ for j in range(2):
125
+ hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
126
+ sd_down_prefix = f"encoder.down.{i}.block.{j}."
127
+ vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
128
+
129
+ if i < 3:
130
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
131
+ sd_downsample_prefix = f"down.{i}.downsample."
132
+ vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
133
+
134
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
135
+ sd_upsample_prefix = f"up.{3-i}.upsample."
136
+ vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
137
+
138
+ # up_blocks have three resnets
139
+ # also, up blocks in hf are numbered in reverse from sd
140
+ for j in range(3):
141
+ hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
142
+ sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
143
+ vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
144
+
145
+ # this part accounts for mid blocks in both the encoder and the decoder
146
+ for i in range(2):
147
+ hf_mid_res_prefix = f"mid_block.resnets.{i}."
148
+ sd_mid_res_prefix = f"mid.block_{i+1}."
149
+ vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
150
+
151
+
152
+ vae_conversion_map_attn = [
153
+ # (stable-diffusion, HF Diffusers)
154
+ ("norm.", "group_norm."),
155
+ ("q.", "query."),
156
+ ("k.", "key."),
157
+ ("v.", "value."),
158
+ ("proj_out.", "proj_attn."),
159
+ ]
160
+
161
+
162
+ def reshape_weight_for_sd(w):
163
+ # convert HF linear weights to SD conv2d weights
164
+ return w.reshape(*w.shape, 1, 1)
165
+
166
+
167
+ def convert_vae_state_dict(vae_state_dict):
168
+ mapping = {k: k for k in vae_state_dict.keys()}
169
+ for k, v in mapping.items():
170
+ for sd_part, hf_part in vae_conversion_map:
171
+ v = v.replace(hf_part, sd_part)
172
+ mapping[k] = v
173
+ for k, v in mapping.items():
174
+ if "attentions" in k:
175
+ for sd_part, hf_part in vae_conversion_map_attn:
176
+ v = v.replace(hf_part, sd_part)
177
+ mapping[k] = v
178
+ new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
179
+ weights_to_convert = ["q", "k", "v", "proj_out"]
180
+ for k, v in new_state_dict.items():
181
+ for weight_name in weights_to_convert:
182
+ if f"mid.attn_1.{weight_name}.weight" in k:
183
+ print(f"Reshaping {k} for SD format")
184
+ new_state_dict[k] = reshape_weight_for_sd(v)
185
+ return new_state_dict
186
+
187
+
188
+ # =========================#
189
+ # Text Encoder Conversion #
190
+ # =========================#
191
+ # pretty much a no-op
192
+
193
+
194
+ def convert_text_enc_state_dict(text_enc_dict):
195
+ return text_enc_dict
196
+
197
+
198
+ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, progress=gr.Progress(track_tqdm=True)):
199
+ progress(0, desc="Start converting...")
200
+ from safetensors.torch import load_file
201
+ input_safetensors = False
202
+ unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
203
+ if not osp.exists(unet_path):
204
+ unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
205
+ input_safetensors = True
206
+ vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
207
+ if not osp.exists(vae_path):
208
+ vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
209
+ input_safetensors = True
210
+ text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
211
+ if not osp.exists(text_enc_path):
212
+ text_enc_path = osp.join(model_path, "text_encoder", "model.safetensors")
213
+ input_safetensors = True
214
+
215
+ # Convert the UNet model
216
+ unet_state_dict = torch.load(unet_path, map_location='cpu') if not input_safetensors else load_file(unet_path)
217
+ unet_state_dict = convert_unet_state_dict(unet_state_dict)
218
+ unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
219
+
220
+ # Convert the VAE model
221
+ vae_state_dict = torch.load(vae_path, map_location='cpu') if not input_safetensors else load_file(vae_path)
222
+ vae_state_dict = convert_vae_state_dict(vae_state_dict)
223
+ vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
224
+
225
+ # Convert the text encoder model
226
+ text_enc_dict = torch.load(text_enc_path, map_location='cpu') if not input_safetensors else load_file(text_enc_path)
227
+ text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
228
+ text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
229
+
230
+ # Put together new checkpoint
231
+ state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
232
+ if half:
233
+ state_dict = {k:v.half() for k,v in state_dict.items()}
234
+ state_dict = {"state_dict": state_dict}
235
+ torch.save(state_dict, checkpoint_path)
236
+
237
+ progress(1, desc="Converted.")
238
+
239
+
240
+ def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
241
+ from huggingface_hub import snapshot_download
242
+ try:
243
+ snapshot_download(repo_id=repo_id, local_dir=dir_path)
244
+ except Exception as e:
245
+ print(f"Error: Failed to download {repo_id}. ")
246
+ return
247
+
248
+
249
+ def upload_safetensors_to_repo(filename, progress=gr.Progress(track_tqdm=True)):
250
+ from huggingface_hub import HfApi, hf_hub_url
251
+ import os
252
+ from pathlib import Path
253
+ output_filename = Path(filename).name
254
+ hf_token = os.environ.get("HF_TOKEN")
255
+ repo_id = os.environ.get("HF_OUTPUT_REPO")
256
+ api = HfApi()
257
+ try:
258
+ progress(0, desc="Start uploading...")
259
+ api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_id=repo_id, token=hf_token)
260
+ progress(1, desc="Uploaded.")
261
+ url = hf_hub_url(repo_id=repo_id, filename=output_filename)
262
+ except Exception as e:
263
+ print(f"Error: Failed to upload to {repo_id}. ")
264
+ return None
265
+ return url
266
+
267
+
268
+ def convert_repo_to_safetensors(repo_id, half = True, progress=gr.Progress(track_tqdm=True)):
269
+ download_dir = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}"
270
+ output_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}.safetensors"
271
+ download_repo(repo_id, download_dir)
272
+ convert_diffusers_to_safetensors(download_dir, output_filename, half)
273
+ return output_filename
274
+
275
+
276
+ def convert_repo_to_safetensors_multi_sd(repo_id, files, is_upload, urls, half=True, progress=gr.Progress(track_tqdm=True)):
277
+ file = convert_repo_to_safetensors(repo_id, half)
278
+ if not urls: urls = []
279
+ url = ""
280
+ if is_upload:
281
+ url = upload_safetensors_to_repo(file, half)
282
+ if url: urls.append(url)
283
+ md = ""
284
+ for u in urls:
285
+ md += f"[Download {str(u).split('/')[-1]}]({str(u)})<br>"
286
+ if not files: files = []
287
+ files.append(file)
288
+ return gr.update(value=files), gr.update(value=urls, choices=urls), gr.update(value=md)
289
+
290
+
291
+ if __name__ == "__main__":
292
+ parser = argparse.ArgumentParser()
293
+
294
+ parser.add_argument("--repo_id", default=None, type=str, required=True, help="HF Repo ID of the model to convert.")
295
+ parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
296
+
297
+ args = parser.parse_args()
298
+ assert args.repo_id is not None, "Must provide a Repo ID!"
299
+
300
+ convert_repo_to_safetensors(args.repo_id, args.half)
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch
2
+ safetensors
3
+ huggingface-hub