diffusers/scripts/extract_lora_from_model.py at main · toshas/diffusers · GitHub

History

151 lines (121 loc) · 4.76 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

"""

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(

"--base_subfolder",

default="transformer",

type=str,

help="subfolder to load the base checkpoint from if any.",

)

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(

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

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):

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__":

args = parse_args()

main(args)

diffusers/scripts/extract_lora_from_model.py at main · toshas/diffusers · GitHub

Expand file tree

Breadcrumbs

extract_lora_from_model.py

Latest commit

History

Breadcrumbs

extract_lora_from_model.py

File metadata and controls