From a67d339e3bf8abd25d224818f6fd6512078f6f37 Mon Sep 17 00:00:00 2001 From: Wang Qiang <37444407+XDUWQ@users.noreply.github.com> Date: Tue, 15 Aug 2023 20:04:32 +0800 Subject: [PATCH] support float16 training and pipeline for stable diffusion (#447) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * support float16 traing and pipeline for stable diffusion * pre commit * fix bugs * add torch type example * fix bugs of torch type * support type float16 * fix bugs of load pipeline * change type to fp16 * lora rank --------- Co-authored-by: 翊靖 --- .../finetune_stable_diffusion_custom.py | 11 +++++++++- .../custom/run_train_custom.sh | 3 ++- .../finetune_stable_diffusion_dreambooth.py | 9 ++++++++ .../dreambooth/run_train_dreambooth.sh | 1 + .../lora/finetune_stable_diffusion_lora.py | 9 ++++++++ .../stable_diffusion/lora/run_train_lora.sh | 3 ++- .../stable_diffusion/stable_diffusion.py | 9 ++++---- .../stable_diffusion_pipeline.py | 2 +- .../custom_diffusion_trainer.py | 22 +++++++++++++------ .../dreambooth_diffusion_trainer.py | 17 +++++++++++--- .../lora_diffusion/lora_diffusion_trainer.py | 16 ++++++++++++-- 11 files changed, 82 insertions(+), 20 deletions(-) diff --git a/examples/pytorch/stable_diffusion/custom/finetune_stable_diffusion_custom.py b/examples/pytorch/stable_diffusion/custom/finetune_stable_diffusion_custom.py index 007ea82b..76a050c4 100644 --- a/examples/pytorch/stable_diffusion/custom/finetune_stable_diffusion_custom.py +++ b/examples/pytorch/stable_diffusion/custom/finetune_stable_diffusion_custom.py @@ -2,6 +2,7 @@ import os from dataclasses import dataclass, field import cv2 +import torch from modelscope.metainfo import Trainers from modelscope.msdatasets import MsDataset @@ -95,6 +96,12 @@ class StableDiffusionCustomArguments(TrainingArgs): 'help': 'Path to json containing multiple concepts.', }) + torch_type: str = field( + default='float32', + metadata={ + 'help': ' The torch type, default is float32.', + }) + training_args = StableDiffusionCustomArguments( task='text-to-image-synthesis').parse_cli() @@ -148,6 +155,8 @@ kwargs = dict( work_dir=training_args.work_dir, train_dataset=train_dataset, eval_dataset=validation_dataset, + torch_type=torch.float16 + if args.torch_type == 'float16' else torch.float32, cfg_modify_fn=cfg_modify_fn) # build trainer and training @@ -159,7 +168,7 @@ pipe = pipeline( task=Tasks.text_to_image_synthesis, model=training_args.model, custom_dir=training_args.work_dir + '/output', - modifier_token='+', + modifier_token=args.modifier_token, model_revision=args.model_revision) output = pipe({'text': args.instance_prompt}) diff --git a/examples/pytorch/stable_diffusion/custom/run_train_custom.sh b/examples/pytorch/stable_diffusion/custom/run_train_custom.sh index fab8e059..7f9cb500 100644 --- a/examples/pytorch/stable_diffusion/custom/run_train_custom.sh +++ b/examples/pytorch/stable_diffusion/custom/run_train_custom.sh @@ -7,11 +7,12 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/custom/finetune_stable_d --class_data_dir './tmp/class_data' \ --train_dataset_name 'buptwq/lora-stable-diffusion-finetune-dog' \ --max_epochs 250 \ - --modifier_token "+" \ + --modifier_token "" \ --num_class_images=200 \ --save_ckpt_strategy 'by_epoch' \ --logging_interval 1 \ --train.dataloader.workers_per_gpu 0 \ --evaluation.dataloader.workers_per_gpu 0 \ --train.optimizer.lr 1e-5 \ + --torch_type 'float32' \ --use_model_config true diff --git a/examples/pytorch/stable_diffusion/dreambooth/finetune_stable_diffusion_dreambooth.py b/examples/pytorch/stable_diffusion/dreambooth/finetune_stable_diffusion_dreambooth.py index 760396d0..5659a105 100644 --- a/examples/pytorch/stable_diffusion/dreambooth/finetune_stable_diffusion_dreambooth.py +++ b/examples/pytorch/stable_diffusion/dreambooth/finetune_stable_diffusion_dreambooth.py @@ -2,6 +2,7 @@ import os from dataclasses import dataclass, field import cv2 +import torch from modelscope.metainfo import Trainers from modelscope.msdatasets import MsDataset @@ -59,6 +60,12 @@ class StableDiffusionDreamboothArguments(TrainingArgs): 'help': 'The pipeline prompt.', }) + torch_type: str = field( + default='float32', + metadata={ + 'help': ' The torch type, default is float32.', + }) + training_args = StableDiffusionDreamboothArguments( task='text-to-image-synthesis').parse_cli() @@ -106,6 +113,8 @@ kwargs = dict( resolution=args.resolution, prior_loss_weight=args.prior_loss_weight, prompt=args.prompt, + torch_type=torch.float16 + if args.torch_type == 'float16' else torch.float32, cfg_modify_fn=cfg_modify_fn) # build trainer and training diff --git a/examples/pytorch/stable_diffusion/dreambooth/run_train_dreambooth.sh b/examples/pytorch/stable_diffusion/dreambooth/run_train_dreambooth.sh index 03c51ffb..461434ee 100644 --- a/examples/pytorch/stable_diffusion/dreambooth/run_train_dreambooth.sh +++ b/examples/pytorch/stable_diffusion/dreambooth/run_train_dreambooth.sh @@ -17,4 +17,5 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/dreambooth/finetune_stab --train.dataloader.workers_per_gpu 0 \ --evaluation.dataloader.workers_per_gpu 0 \ --train.optimizer.lr 5e-6 \ + --torch_type 'float32' \ --use_model_config true diff --git a/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py b/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py index 6001af48..f0c40be7 100644 --- a/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py +++ b/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py @@ -2,6 +2,7 @@ import os from dataclasses import dataclass, field import cv2 +import torch from modelscope.metainfo import Trainers from modelscope.msdatasets import MsDataset @@ -25,6 +26,12 @@ class StableDiffusionLoraArguments(TrainingArgs): 'help': 'The rank size of lora intermediate linear.', }) + torch_type: str = field( + default='float32', + metadata={ + 'help': ' The torch type, default is float32.', + }) + training_args = StableDiffusionLoraArguments( task='text-to-image-synthesis').parse_cli() @@ -66,6 +73,8 @@ kwargs = dict( train_dataset=train_dataset, eval_dataset=validation_dataset, lora_rank=args.lora_rank, + torch_type=torch.float16 + if args.torch_type == 'float16' else torch.float32, cfg_modify_fn=cfg_modify_fn) # build trainer and training diff --git a/examples/pytorch/stable_diffusion/lora/run_train_lora.sh b/examples/pytorch/stable_diffusion/lora/run_train_lora.sh index d4f1b07d..82e31aad 100644 --- a/examples/pytorch/stable_diffusion/lora/run_train_lora.sh +++ b/examples/pytorch/stable_diffusion/lora/run_train_lora.sh @@ -5,10 +5,11 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/lora/finetune_stable_dif --work_dir './tmp/lora_diffusion' \ --train_dataset_name 'buptwq/lora-stable-diffusion-finetune' \ --max_epochs 100 \ - --lora_rank 4 \ + --lora_rank 16 \ --save_ckpt_strategy 'by_epoch' \ --logging_interval 1 \ --train.dataloader.workers_per_gpu 0 \ --evaluation.dataloader.workers_per_gpu 0 \ --train.optimizer.lr 1e-4 \ + --torch_type 'float16' \ --use_model_config true diff --git a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py index 6b829485..80d8ab28 100644 --- a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py +++ b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py @@ -39,7 +39,7 @@ class StableDiffusion(TorchModel): self.lora_tune = kwargs.pop('lora_tune', False) self.dreambooth_tune = kwargs.pop('dreambooth_tune', False) - self.weight_dtype = torch.float32 + self.weight_dtype = kwargs.pop('torch_type', torch.float32) self.device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu') @@ -59,14 +59,15 @@ class StableDiffusion(TorchModel): # Freeze gradient calculation and move to device if self.vae is not None: self.vae.requires_grad_(False) - self.vae = self.vae.to(self.device) + self.vae = self.vae.to(self.device, dtype=self.weight_dtype) if self.text_encoder is not None: self.text_encoder.requires_grad_(False) - self.text_encoder = self.text_encoder.to(self.device) + self.text_encoder = self.text_encoder.to( + self.device, dtype=self.weight_dtype) if self.unet is not None: if self.lora_tune: self.unet.requires_grad_(False) - self.unet = self.unet.to(self.device) + self.unet = self.unet.to(self.device, dtype=self.weight_dtype) # xformers accelerate memory efficient attention if xformers_enable: diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py index 7e56f24c..e5345543 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py @@ -40,6 +40,7 @@ class StableDiffusionPipeline(DiffusersPipeline): use_safetensors: load safetensors weights. """ use_safetensors = kwargs.pop('use_safetensors', False) + torch_type = kwargs.pop('torch_type', torch.float32) # check custom diffusion input value if custom_dir is None and modifier_token is not None: raise ValueError( @@ -50,7 +51,6 @@ class StableDiffusionPipeline(DiffusersPipeline): self.device = 'cuda' if torch.cuda.is_available() else 'cpu' # load pipeline - torch_type = torch.float16 if self.device == 'cuda' else torch.float32 self.pipeline = DiffusionPipeline.from_pretrained( model, use_safetensors=use_safetensors, torch_dtype=torch_type) self.pipeline = self.pipeline.to(self.device) diff --git a/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py b/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py index 28140fb2..435fd2e3 100644 --- a/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py +++ b/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py @@ -37,24 +37,31 @@ from modelscope.utils.torch_utils import is_dist class CustomCheckpointProcessor(CheckpointProcessor): - def __init__(self, modifier_token, modifier_token_id): + def __init__(self, + modifier_token, + modifier_token_id, + torch_type=torch.float32): """Checkpoint processor for custom diffusion. Args: modifier_token: The token to use as a modifier for the concept. modifier_token_id: The modifier token id for the concept. + torch_type: The torch type, default is float32. + """ self.modifier_token = modifier_token self.modifier_token_id = modifier_token_id + self.torch_type = torch_type def save_checkpoints(self, trainer, checkpoint_path_prefix, output_dir, - meta=None): + meta=None, + save_optimizers=True): """Save the state dict for custom diffusion model. """ - trainer.model.unet = trainer.model.unet.to(torch.float32) + trainer.model.unet = trainer.model.unet.to(self.torch_type) trainer.model.unet.save_attn_procs(output_dir) learned_embeds = trainer.model.text_encoder.get_input_embeddings( @@ -281,6 +288,7 @@ class CustomDiffusionTrainer(EpochBasedTrainer): instance_prompt = kwargs.pop('instance_prompt', 'a photo of sks dog') class_prompt = kwargs.pop('class_prompt', 'dog') class_data_dir = kwargs.pop('class_data_dir', '/tmp/class_data') + self.torch_type = kwargs.pop('torch_type', torch.float32) self.real_prior = kwargs.pop('real_prior', False) self.num_class_images = kwargs.pop('num_class_images', 200) self.resolution = kwargs.pop('resolution', 512) @@ -387,7 +395,7 @@ class CustomDiffusionTrainer(EpochBasedTrainer): self.hooks))[0] ckpt_hook.set_processor( CustomCheckpointProcessor(self.modifier_token, - self.modifier_token_id)) + self.modifier_token_id, self.torch_type)) # Add new Custom Diffusion weights to the attention layers attention_class = CustomDiffusionAttnProcessor @@ -477,7 +485,7 @@ class CustomDiffusionTrainer(EpochBasedTrainer): size=self.resolution, mask_size=self.model.vae.encode( torch.randn(1, 3, self.resolution, - self.resolution).to(dtype=torch.float32).to( + self.resolution).to(dtype=self.torch_type).to( self.device)).latent_dist.sample().size()[-1], center_crop=self.center_crop, num_class_images=self.num_class_images, @@ -534,8 +542,8 @@ class CustomDiffusionTrainer(EpochBasedTrainer): if cur_class_images < self.num_class_images: pipeline = DiffusionPipeline.from_pretrained( self.model_dir, - torch_dtype=torch.float32, safety_checker=None, + torch_dtype=self.torch_type, revision=None, ) pipeline.set_progress_bar_config(disable=True) @@ -656,7 +664,7 @@ class CustomDiffusionTrainer(EpochBasedTrainer): batch = next(self.iter_train_dataloader) # Convert images to latent space latents = self.model.vae.encode(batch['pixel_values'].to( - dtype=torch.float32).to(self.device)).latent_dist.sample() + dtype=self.torch_type).to(self.device)).latent_dist.sample() latents = latents * self.model.vae.config.scaling_factor # Sample noise that we'll add to the latents diff --git a/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py b/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py index 3b300ea4..f391f87a 100644 --- a/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py +++ b/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py @@ -32,8 +32,16 @@ from modelscope.utils.torch_utils import is_dist class DreamboothCheckpointProcessor(CheckpointProcessor): - def __init__(self, model_dir): + def __init__(self, model_dir, torch_type=torch.float32): + """Checkpoint processor for dreambooth diffusion. + + Args: + model_dir: The model id or local model dir. + torch_type: The torch type, default is float32. + + """ self.model_dir = model_dir + self.torch_type = torch_type def save_checkpoints(self, trainer, @@ -49,6 +57,7 @@ class DreamboothCheckpointProcessor(CheckpointProcessor): pipeline = DiffusionPipeline.from_pretrained( self.model_dir, unet=trainer.model.unet, + torch_type=self.torch_type, **pipeline_args, ) scheduler_args = {} @@ -174,6 +183,7 @@ class DreamboothDiffusionTrainer(EpochBasedTrainer): prior_loss_weight: the weight of the prior loss. """ + self.torch_type = kwargs.pop('torch_type', torch.float32) self.with_prior_preservation = kwargs.pop('with_prior_preservation', False) self.instance_prompt = kwargs.pop('instance_prompt', @@ -219,7 +229,7 @@ class DreamboothDiffusionTrainer(EpochBasedTrainer): warnings.warn('Multiple GPU inference not yet supported.') pipeline = DiffusionPipeline.from_pretrained( self.model_dir, - torch_dtype=torch.float32, + torch_dtype=self.torch_type, safety_checker=None, revision=None, ) @@ -309,7 +319,8 @@ class DreamboothDiffusionTrainer(EpochBasedTrainer): input_ids = batch['input_ids'].to(self.device) with torch.no_grad(): latents = self.model.vae.encode( - target_prior.to(dtype=torch.float32)).latent_dist.sample() + target_prior.to( + dtype=self.torch_type)).latent_dist.sample() latents = latents * self.model.vae.config.scaling_factor # Sample noise that we'll add to the latents diff --git a/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py b/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py index 7c6644bd..2e4d4090 100644 --- a/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py +++ b/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py @@ -17,6 +17,15 @@ from modelscope.utils.config import ConfigDict class LoraDiffusionCheckpointProcessor(CheckpointProcessor): + def __init__(self, torch_type=torch.float32): + """Checkpoint processor for lora diffusion. + + Args: + torch_type: The torch type, default is float32. + + """ + self.torch_type = torch_type + def save_checkpoints(self, trainer, checkpoint_path_prefix, @@ -25,7 +34,7 @@ class LoraDiffusionCheckpointProcessor(CheckpointProcessor): save_optimizers=True): """Save the state dict for lora tune model. """ - trainer.model.unet = trainer.model.unet.to(torch.float32) + trainer.model.unet = trainer.model.unet.to(self.torch_type) trainer.model.unet.save_attn_procs(output_dir) @@ -38,15 +47,18 @@ class LoraDiffusionTrainer(EpochBasedTrainer): Args: lora_rank: The rank size of lora intermediate linear. + torch_type: The torch type, default is float32. """ lora_rank = kwargs.pop('lora_rank', 4) + torch_type = kwargs.pop('torch_type', torch.float32) # set lora save checkpoint processor ckpt_hook = list( filter(lambda hook: isinstance(hook, CheckpointHook), self.hooks))[0] - ckpt_hook.set_processor(LoraDiffusionCheckpointProcessor()) + ckpt_hook.set_processor( + LoraDiffusionCheckpointProcessor(torch_type=torch_type)) # Set correct lora layers lora_attn_procs = {} for name in self.model.unet.attn_processors.keys():