from numpy import dtype from safetensors import safe_open import torch import argparse import sys import re folder = 'R:/SD/models/stable-diffusion-v1-5/' #folder = '/home/debreuil/SD/models/stable-diffusion-v1-5/' enc_file = f'{folder}text_encoder/model.safetensors' enc16_file = f'{folder}text_encoder/pytorch_model.safetensors' vae_file = f'{folder}vae/diffusion_pytorch_model.safetensors' v15_file = f'{folder}v1-5-pruned-emaonly.safetensors' v15_deliberate_v2_file = f'{folder}deliberate_v2.safetensors' realisticVision = f'{folder}realisticVisionV30_v30VAE.safetensors' revAnimated_v122 = f'{folder}revAnimated_v122.safetensors' lora_RA_file = f'{folder}Lora/LowRA.safetensors' lora_more_details_file = f'{folder}Lora/more_details.safetensors' lora_3DMM_V11_file = f'{folder}Lora/3DMM_V11.safetensors' lora_epi_noiseoffset2_file = f'{folder}Lora/epi_noiseoffset2.safetensors' control_canny_file = f'{folder}controlnet/control_canny-fp16.safetensors' control_depth_file = f'{folder}controlnet/control_depth-fp16.safetensors' target = control_depth_file class ConvertSafetensors: def __init__(self): self.is_model = False self.is_controlnet = False self.is_lora = False self.diffusion_key = "model.diffusion_model." self.lora_key = "lora_" self.unet_tensors = {"input_blocks":{},"middle_block":{}, "output_blocks":{}, "time_embed":{}, "out":{}, 'zero_convs':{}, 'input_hint_block':{}, 'middle_block_out':{},} # controlnet self.lora_tensors = {"te_text_model_encoder_layers":{}, "unet_down_blocks":{}, "unet_mid_block":{}, "unet_up_blocks":{}} self.converted = [] def parse(self, path:str): self.converted = [] self.get_tensors(path) if self.is_model: self.transform_full_sd() assert len(self.converted) == 686, f"tensor len should be 686 but is {len(self.converted)}." elif self.is_lora: self.transform_lora_sd() assert len(self.converted) == 792, f"tensor len should be 792 but is {len(self.converted)}." elif self.is_controlnet: self.transform_controlnet() assert len(self.converted) == 340, f"tensor len should be 340 but is {len(self.converted)}." out_file_path = target.replace(".safetensors", "_f16.bin") torch.save(self.converted, out_file_path) del self.converted def get_tensors(self, path:str): with safe_open(target, framework="pt", device=0) as st: txt_file = target.replace(".safetensors", ".txt") with open(txt_file, 'w') as file: for key in st.keys(): if key.startswith(self.diffusion_key) : self.is_model = True tail = key[len(self.diffusion_key):] elif key.startswith(self.lora_key): self.is_lora = True tail = key[len(self.lora_key):] else: self.is_controlnet = True tail = key tensor = st.get_tensor(key) tensor = tensor.to('cuda', dtype=torch.float16) if(self.is_lora): self.parse_lora_key(tail, tensor, file) else: self.parse_model_key(tail, tensor, file) def parse_model_key(self, tail, tensor, file): sp = tail.split(".") for block in self.unet_tensors.keys(): if(block == sp[0]): sub_key = tail[len(block) + 1:] match = re.search(r'\d+(\.\d+)?', sub_key) block_num = match.group() rem = sub_key[len(block_num) + 1:] if rem.startswith("transformer_blocks.0."): rem = rem[len("transformer_blocks.0."):] if block_num not in self.unet_tensors[block]: self.unet_tensors[block][block_num] = {} self.unet_tensors[block][block_num][rem] = tensor file.write("{:<40} {} {} {}\n".format(str(tensor.shape), block, block_num, rem)) break def parse_lora_key(self, tail, tensor, file): for block in self.lora_tensors.keys(): if tail.startswith(block): sub_key = tail[len(block) + 1:] match = re.search(r'^\d+', sub_key) if match: block_num = match.group() rem = sub_key[len(block_num) + 1:] else: block_num = "*" rem = sub_key # if rem.startswith("transformer_blocks.0."): # rem = rem[len("transformer_blocks.0."):] if block_num not in self.lora_tensors[block]: self.lora_tensors[block][block_num] = {} self.lora_tensors[block][block_num][rem] = tensor file.write("{:<30} {} {} {}\n".format(str(tensor.shape), block, block_num, rem)) break def transform_full_sd(self): self.transform_timesteps() self.transform_input_blocks() self.transform_middle_block() self.transform_output_blocks() self.transform_out() def transform_controlnet(self): self.transform_timesteps() self.transform_input_blocks() self.transform_middle_block() self.transform_zero_convs() self.transform_middle_block_out() self.transform_input_hint_block() def transform_lora_sd(self): self.transform_lora_text_encoder_layers() assert len(self.converted) == 216, f"text encoder tensors len should be 216 but is {len(self.converted)}." self.transform_lora_unet() def transform_lora_text_encoder_layers(self): tensor_dict = self.lora_tensors["te_text_model_encoder_layers"] self.add_MLPTransformer(tensor_dict["0"]) self.add_MLPTransformer(tensor_dict["1"]) self.add_MLPTransformer(tensor_dict["2"]) self.add_MLPTransformer(tensor_dict["3"]) self.add_MLPTransformer(tensor_dict["4"]) self.add_MLPTransformer(tensor_dict["5"]) self.add_MLPTransformer(tensor_dict["6"]) self.add_MLPTransformer(tensor_dict["7"]) self.add_MLPTransformer(tensor_dict["8"]) self.add_MLPTransformer(tensor_dict["9"]) self.add_MLPTransformer(tensor_dict["10"]) self.add_MLPTransformer(tensor_dict["11"]) def transform_lora_unet(self): down_dict = self.lora_tensors["unet_down_blocks"] self.add_lora_unet_block(down_dict["0"], 2) self.add_lora_unet_block(down_dict["1"], 2) self.add_lora_unet_block(down_dict["2"], 2) mid_dict = self.lora_tensors["unet_mid_block"] self.add_lora_unet_block(mid_dict["*"], 1) up_dict = self.lora_tensors["unet_up_blocks"] self.add_lora_unet_block(up_dict["1"], 3) self.add_lora_unet_block(up_dict["2"], 3) self.add_lora_unet_block(up_dict["3"], 3) def transform_zero_convs(self): tensor_dict = self.unet_tensors["zero_convs"] self.add_weight_bias(tensor_dict["0.0"]) self.add_weight_bias(tensor_dict["1.0"]) self.add_weight_bias(tensor_dict["2.0"]) self.add_weight_bias(tensor_dict["3.0"]) self.add_weight_bias(tensor_dict["4.0"]) self.add_weight_bias(tensor_dict["5.0"]) self.add_weight_bias(tensor_dict["6.0"]) self.add_weight_bias(tensor_dict["7.0"]) self.add_weight_bias(tensor_dict["8.0"]) self.add_weight_bias(tensor_dict["9.0"]) self.add_weight_bias(tensor_dict["10.0"]) self.add_weight_bias(tensor_dict["11.0"]) def transform_middle_block_out(self): tensor_dict = self.unet_tensors["middle_block_out"] self.add_weight_bias(tensor_dict["0"]) def transform_input_hint_block(self): tensor_dict = self.unet_tensors["input_hint_block"] self.add_weight_bias(tensor_dict["0"]) self.add_weight_bias(tensor_dict["2"]) self.add_weight_bias(tensor_dict["4"]) self.add_weight_bias(tensor_dict["6"]) self.add_weight_bias(tensor_dict["8"]) self.add_weight_bias(tensor_dict["10"]) self.add_weight_bias(tensor_dict["12"]) self.add_weight_bias(tensor_dict["14"]) def transform_timesteps(self): t_temb = self.unet_tensors["time_embed"] self.add_timesteps(t_temb["0"]) self.add_timesteps(t_temb["2"]) def transform_input_blocks(self): t_input = self.unet_tensors["input_blocks"] self.add_weight_bias(t_input["0.0"]) self.add_resnet(t_input["1.0"]) self.add_transformer(t_input["1.1"]) self.add_resnet(t_input["2.0"]) self.add_transformer(t_input["2.1"]) self.add_downsample(t_input["3.0"]) self.add_resnet(t_input["4.0"]) self.add_transformer(t_input["4.1"]) self.add_resnet(t_input["5.0"]) self.add_transformer(t_input["5.1"]) self.add_downsample(t_input["6.0"]) self.add_resnet(t_input["7.0"]) self.add_transformer(t_input["7.1"]) self.add_resnet(t_input["8.0"]) self.add_transformer(t_input["8.1"]) self.add_downsample(t_input["9.0"]) self.add_resnet(t_input["10.0"]) self.add_resnet(t_input["11.0"]) def transform_middle_block(self): t_middle = self.unet_tensors["middle_block"] self.add_resnet(t_middle["0"]) self.add_transformer(t_middle["1"]) self.add_resnet(t_middle["2"]) def transform_output_blocks(self): t_output = self.unet_tensors["output_blocks"] self.add_resnet(t_output["0.0"]) self.add_resnet(t_output["1.0"]) self.add_resnet(t_output["2.0"]) self.add_upsample(t_output["2.1"]) self.add_resnet(t_output["3.0"]) self.add_transformer(t_output["3.1"]) self.add_resnet(t_output["4.0"]) self.add_transformer(t_output["4.1"]) self.add_resnet(t_output["5.0"]) self.add_transformer(t_output["5.1"]) self.add_upsample(t_output["5.2"]) self.add_resnet(t_output["6.0"]) self.add_transformer(t_output["6.1"]) self.add_resnet(t_output["7.0"]) self.add_transformer(t_output["7.1"]) self.add_resnet(t_output["8.0"]) self.add_transformer(t_output["8.1"]) self.add_upsample(t_output["8.2"]) self.add_resnet(t_output["9.0"]) self.add_transformer(t_output["9.1"]) self.add_resnet(t_output["10.0"]) self.add_transformer(t_output["10.1"]) self.add_resnet(t_output["11.0"]) self.add_transformer(t_output["11.1"]) def transform_out(self): t_out = self.unet_tensors["out"] self.add_weight_bias(t_out["0"]) self.add_weight_bias(t_out["2"]) def add_timesteps(self, block:dict): assert len(block) == 2, "timestep length error" # self.ordered_modules = [self.linear_1, self.linear_2] self.converted.append(block["weight"]) self.converted.append(block["bias"]) def add_weight_bias(self, block:dict): assert len(block) == 2, "weight_bias length error" self.converted.append(block["weight"]) self.converted.append(block["bias"]) def add_resnet(self, block:dict): assert len(block) == 10 or len(block) == 12, "resnet length error" start_len = len(self.converted) # self.ordered_modules = [self.norm1, self.conv1, self.time_emb_proj, self.norm2, self.conv2] self.converted.append(block["in_layers.0.weight"]) self.converted.append(block["in_layers.0.bias"]) self.converted.append(block["in_layers.2.weight"]) self.converted.append(block["in_layers.2.bias"]) self.converted.append(block["emb_layers.1.weight"]) self.converted.append(block["emb_layers.1.bias"]) self.converted.append(block["out_layers.0.weight"]) self.converted.append(block["out_layers.0.bias"]) self.converted.append(block["out_layers.3.weight"]) self.converted.append(block["out_layers.3.bias"]) if "skip_connection.bias" in block.keys(): self.converted.append(block["skip_connection.weight"]) self.converted.append(block["skip_connection.bias"]) assert len(self.converted) - start_len == len(block), "Not all tensors parsed in resnet." def add_transformer(self, block:dict): assert len(block) == 26, "transform length error" start_len = len(self.converted) # self.ordered_modules = [ # self.norm, self.proj_in, # self.norm1, self.cross0_q, self.cross0_k, self.cross0_v, self.cross0_out, # self.norm2, self.cross1_q, self.cross1_k, self.cross1_v, self.cross1_out, # self.norm3, self.ff_GEGLU.proj, self.ff_out, # self.proj_out] self.converted.append(block["norm.weight"]) self.converted.append(block["norm.bias"]) self.converted.append(block["proj_in.weight"]) self.converted.append(block["proj_in.bias"]) self.converted.append(block["norm1.weight"]) self.converted.append(block["norm1.bias"]) self.converted.append(block["attn1.to_q.weight"]) self.converted.append(block["attn1.to_k.weight"]) self.converted.append(block["attn1.to_v.weight"]) self.converted.append(block["attn1.to_out.0.weight"]) self.converted.append(block["attn1.to_out.0.bias"]) self.converted.append(block["norm2.weight"]) self.converted.append(block["norm2.bias"]) self.converted.append(block["attn2.to_q.weight"]) self.converted.append(block["attn2.to_k.weight"]) self.converted.append(block["attn2.to_v.weight"]) self.converted.append(block["attn2.to_out.0.weight"]) self.converted.append(block["attn2.to_out.0.bias"]) self.converted.append(block["norm3.weight"]) self.converted.append(block["norm3.bias"]) self.converted.append(block["ff.net.0.proj.weight"]) self.converted.append(block["ff.net.0.proj.bias"]) self.converted.append(block["ff.net.2.weight"]) self.converted.append(block["ff.net.2.bias"]) self.converted.append(block["proj_out.weight"]) self.converted.append(block["proj_out.bias"]) assert len(self.converted) - start_len == len(block), "Not all tensors parsed in transformer." def add_downsample(self, block:dict): assert len(block) == 2, "downsample length error" self.converted.append(block["op.weight"]) self.converted.append(block["op.bias"]) def add_upsample(self, block:dict): assert len(block) == 2, "upsample length error" self.converted.append(block["conv.weight"]) self.converted.append(block["conv.bias"]) def append_up_down_alpha(self, block:dict, prefix:str): self.converted.append(block[f"{prefix}.lora_up.weight"]) self.converted.append(block[f"{prefix}.lora_down.weight"]) self.converted.append(block[f"{prefix}.alpha"]) def add_MLPTransformer(self, block:dict): assert len(block) == 18, "mlp transform length error" start_len = len(self.converted) self.append_up_down_alpha(block, "mlp_fc1") self.append_up_down_alpha(block, "mlp_fc2") self.append_up_down_alpha(block, "self_attn_q_proj") self.append_up_down_alpha(block, "self_attn_k_proj") self.append_up_down_alpha(block, "self_attn_v_proj") self.append_up_down_alpha(block, "self_attn_out_proj") assert len(self.converted) - start_len == len(block), "Not all tensors parsed in mlp transformer." def add_lora_unet_block(self, block:dict, count): assert len(block) == 36*count, "lora transform length error" start_len = len(self.converted) for i in range(count): prefix = f"attentions_{i}_" self.append_up_down_alpha(block, f"{prefix}proj_in") for j in range(2): tprefix = f"{prefix}transformer_blocks_0_attn{j+1}_" self.append_up_down_alpha(block, f"{tprefix}to_q") self.append_up_down_alpha(block, f"{tprefix}to_k") self.append_up_down_alpha(block, f"{tprefix}to_v") self.append_up_down_alpha(block, f"{tprefix}to_out_0") self.append_up_down_alpha(block, f"{prefix}transformer_blocks_0_ff_net_0_proj") self.append_up_down_alpha(block, f"{prefix}transformer_blocks_0_ff_net_2") self.append_up_down_alpha(block, f"{prefix}proj_out") assert len(self.converted) - start_len == len(block), "Not all tensors parsed in mlp transformer." if __name__ == "__main__": if len(sys.argv) > 1: parser = argparse.ArgumentParser(description='Process a path.') parser.add_argument('path', type=str, help='The path to process') args = parser.parse_args() target = args.path convertSD = ConvertSafetensors() convertSD.parse(target)