|
""" |
|
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 |
|
|
|
|
|
|
|
|
|
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.") |
