comfyanonymous commited on
Commit
cc8fdd4
1 Parent(s): 5a31360

Create convert.py

Browse files
Files changed (1) hide show
  1. convert.py +47 -0
convert.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import safetensors.torch
3
+ from transformers import T5Tokenizer, T5EncoderModel
4
+
5
+ #https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/blob/main/t2i/model/pytorch_model_ema.pt
6
+ input_diffusion = "pytorch_model_ema.pt"
7
+
8
+ #https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2/tree/main/t2i/clip_text_encoder
9
+ input_bert = "./clip_text_encoder/pytorch_model.bin"
10
+
11
+ #https://huggingface.co/stabilityai/sdxl-vae/blob/main/sdxl_vae.safetensors
12
+ # or
13
+ #https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors
14
+ input_vae = "sdxl_vae.safetensors"
15
+
16
+ output = "hunyuan_dit_1.2.safetensors"
17
+
18
+ bert_sd = torch.load(input_bert, weights_only=True)
19
+
20
+ mt5 = T5EncoderModel.from_pretrained("google/mt5-xl")
21
+ tokenizer = T5Tokenizer.from_pretrained("google/mt5-xl")
22
+
23
+ sp_model = torch.ByteTensor(list(tokenizer.sp_model.serialized_model_proto()))
24
+ t5_sd = mt5.state_dict()
25
+
26
+ out_sd = {}
27
+
28
+ out_sd["text_encoders.mt5xl.spiece_model"] = sp_model
29
+
30
+ for k in t5_sd:
31
+ out_sd["text_encoders.mt5xl.transformer.{}".format(k)] = t5_sd[k].half()
32
+
33
+ for k in bert_sd:
34
+ if not k.startswith("visual."):
35
+ out_sd["text_encoders.hydit_clip.transformer.{}".format(k)] = bert_sd[k].half()
36
+
37
+ hydit = torch.load(input_diffusion, weights_only=True)
38
+ for k in hydit:
39
+ out_sd["model.{}".format(k)] = hydit[k].half()
40
+
41
+
42
+ vae_sd = safetensors.torch.load_file(input_vae)
43
+
44
+ for k in vae_sd:
45
+ out_sd["vae.{}".format(k)] = vae_sd[k].half()
46
+
47
+ safetensors.torch.save_file(out_sd, output)