Stable-Diffusion-Vertex/Diffusers/train.py

# Copyright 2023 Google LLC
#
# 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.

import subprocess
import os
import argparse
import re
import torch
from safetensors.torch import save_file

def bin_to_safetensors(output_path):
    newDict = dict();
    checkpoint = torch.load(output_path + '/pytorch_lora_weights.bin');
    for idx, key in enumerate(checkpoint):
      newKey = re.sub('\.processor\.', '_', key);
      newKey = re.sub('mid_block\.', 'mid_block_', newKey);
      newKey = re.sub('_lora.up.', '.lora_up.', newKey);
      newKey = re.sub('_lora.down.', '.lora_down.', newKey);
      newKey = re.sub('\.(\d+)\.', '_\\1_', newKey);
      newKey = re.sub('to_out', 'to_out_0', newKey);
      newKey = 'lora_unet_'+newKey;

      newDict[newKey] = checkpoint[key];

    newLoraName = 'pytorch_lora_weights.safetensors';
    print("Saving " + newLoraName);
    save_file(newDict, output_path + '/' + newLoraName);

def main(args):
    
    subprocess.run("accelerate config default", shell=True)
    subprocess.run("cat /root/.cache/huggingface/accelerate/default_config.yaml", shell=True)

    METHOD = args.method
    MODEL_NAME= args.model_name #"runwayml/stable-diffusion-v1-5"
    INSTANCE_DIR= args.input_storage
    OUTPUT_DIR= args.output_storage
    PROMPT = args.prompt
    CLASS_PROMPT = args.class_prompt
    NUM_CLASS_IMAGES = int(args.num_class_images)
    STEPS = int(args.max_train_steps)
    TEXT_ENCODER = bool(args.text_encoder)
    SET_GRADS_TO_NONE = bool(args.set_grads_to_none)
    
    RESOLUTION = int(args.resolution)
    BATCH_SIZE = int(args.batch_size)
    USE_8BIT = bool(args.use_8bit)
    LR = float(args.lr)
    GRADIENT_ACCU_STEPS = int(args.gradient_accumulation_steps)
    NUM_VALID_IMG = int(args.num_validation_images)
    VALID_PRMOP = args.validation_prompt
    

    # Note the constraint: raise error: (args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1)
    
    if METHOD == "diffuser_dreambooth":
        os.chdir("/root/diffusers/examples/dreambooth")
        # for complex commands, with many args, use string + `shell=True`:
        cmd_str = (f'accelerate launch train_dreambooth.py '
                   f'--pretrained_model_name_or_path="{MODEL_NAME}" '
                   f'--instance_data_dir="{INSTANCE_DIR}" '
                   f'--output_dir="{OUTPUT_DIR}" '
                   f'--instance_prompt="{PROMPT}" '
                   f'--class_data_dir="{OUTPUT_DIR}/class_data" '
                   f'--with_prior_preservation --prior_loss_weight=1.0 '
                   f'--class_prompt="{CLASS_PROMPT}" '
                   f'--resolution={RESOLUTION} '
                   f'--train_batch_size={BATCH_SIZE} '                   
                   f'--gradient_checkpointing '
                   f'--gradient_accumulation_steps={GRADIENT_ACCU_STEPS} '
                   f'--mixed_precision="fp16" '
                   f'--learning_rate={LR} '
                   f'--lr_scheduler="constant" '
                   f'--lr_warmup_steps=0 '
                   f'--num_class_images={NUM_CLASS_IMAGES} '
                   f'--enable_xformers_memory_efficient_attention '
                   f'--max_train_steps={STEPS}')
        
        if TEXT_ENCODER == True:
            cmd_str += f' --train_text_encoder'
        if SET_GRADS_TO_NONE == True:
            cmd_str += f' --set_grads_to_none'
        if USE_8BIT == True:
            cmd_str += f' --use_8bit_adam'
        
    elif METHOD == "diffuser_dreambooth_lora":
        os.chdir("/root/diffusers/examples/dreambooth")
        # for complex commands, with many args, use string + `shell=True`:
        cmd_str = (f'accelerate launch train_dreambooth_lora.py '
                   f'--pretrained_model_name_or_path="{MODEL_NAME}" '
                   f'--instance_data_dir="{INSTANCE_DIR}" '
                   f'--output_dir="{OUTPUT_DIR}" '
                   f'--instance_prompt="{PROMPT}" '
                   f'--resolution={RESOLUTION} '
                   f'--train_batch_size={BATCH_SIZE} '
                   f'--mixed_precision="fp16" '
                   f'--gradient_accumulation_steps={GRADIENT_ACCU_STEPS} '
                   f'--learning_rate={LR} '
                   f'--lr_scheduler="constant" '
                   f'--lr_warmup_steps=0 '
                   f'--max_train_steps={STEPS}')
    
        if USE_8BIT == True:
            cmd_str += f' --use_8bit_adam'

    elif METHOD == "diffuser_text_to_image":
        os.chdir("/root/diffusers/examples/text_to_image")
        cmd_str = (f'accelerate launch --mixed_precision="fp16" train_text_to_image.py '
                   f'--pretrained_model_name_or_path="{MODEL_NAME}" '
                   f'--train_data_dir="{INSTANCE_DIR}" '
                   f'--use_ema '
                   f'--resolution={RESOLUTION} --center_crop --random_flip '
                   f'--mixed_precision="fp16" '
                   f'--train_batch_size={BATCH_SIZE} '
                   f'--gradient_accumulation_steps={GRADIENT_ACCU_STEPS} '
                   f'--gradient_checkpointing '
                   f'--max_train_steps={STEPS} '
                   f'--learning_rate={LR} '
                   f'--max_grad_norm=1 '
                   f'--lr_scheduler="constant" '
                   f'--lr_warmup_steps=0 '
                   f'--enable_xformers_memory_efficient_attention '
                   f'--output_dir="{OUTPUT_DIR}"')

        if USE_8BIT == True:
            cmd_str += f' --use_8bit_adam'

    elif METHOD == "diffuser_text_to_image_lora":
        os.chdir("/root/diffusers/examples/text_to_image")
        cmd_str = (f'accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py '
                   f'--pretrained_model_name_or_path="{MODEL_NAME}" '
                   f'--train_data_dir="{INSTANCE_DIR}" '
                   f'--resolution={RESOLUTION} --center_crop --random_flip '
                   f'--mixed_precision="fp16" '
                   f'--train_batch_size={BATCH_SIZE} '
                   f'--gradient_accumulation_steps={GRADIENT_ACCU_STEPS} '
                   f'--gradient_checkpointing '
                   f'--max_train_steps={STEPS} '
                   f'--learning_rate={LR} '
                   f'--max_grad_norm=1 '
                   f'--lr_scheduler="constant" --lr_warmup_steps=0 '
                   f'--seed=42 '
                   f'--num_validation_images={NUM_VALID_IMG} '
                   f'--validation_prompt="{VALID_PRMOP}" '
                   f'--output_dir="{OUTPUT_DIR}"')
        
    elif METHOD == "peft_lora":
        os.chdir("/root/peft/examples/lora_dreambooth")
                
        #create class data diretory
        if not os.path.exists(OUTPUT_DIR):
            os.mkdir(OUTPUT_DIR)
            print(f"{OUTPUT_DIR} has been created.")
        else:
            print(f"{OUTPUT_DIR} already exists.")
            
        class_directory = OUTPUT_DIR + '/class_data'
        if not os.path.exists(class_directory):
            os.mkdir(class_directory)
            print(f"{class_directory} has been created.")
        else:
            print(f"{class_directory} already exists.")

        # for complex commands, with many args, use string + `shell=True`:
        cmd_str = (f'accelerate launch train_dreambooth.py '
                   f'--pretrained_model_name_or_path="{MODEL_NAME}" '
                   f'--instance_data_dir="{INSTANCE_DIR}" '
                   f'--output_dir="{OUTPUT_DIR}" '
                   f'--with_prior_preservation '
                   f'--prior_loss_weight=1 '
                   f'--num_class_images={NUM_CLASS_IMAGES} '
                   f'--class_prompt="{CLASS_PROMPT}" '
                   f'--class_data_dir="{OUTPUT_DIR}/class_data" '
                   f'--instance_prompt="{PROMPT}" '
                   f'--use_lora '
                   f'--lora_r=4 '
                   f'--lora_alpha=4 '
                   f'--lora_bias=none '
                   f'--lora_dropout=0.0 '
                   f'--lora_text_encoder_r=4 '
                   f'--lora_text_encoder_alpha=4 '
                   f'--lora_text_encoder_bias=none '
                   f'--lora_text_encoder_dropout=0.0 '
                   f'--gradient_checkpointing '
                   f'--resolution=512 '
                   f'--train_batch_size=1 '
                   f'--use_8bit_adam '
                   f'--mixed_precision="fp16" '
                   f'--gradient_accumulation_steps=1 '
                   f'--learning_rate=1e-4 '
                   f'--lr_scheduler="constant" '
                   f'--lr_warmup_steps=0 '
                   f'--enable_xformers_memory_efficient_attention '
                   f'--max_train_steps={STEPS}')
        if TEXT_ENCODER == True:
            cmd_str += f' --train_text_encoder '


    
    # start training
    subprocess.run(cmd_str, shell=True)

    # convert to safetensors
    if (METHOD == "diffuser_dreambooth_lora") or (METHOD == "diffuser_text_to_image_lora"):
        bin_to_safetensors(args.output_storage)

    if (METHOD == "diffuser_dreambooth") or (METHOD == "diffuser_text_to_image"):
        subprocess.run(f'python3 /root/diffusers/scripts/convert_diffusers_to_original_stable_diffusion.py --model_path {OUTPUT_DIR} --checkpoint_path {OUTPUT_DIR}/dreambooth.safetensors --use_safetensors', shell=True)

    if bool(args.save_nfs) == True:
        nfs_path = args.nfs_mnt_dir

        if not os.path.exists(nfs_path):
            print("nfs not exist")
        else:
            if not os.path.exists(nfs_path + '/' + args.method):
               os.mkdir(nfs_path + '/' + args.method)
               print(f"{nfs_path}/{args.method} has been created.")
            else:
               print(f"{nfs_path}/{args.method} already exists.")
            copy_cmd = f'cp {OUTPUT_DIR}/*.safetensors {nfs_path}/{args.method}'
            subprocess.run(copy_cmd, shell=True)
            subprocess.run(f'ls {nfs_path}/{args.method}', shell=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--method", type=str, default="diffuser_dreambooth_lora", help="dreambooth/dreambooth_lora/peft_lora")
    parser.add_argument("--model_name", type=str, default="runwayml/stable-diffusion-v1-5", help="bucket_name/model_folder")
    parser.add_argument("--input_storage", type=str,default="abc", help="/gcs/bucket_name/input_image_folder")
    parser.add_argument("--output_storage", type=str, default="abc",help="/gcs/bucket_name/output_folder")
    parser.add_argument("--prompt", type=str, default="a photo of sks dog",help="instance prompt")
    parser.add_argument("--class_prompt", type=str, required=False,default="a photo of dog",help="instance prompt")
    parser.add_argument("--max_train_steps", type=int, default=400,help="training steps")
    parser.add_argument("--text_encoder", type=bool, default=True,help="train text encoder")
    parser.add_argument("--set_grads_to_none", type=bool, default=False,help="set grads to none if CUDA memory if very low <16GB")
    parser.add_argument("--num_class_images", type=int, default=50, help="generate image number")
    
    parser.add_argument("--resolution", type=int, default=512, help="resize input image resolution to")
    parser.add_argument("--batch_size", type=int, default=1, help="training batch size")
    parser.add_argument("--use_8bit", type=bool, default=True, help="use 8bit adam optimizer")
    parser.add_argument("--lr", type=float, default=1e-4, help="lora=1e-4,others=1e-5")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="gradient accumulation steps")
    parser.add_argument("--num_validation_images", type=int, default=4, help="num_validation_images, text_to_image_lora")
    parser.add_argument("--validation_prompt", type=str, default="a photo of sks dog in the forest", help="validation prompt")
    parser.add_argument("--save_nfs", type=bool, default=False, help="if save the model to file store")
    parser.add_argument("--save_nfs_only", type=bool, default=False, help="only copy file from gcs to filestore, no training")
    parser.add_argument("--nfs_mnt_dir", type=str, default="/mnt/nfs/model_repo", help="Filestore's mount directory")

    args = parser.parse_args()
    print(args)
    if bool(args.save_nfs_only) == True:
        nfs_path = args.nfs_mnt_dir
        if not os.path.exists(nfs_path):
            print("nfs not exist")
        else:
            if not os.path.exists(nfs_path + '/' + args.method):
               os.mkdir(nfs_path + '/' + args.method)
               print(f"{nfs_path}/{args.method} has been created.")
            else:
               print(f"{nfs_path}/{args.method} already exists.")
            copy_cmd = f'cp {args.output_storage}/*.safetensors {nfs_path}/{args.method}'
            subprocess.run(copy_cmd, shell=True)
            subprocess.run(f'ls {nfs_path}/{args.method}', shell=True) 
    else:
       main(args)