merge_non_lora_weight.py 730 Bytes
Newer Older
Zhekai Zhang's avatar
Zhekai Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from safetensors.torch import safe_open, save_file


def main():
    input_path1 = "app/i2i/pretrained/converted/sketch.safetensors"
    input_path2 = "app/i2i/pretrained/original/flux-lora-sketch2image-bf16.safetensors"

    sd1 = {}
    with safe_open(input_path1, framework="pt") as f:
        for k in f.keys():
            sd1[k] = f.get_tensor(k)

    sd2 = {}
    with safe_open(input_path2, framework="pt") as f:
        for k in f.keys():
            sd2[k] = f.get_tensor(k)

    for k in sd1.keys():
        if "lora" not in k:
            print(k)
            sd2[k.replace("transformer.", "")] = sd1[k]

    save_file(sd2, "svdq-flux.1-pix2pix-turbo-sketch2image.safetensors")


if __name__ == "__main__":
    main()