File size: 2,645 Bytes
ba0bbec |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 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.") |