Upload create_lora.py with huggingface_hub

create_lora.py

@@ -0,0 +1,80 @@
+"""
+Adapted from 
+https://github.com/Stability-AI/stability-ComfyUI-nodes/blob/001154622564b17223ce0191803c5fff7b87146c/control_lora_create.py
+"""
+
+from diffusers import CogVideoXTransformer3DModel
+from tqdm.auto import tqdm
+from safetensors.torch import save_file
+import torch
+
+RANK = 64
+CLAMP_QUANTILE = 0.99
+
+
+# Comes from
+# https://github.com/Stability-AI/stability-ComfyUI-nodes/blob/001154622564b17223ce0191803c5fff7b87146c/control_lora_create.py#L9
+def extract_lora(diff, rank):
+    if torch.cuda.is_available():
+        diff = diff.to("cuda")
+    
+    is_conv2d = (len(diff.shape) == 4)
+    kernel_size = None if not is_conv2d else diff.size()[2:4]
+    is_conv2d_3x3 = is_conv2d and kernel_size != (1, 1)
+    out_dim, in_dim = diff.size()[0:2]
+    rank = min(rank, in_dim, out_dim)
+
+    if is_conv2d:
+        if is_conv2d_3x3:
+            diff = diff.flatten(start_dim=1)
+        else:
+            diff = diff.squeeze()
+
+    U, S, Vh = torch.linalg.svd(diff.float())
+    U = U[:, :rank]
+    S = S[:rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:rank, :]
+
+    dist = torch.cat([U.flatten(), Vh.flatten()])
+    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+    low_val = -hi_val
+
+    U = U.clamp(low_val, hi_val)
+    Vh = Vh.clamp(low_val, hi_val)
+    if is_conv2d:
+        U = U.reshape(out_dim, rank, 1, 1)
+        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
+    return (U.cpu(), Vh.cpu())
+
+
+transformer_finetuned = CogVideoXTransformer3DModel.from_pretrained(
+    "cogvideox-cakeify", subfolder="transformer", torch_dtype=torch.bfloat16
+)
+state_dict_ft = transformer_finetuned.state_dict()
+
+transformer = CogVideoXTransformer3DModel.from_pretrained(
+    "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
+)
+state_dict = transformer.state_dict()
+output_dict = {}
+
+for k in tqdm(state_dict, desc="Extracting LoRA..."):
+    original_param = state_dict[k]
+    finetuned_param = state_dict_ft[k]
+    if len(original_param.shape) >= 2:
+        diff = finetuned_param.float() - original_param.float()
+        out = extract_lora(diff, RANK)
+        name = k
+        
+        if name.endswith(".weight"):
+            name = name[:-len(".weight")]
+        down_key = "{}.lora_A.weight".format(name)
+        up_key = "{}.lora_B.weight".format(name)
+        
+        output_dict[up_key] = out[0].contiguous().to(finetuned_param.dtype)
+        output_dict[down_key] = out[1].contiguous().to(finetuned_param.dtype)
+
+output_dict = {f"transformer.{k}": v for k, v in output_dict.items()}
+save_file(output_dict, "extracted_cakeify_lora_64.safetensors")
+print(f"LoRA saved and it contains {len(output_dict)} keys.")