FLUX.2-klein-4B-Distill-LoRA-Extract

PJMixers-Images/FLUX.2-klein-4B-Distill-LoRA-Extract

Extracted LoRAs of the distilled FLUX.2-klein-4B model for use with the FLUX.2-klein-base-4B model.

Downloads

Conversion Script

import torch
from safetensors import safe_open
from safetensors.torch import save_file
from tqdm import tqdm


BASE_MODEL_SAFETENSORS = "./flux-2-klein-base-4b.safetensors"
TUNED_MODEL_SAFETENSORS = "./flux-2-klein-4b.safetensors"
OUTPUT_DIR = "./extracted_loras"

RANKS = [8, 16, 32, 64, 128, 256]
CLAMP_QUANTILE = 0.99
SAVE_DTYPES = [torch.float16, torch.bfloat16, torch.float32]

DTYPE_STR_MAP = {
    torch.float16: "fp16",
    torch.bfloat16: "bf16",
    torch.float32: "fp32",
}


with safe_open(BASE_MODEL_SAFETENSORS, framework="pt", device="cuda") as f_org, \
     safe_open(TUNED_MODEL_SAFETENSORS, framework="pt", device="cuda") as f_tuned:

    org_keys = set(f_org.keys())
    tuned_keys = set(f_tuned.keys())
    shared_keys = sorted(org_keys & tuned_keys)

    candidate_keys = [k for k in shared_keys if k.endswith(".weight")]
    print(f"Found {len(candidate_keys)} shared weight keys")

    for target_rank in RANKS:
        print(f"\nProcessing Rank: {target_rank}")
        lora_sds = {dtype: {} for dtype in SAVE_DTYPES}

        for key in tqdm(candidate_keys):
            v_org = f_org.get_tensor(key)
            v_tuned = f_tuned.get_tensor(key)

            if v_org.shape != v_tuned.shape or v_org.ndim != 2:
                del v_org, v_tuned
                continue

            diff = (v_tuned.to(torch.float32) - v_org.to(torch.float32))
            del v_org, v_tuned

            out_dim, in_dim = diff.shape
            rank = min(target_rank, in_dim, out_dim)

            U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
            del diff

            U = U[:, :rank]
            S = S[:rank]
            Vh = Vh[:rank, :]

            U = U @ torch.diag(S)
            del S

            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi = torch.quantile(dist, CLAMP_QUANTILE)
            del dist
            U = U.clamp(-hi, hi)
            Vh = Vh.clamp(-hi, hi)

            lora_name = "lora_unet_" + key.replace(".weight", "").replace(".", "_")

            for save_dtype in SAVE_DTYPES:
                lora_sds[save_dtype][lora_name + ".lora_up.weight"] = U.to(save_dtype).contiguous()
                lora_sds[save_dtype][lora_name + ".lora_down.weight"] = Vh.to(save_dtype).contiguous()
                lora_sds[save_dtype][lora_name + ".alpha"] = torch.tensor(rank, dtype=save_dtype)

            del U, Vh

        metadata = {}
        for save_dtype in SAVE_DTYPES:
            dtype_name = DTYPE_STR_MAP.get(save_dtype, str(save_dtype).split(".")[-1])
            output_path = f"{OUTPUT_DIR}/flux-2-klein-4b-lora-rank{target_rank}-{CLAMP_QUANTILE}-{dtype_name}.safetensors"
            save_file(lora_sds[save_dtype], output_path, metadata=metadata)
            print(f"Saved extracted LoRA to {output_path}")