Views
No views yet
| Precision/Rank | rank=8 | rank=16 | rank=32 | rank=64 | rank=128 | rank=256 |
|---|---|---|---|---|---|---|
| FP16 | rank8-fp16 | rank16-fp16 | rank32-fp16 | rank64-fp16 | rank128-fp16 | rank256-fp16 |
| BF16 | rank8-bf16 | rank16-bf16 | rank32-bf16 | rank64-bf16 | rank128-bf16 | rank256-bf16 |
| FP32 | rank8-fp32 | rank16-fp32 | rank32-fp32 | rank64-fp32 | rank128-fp32 | rank256-fp32 |
1import torch
2from safetensors import safe_open
3from safetensors.torch import save_file
4from tqdm import tqdm
5
6
7BASE_MODEL_SAFETENSORS = "./flux-2-klein-base-4b.safetensors"
8TUNED_MODEL_SAFETENSORS = "./flux-2-klein-4b.safetensors"
9OUTPUT_DIR = "./extracted_loras"
10
11RANKS = [8, 16, 32, 64, 128, 256]
12CLAMP_QUANTILE = 0.99
13SAVE_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
14
15DTYPE_STR_MAP = {
16 torch.float16: "fp16",
17 torch.bfloat16: "bf16",
18 torch.float32: "fp32",
19}
20
21
22with safe_open(BASE_MODEL_SAFETENSORS, framework="pt", device="cuda") as f_org, \
23 safe_open(TUNED_MODEL_SAFETENSORS, framework="pt", device="cuda") as f_tuned:
24
25 org_keys = set(f_org.keys())
26 tuned_keys = set(f_tuned.keys())
27 shared_keys = sorted(org_keys & tuned_keys)
28
29 candidate_keys = [k for k in shared_keys if k.endswith(".weight")]
30 print(f"Found {len(candidate_keys)} shared weight keys")
31
32 for target_rank in RANKS:
33 print(f"\nProcessing Rank: {target_rank}")
34 lora_sds = {dtype: {} for dtype in SAVE_DTYPES}
35
36 for key in tqdm(candidate_keys):
37 v_org = f_org.get_tensor(key)
38 v_tuned = f_tuned.get_tensor(key)
39
40 if v_org.shape != v_tuned.shape or v_org.ndim != 2:
41 del v_org, v_tuned
42 continue
43
44 diff = (v_tuned.to(torch.float32) - v_org.to(torch.float32))
45 del v_org, v_tuned
46
47 out_dim, in_dim = diff.shape
48 rank = min(target_rank, in_dim, out_dim)
49
50 U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
51 del diff
52
53 U = U[:, :rank]
54 S = S[:rank]
55 Vh = Vh[:rank, :]
56
57 U = U @ torch.diag(S)
58 del S
59
60 dist = torch.cat([U.flatten(), Vh.flatten()])
61 hi = torch.quantile(dist, CLAMP_QUANTILE)
62 del dist
63 U = U.clamp(-hi, hi)
64 Vh = Vh.clamp(-hi, hi)
65
66 lora_name = "lora_unet_" + key.replace(".weight", "").replace(".", "_")
67
68 for save_dtype in SAVE_DTYPES:
69 lora_sds[save_dtype][lora_name + ".lora_up.weight"] = U.to(save_dtype).contiguous()
70 lora_sds[save_dtype][lora_name + ".lora_down.weight"] = Vh.to(save_dtype).contiguous()
71 lora_sds[save_dtype][lora_name + ".alpha"] = torch.tensor(rank, dtype=save_dtype)
72
73 del U, Vh
74
75 metadata = {}
76 for save_dtype in SAVE_DTYPES:
77 dtype_name = DTYPE_STR_MAP.get(save_dtype, str(save_dtype).split(".")[-1])
78 output_path = f"{OUTPUT_DIR}/flux-2-klein-4b-lora-rank{target_rank}-{CLAMP_QUANTILE}-{dtype_name}.safetensors"
79 save_file(lora_sds[save_dtype], output_path, metadata=metadata)
80 print(f"Saved extracted LoRA to {output_path}")
81