comdoleger commited on
Commit
668242b
·
verified ·
1 Parent(s): 64e55ce

Upload scripts/convert_lora_to_peft_format.py with huggingface_hub

Browse files
scripts/convert_lora_to_peft_format.py ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # currently only works with flux as support is not quite there yet
2
+
3
+ import argparse
4
+ import os.path
5
+ from collections import OrderedDict
6
+
7
+ parser = argparse.ArgumentParser()
8
+ parser.add_argument(
9
+ 'input_path',
10
+ type=str,
11
+ help='Path to original sdxl model'
12
+ )
13
+ parser.add_argument(
14
+ 'output_path',
15
+ type=str,
16
+ help='output path'
17
+ )
18
+ args = parser.parse_args()
19
+ args.input_path = os.path.abspath(args.input_path)
20
+ args.output_path = os.path.abspath(args.output_path)
21
+
22
+ from safetensors.torch import load_file, save_file
23
+
24
+ meta = OrderedDict()
25
+ meta['format'] = 'pt'
26
+
27
+ state_dict = load_file(args.input_path)
28
+
29
+ # peft doesnt have an alpha so we need to scale the weights
30
+ alpha_keys = [
31
+ 'lora_transformer_single_transformer_blocks_0_attn_to_q.alpha' # flux
32
+ ]
33
+
34
+ # keys where the rank is in the first dimension
35
+ rank_idx0_keys = [
36
+ 'lora_transformer_single_transformer_blocks_0_attn_to_q.lora_down.weight'
37
+ # 'transformer.single_transformer_blocks.0.attn.to_q.lora_A.weight'
38
+ ]
39
+
40
+ alpha = None
41
+ rank = None
42
+
43
+ for key in rank_idx0_keys:
44
+ if key in state_dict:
45
+ rank = int(state_dict[key].shape[0])
46
+ break
47
+
48
+ if rank is None:
49
+ raise ValueError(f'Could not find rank in state dict')
50
+
51
+ for key in alpha_keys:
52
+ if key in state_dict:
53
+ alpha = int(state_dict[key])
54
+ break
55
+
56
+ if alpha is None:
57
+ # set to rank if not found
58
+ alpha = rank
59
+
60
+
61
+ up_multiplier = alpha / rank
62
+
63
+ new_state_dict = {}
64
+
65
+ for key, value in state_dict.items():
66
+ if key.endswith('.alpha'):
67
+ continue
68
+
69
+ orig_dtype = value.dtype
70
+
71
+ new_val = value.float() * up_multiplier
72
+
73
+ new_key = key
74
+ new_key = new_key.replace('lora_transformer_', 'transformer.')
75
+ for i in range(100):
76
+ new_key = new_key.replace(f'transformer_blocks_{i}_', f'transformer_blocks.{i}.')
77
+ new_key = new_key.replace('lora_down', 'lora_A')
78
+ new_key = new_key.replace('lora_up', 'lora_B')
79
+ new_key = new_key.replace('_lora', '.lora')
80
+ new_key = new_key.replace('attn_', 'attn.')
81
+ new_key = new_key.replace('ff_', 'ff.')
82
+ new_key = new_key.replace('context_net_', 'context.net.')
83
+ new_key = new_key.replace('0_proj', '0.proj')
84
+ new_key = new_key.replace('norm_linear', 'norm.linear')
85
+ new_key = new_key.replace('norm_out_linear', 'norm_out.linear')
86
+ new_key = new_key.replace('to_out_', 'to_out.')
87
+
88
+ new_state_dict[new_key] = new_val.to(orig_dtype)
89
+
90
+ save_file(new_state_dict, args.output_path, meta)
91
+ print(f'Saved to {args.output_path}')