Added safetensor converter.

ConvertSafetensors.py

@@ -0,0 +1,406 @@
+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)