Function main

scripts/extract_lora_from_model.py:114–146

(args)

@torch.no_grad()
def main(args):
    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=args.base_subfolder, 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__":