support float16 training and pipeline for stable diffusion (#447)

* 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: 翊靖 <yijing.wq@alibaba-inc.com>

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='<new1>+<new2>',
+    modifier_token=args.modifier_token,
     model_revision=args.model_revision)
 
 output = pipe({'text': args.instance_prompt})

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 "<new1>+<new2>" \
+    --modifier_token "<new1>" \
     --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

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

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

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

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

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:

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)

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

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

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