From c9ef083c2ed22f6a9b7117199ef6b1fe736dd9a8 Mon Sep 17 00:00:00 2001 From: haofanwang Date: Sat, 18 Feb 2023 01:36:04 +0800 Subject: [PATCH 1/4] add lora convertor --- .../convert_lora_safetensor_to_diffusers.py | 133 ++++++++++++++++++ 1 file changed, 133 insertions(+) create mode 100644 scripts/convert_lora_safetensor_to_diffusers.py diff --git a/scripts/convert_lora_safetensor_to_diffusers.py b/scripts/convert_lora_safetensor_to_diffusers.py new file mode 100644 index 000000000000..ae89d3d75c32 --- /dev/null +++ b/scripts/convert_lora_safetensor_to_diffusers.py @@ -0,0 +1,133 @@ +# coding=utf-8 +# Copyright 2023, Haofan Wang, Qixun Wang, All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" Conversion script for the LoRA's safetensors checkpoints. """ + +import torch +from safetensors.torch import load_file +from diffusers import StableDiffusionPipeline +from diffusers import DPMSolverMultistepScheduler + +import argparse + + +def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT_ENCODER, alpha): + + # load base model + pipeline = StableDiffusionPipeline.from_pretrained(base_model_path, torch_dtype=torch.float32) + + # load LoRA weight from .safetensors + state_dict = load_file(checkpoint_path) + + visited = [] + + # directly update weight in diffusers model + for key in state_dict: + + # it is suggested to print out the key, it usually will be something like below + # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" + + # as we have set the alpha beforehand, so just skip + if '.alpha' in key or key in visited: + continue + + if 'text' in key: + layer_infos = key.split('.')[0].split(LORA_PREFIX_TEXT_ENCODER+'_')[-1].split('_') + curr_layer = pipeline.text_encoder + else: + layer_infos = key.split('.')[0].split(LORA_PREFIX_UNET+'_')[-1].split('_') + curr_layer = pipeline.unet + + # find the target layer + temp_name = layer_infos.pop(0) + while len(layer_infos) > -1: + try: + curr_layer = curr_layer.__getattr__(temp_name) + if len(layer_infos) > 0: + temp_name = layer_infos.pop(0) + elif len(layer_infos) == 0: + break + except Exception: + if len(temp_name) > 0: + temp_name += '_'+layer_infos.pop(0) + else: + temp_name = layer_infos.pop(0) + + pair_keys = [] + if 'lora_down' in key: + pair_keys.append(key.replace('lora_down', 'lora_up')) + pair_keys.append(key) + else: + pair_keys.append(key) + pair_keys.append(key.replace('lora_up', 'lora_down')) + + # update weight + if len(state_dict[pair_keys[0]].shape) == 4: + weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) + weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) + curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) + else: + weight_up = state_dict[pair_keys[0]].to(torch.float32) + weight_down = state_dict[pair_keys[1]].to(torch.float32) + curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down) + + # update visited list + for item in pair_keys: + visited.append(item) + + return pipeline + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--base_model_path", default=None, type=str, required=True, help="Path to the base model in diffusers format." + ) + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--dump_path", default=None, type=str, required=True, help="Path to the output model." + ) + parser.add_argument( + "--lora_prefix_unet", default="lora_unet", type=str, help="The prefix of UNet weight in safetensors" + ) + parser.add_argument( + "--lora_prefix_text_encoder", default="lora_te", type=str, help="The prefix of text encoder weight in safetensors" + ) + parser.add_argument( + "--alpha", default=0.75, type=int, help="The merging ratio in W = W0 + alpha * deltaW" + ) + parser.add_argument( + "--to_safetensors", action="store_true", help="Whether to store pipeline in safetensors format or not." + ) + parser.add_argument( + "--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)" + ) + + args = parser.parse_args() + + base_model_path = args.base_model_path + checkpoint_path = args.checkpoint_path + dump_path = args.dump_path + lora_prefix_unet = args.lora_prefix_unet + lora_prefix_text_encoder = args.lora_prefix_text_encoder + alpha = args.alpha + + pipe = convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha) + + pipe = pipe.to(args.device) + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) \ No newline at end of file From 653448311cbe93f95fd58a9d6d0bbb2f3d5dba50 Mon Sep 17 00:00:00 2001 From: Haofan Wang Date: Sat, 18 Feb 2023 01:53:01 +0800 Subject: [PATCH 2/4] Update convert_lora_safetensor_to_diffusers.py --- .../convert_lora_safetensor_to_diffusers.py | 55 +++++++++---------- 1 file changed, 26 insertions(+), 29 deletions(-) diff --git a/scripts/convert_lora_safetensor_to_diffusers.py b/scripts/convert_lora_safetensor_to_diffusers.py index ae89d3d75c32..705ba15f17e1 100644 --- a/scripts/convert_lora_safetensor_to_diffusers.py +++ b/scripts/convert_lora_safetensor_to_diffusers.py @@ -15,12 +15,12 @@ """ Conversion script for the LoRA's safetensors checkpoints. """ +import argparse + import torch from safetensors.torch import load_file -from diffusers import StableDiffusionPipeline -from diffusers import DPMSolverMultistepScheduler -import argparse +from diffusers import StableDiffusionPipeline def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT_ENCODER, alpha): @@ -35,19 +35,19 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT # directly update weight in diffusers model for key in state_dict: - + # it is suggested to print out the key, it usually will be something like below # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" - + # as we have set the alpha beforehand, so just skip - if '.alpha' in key or key in visited: + if ".alpha" in key or key in visited: continue - - if 'text' in key: - layer_infos = key.split('.')[0].split(LORA_PREFIX_TEXT_ENCODER+'_')[-1].split('_') + + if "text" in key: + layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") curr_layer = pipeline.text_encoder else: - layer_infos = key.split('.')[0].split(LORA_PREFIX_UNET+'_')[-1].split('_') + layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_") curr_layer = pipeline.unet # find the target layer @@ -61,18 +61,18 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT break except Exception: if len(temp_name) > 0: - temp_name += '_'+layer_infos.pop(0) + temp_name += "_" + layer_infos.pop(0) else: temp_name = layer_infos.pop(0) - + pair_keys = [] - if 'lora_down' in key: - pair_keys.append(key.replace('lora_down', 'lora_up')) + if "lora_down" in key: + pair_keys.append(key.replace("lora_down", "lora_up")) pair_keys.append(key) else: pair_keys.append(key) - pair_keys.append(key.replace('lora_up', 'lora_down')) - + pair_keys.append(key.replace("lora_up", "lora_down")) + # update weight if len(state_dict[pair_keys[0]].shape) == 4: weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) @@ -82,7 +82,7 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT weight_up = state_dict[pair_keys[0]].to(torch.float32) weight_down = state_dict[pair_keys[1]].to(torch.float32) curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down) - + # update visited list for item in pair_keys: visited.append(item) @@ -99,25 +99,22 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT parser.add_argument( "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." ) - parser.add_argument( - "--dump_path", default=None, type=str, required=True, help="Path to the output model." - ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") parser.add_argument( "--lora_prefix_unet", default="lora_unet", type=str, help="The prefix of UNet weight in safetensors" ) parser.add_argument( - "--lora_prefix_text_encoder", default="lora_te", type=str, help="The prefix of text encoder weight in safetensors" - ) - parser.add_argument( - "--alpha", default=0.75, type=int, help="The merging ratio in W = W0 + alpha * deltaW" + "--lora_prefix_text_encoder", + default="lora_te", + type=str, + help="The prefix of text encoder weight in safetensors", ) + parser.add_argument("--alpha", default=0.75, type=int, help="The merging ratio in W = W0 + alpha * deltaW") parser.add_argument( "--to_safetensors", action="store_true", help="Whether to store pipeline in safetensors format or not." ) - parser.add_argument( - "--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)" - ) - + parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") + args = parser.parse_args() base_model_path = args.base_model_path @@ -130,4 +127,4 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT pipe = convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha) pipe = pipe.to(args.device) - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) \ No newline at end of file + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) From 6829a97a212cfa462b9a6ff722f68ab77b8fbf4f Mon Sep 17 00:00:00 2001 From: Haofan Wang Date: Sat, 18 Feb 2023 01:59:20 +0800 Subject: [PATCH 3/4] Update README.md --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index fc384c9f8fb2..3f5a42538dd9 100644 --- a/README.md +++ b/README.md @@ -467,12 +467,12 @@ image.save("ddpm_generated_image.png") - [Unconditional Diffusion with continuous scheduler](https://huggingface.co/google/ncsnpp-ffhq-1024) **Other Image Notebooks**: -* [image-to-image generation with Stable Diffusion](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb) ![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg), -* [tweak images via repeated Stable Diffusion seeds](https://colab.research.google.com/github/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb) ![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg), +* [image-to-image generation with Stable Diffusion](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb) ![Open In Colab](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb), +* [tweak images via repeated Stable Diffusion seeds](https://colab.research.google.com/github/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb) ![Open In Colab](https://colab.research.google.com/github/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb), **Diffusers for Other Modalities**: -* [Molecule conformation generation](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/geodiff_molecule_conformation.ipynb) ![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg), -* [Model-based reinforcement learning](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/reinforcement_learning_with_diffusers.ipynb) ![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg), +* [Molecule conformation generation](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/geodiff_molecule_conformation.ipynb) ![Open In Colab](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/geodiff_molecule_conformation.ipynb), +* [Model-based reinforcement learning](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/reinforcement_learning_with_diffusers.ipynb) ![Open In Colab](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/reinforcement_learning_with_diffusers.ipynb), ### Web Demos If you just want to play around with some web demos, you can try out the following 🚀 Spaces: From 75501a37157da4968291a7929bb8cb374eb57f22 Mon Sep 17 00:00:00 2001 From: Haofan Wang Date: Mon, 20 Feb 2023 15:52:08 +0800 Subject: [PATCH 4/4] Update convert_lora_safetensor_to_diffusers.py --- scripts/convert_lora_safetensor_to_diffusers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/convert_lora_safetensor_to_diffusers.py b/scripts/convert_lora_safetensor_to_diffusers.py index 705ba15f17e1..ab1023bcca70 100644 --- a/scripts/convert_lora_safetensor_to_diffusers.py +++ b/scripts/convert_lora_safetensor_to_diffusers.py @@ -109,7 +109,7 @@ def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT type=str, help="The prefix of text encoder weight in safetensors", ) - parser.add_argument("--alpha", default=0.75, type=int, help="The merging ratio in W = W0 + alpha * deltaW") + parser.add_argument("--alpha", default=0.75, type=float, help="The merging ratio in W = W0 + alpha * deltaW") parser.add_argument( "--to_safetensors", action="store_true", help="Whether to store pipeline in safetensors format or not." )