From c68e5bdbbc8a71e9258ebd79e771e36016fe10a0 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 23 Jan 2025 08:20:54 +0530 Subject: [PATCH 1/2] feat: add a lora extraction script. --- scripts/extract_lora_from_model.py | 138 +++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 scripts/extract_lora_from_model.py diff --git a/scripts/extract_lora_from_model.py b/scripts/extract_lora_from_model.py new file mode 100644 index 000000000000..7cb8a5c8c61d --- /dev/null +++ b/scripts/extract_lora_from_model.py @@ -0,0 +1,138 @@ +""" +This script demonstrates how to extract a LoRA checkpoint from a fully finetuned model with the CogVideoX model. + +To make it work for other models: + +* Change the model class. Here we use `CogVideoXTransformer3DModel`. For Flux, it would be `FluxTransformer2DModel`, +for example. (TODO: more reason to add `AutoModel`). +* Spply path to the base checkpoint via `base_ckpt_path`. +* Supply path to the fully fine-tuned checkpoint via `--finetune_ckpt_path`. +* Change the `--rank` as needed. + +Example usage: + +```bash +python extract_lora_from_model.py \ + --base_ckpt_path=THUDM/CogVideoX-5b \ + --finetune_ckpt_path=finetrainers/cakeify-v0 \ + --lora_out_path=cakeify_lora.safetensors +``` + +Script is adapted from +https://github.com/Stability-AI/stability-ComfyUI-nodes/blob/001154622564b17223ce0191803c5fff7b87146c/control_lora_create.py +""" + +import argparse + +import torch +from safetensors.torch import save_file +from tqdm.auto import tqdm + +from diffusers import CogVideoXTransformer3DModel + + +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): + # Important to use CUDA otherwise, very slow! + 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()) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--base_ckpt_path", + default=None, + type=str, + required=True, + help="Base checkpoint path from which the model was finetuned. Can be a model ID on the Hub.", + ) + parser.add_argument( + "--finetune_ckpt_path", + default=None, + type=str, + required=True, + help="Fully fine-tuned checkpoint path. Can be a model ID on the Hub.", + ) + parser.add_argument("--rank", default=64, type=int) + parser.add_argument("--lora_out_path", default=None, type=str, required=True) + args = parser.parse_args() + + if not args.lora_out_path.endswith(".safetensors"): + raise ValueError("`lora_out_path` must end with `.safetensors`.") + + return args + + +@torch.no_grad() +def main(args): + # Fully fine-tuned checkpoints usually don't have any other components. So, we + # don't need the `subfolder`. You can add that if needed. + model_finetuned = CogVideoXTransformer3DModel.from_pretrained(args.finetune_ckpt_path, torch_dtype=torch.bfloat16) + state_dict_ft = model_finetuned.state_dict() + + # Change the `subfolder` as needed. + base_model = CogVideoXTransformer3DModel.from_pretrained( + args.base_ckpt_path, subfolder="transformer", torch_dtype=torch.bfloat16 + ) + state_dict = base_model.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) + + prefix = "transformer" if "transformer" in base_model.__class__.__name__.lower() else "unet" + output_dict = {f"{prefix}.{k}": v for k, v in output_dict.items()} + save_file(output_dict, args.lora_out_path) + print(f"LoRA saved and it contains {len(output_dict)} keys.") + + +if __name__ == "__main__": + main() From 1cad3c7b1a930fd796be338b952c741b9f7bc82a Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 23 Jan 2025 08:51:54 +0530 Subject: [PATCH 2/2] updates --- scripts/extract_lora_from_model.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/scripts/extract_lora_from_model.py b/scripts/extract_lora_from_model.py index 7cb8a5c8c61d..0e01ddea47f9 100644 --- a/scripts/extract_lora_from_model.py +++ b/scripts/extract_lora_from_model.py @@ -81,6 +81,12 @@ def parse_args(): required=True, help="Base checkpoint path from which the model was finetuned. Can be a model ID on the Hub.", ) + parser.add_argument( + "--base_subfolder", + default="transformer", + type=str, + help="subfolder to load the base checkpoint from if any.", + ) parser.add_argument( "--finetune_ckpt_path", default=None, @@ -88,6 +94,12 @@ def parse_args(): required=True, help="Fully fine-tuned checkpoint path. Can be a model ID on the Hub.", ) + parser.add_argument( + "--finetune_subfolder", + default=None, + type=str, + help="subfolder to load the fulle finetuned checkpoint from if any.", + ) parser.add_argument("--rank", default=64, type=int) parser.add_argument("--lora_out_path", default=None, type=str, required=True) args = parser.parse_args() @@ -100,14 +112,14 @@ def parse_args(): @torch.no_grad() def main(args): - # Fully fine-tuned checkpoints usually don't have any other components. So, we - # don't need the `subfolder`. You can add that if needed. - model_finetuned = CogVideoXTransformer3DModel.from_pretrained(args.finetune_ckpt_path, torch_dtype=torch.bfloat16) + model_finetuned = CogVideoXTransformer3DModel.from_pretrained( + args.finetune_ckpt_path, subfolder=args.finetune_subfolder, torch_dtype=torch.bfloat16 + ) state_dict_ft = model_finetuned.state_dict() # Change the `subfolder` as needed. base_model = CogVideoXTransformer3DModel.from_pretrained( - args.base_ckpt_path, subfolder="transformer", torch_dtype=torch.bfloat16 + args.base_ckpt_path, subfolder=args.base_subfolder, torch_dtype=torch.bfloat16 ) state_dict = base_model.state_dict() output_dict = {} @@ -135,4 +147,5 @@ def main(args): if __name__ == "__main__": - main() + args = parse_args() + main(args)