From a018cd6107a7661c12c87fc643903a00c14747b1 Mon Sep 17 00:00:00 2001 From: Wang Qiang <37444407+XDUWQ@users.noreply.github.com> Date: Wed, 28 Jun 2023 20:10:28 +0800 Subject: [PATCH] Dreambooth method for finetuning stable diffusions (#339) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Copyright * dreambooth * dreambooth test trainer * fix bugs * pre-commit --------- Co-authored-by: 翊靖 --- modelscope/metainfo.py | 1 + .../stable_diffusion/stable_diffusion.py | 30 +- modelscope/preprocessors/multi_modal.py | 2 +- .../dreambooth_diffusion/__init__.py | 2 + .../dreambooth_diffusion_trainer.py | 384 ++++++++++++++++++ .../multi_modal/lora_diffusion/__init__.py | 1 + .../test_diffusers_stable_diffusion.py | 2 +- .../test_dreambooth_diffusion_trainer.py | 98 +++++ tests/trainers/test_lora_diffusion_trainer.py | 4 +- 9 files changed, 505 insertions(+), 19 deletions(-) create mode 100644 modelscope/trainers/multi_modal/dreambooth_diffusion/__init__.py create mode 100644 modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py create mode 100644 tests/trainers/test_dreambooth_diffusion_trainer.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index f2529be2..d3365b7c 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -895,6 +895,7 @@ class MultiModalTrainers(object): efficient_diffusion_tuning = 'efficient-diffusion-tuning' stable_diffusion = 'stable-diffusion' lora_diffusion = 'lora-diffusion' + dreambooth_diffusion = 'dreambooth-diffusion' class AudioTrainers(object): diff --git a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py index 72f29b56..88cb4969 100644 --- a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py +++ b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py @@ -30,13 +30,12 @@ class StableDiffusion(TorchModel): """ Initialize a vision efficient diffusion tuning model. Args: - model_dir: model id or path, where model_dir/pytorch_model.bin + model_dir: model id or path """ super().__init__(model_dir, *args, **kwargs) - pretrained_model_name_or_path = kwargs.pop( - 'pretrained_model_name_or_path', 'runwayml/stable-diffusion-v1-5') revision = kwargs.pop('revision', None) - self.lora_tune = kwargs.pop('lora_tune', True) + self.lora_tune = kwargs.pop('lora_tune', False) + self.dreambooth_tune = kwargs.pop('dreambooth_tune', False) self.weight_dtype = torch.float32 self.device = torch.device( @@ -44,19 +43,15 @@ class StableDiffusion(TorchModel): # Load scheduler, tokenizer and models self.noise_scheduler = DDPMScheduler.from_pretrained( - pretrained_model_name_or_path, subfolder='scheduler') + model_dir, subfolder='scheduler') self.tokenizer = CLIPTokenizer.from_pretrained( - pretrained_model_name_or_path, - subfolder='tokenizer', - revision=revision) + model_dir, subfolder='tokenizer', revision=revision) self.text_encoder = CLIPTextModel.from_pretrained( - pretrained_model_name_or_path, - subfolder='text_encoder', - revision=revision) + model_dir, subfolder='text_encoder', revision=revision) self.vae = AutoencoderKL.from_pretrained( - pretrained_model_name_or_path, subfolder='vae', revision=revision) + model_dir, subfolder='vae', revision=revision) self.unet = UNet2DConditionModel.from_pretrained( - pretrained_model_name_or_path, subfolder='unet', revision=revision) + model_dir, subfolder='unet', revision=revision) self.safety_checker = None # Freeze gradient calculation and move to device @@ -90,6 +85,7 @@ class StableDiffusion(TorchModel): self.unet.train() self.unet = self.unet.to(self.device) + # Convert to latent space with torch.no_grad(): latents = self.vae.encode( target.to(dtype=self.weight_dtype)).latent_dist.sample() @@ -131,6 +127,9 @@ class StableDiffusion(TorchModel): model_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states).sample + if model_pred.shape[1] == 6: + model_pred, _ = torch.chunk(model_pred, 2, dim=1) + loss = F.mse_loss(model_pred.float(), target.float(), reduction='mean') output = {OutputKeys.LOSS: loss} @@ -144,8 +143,9 @@ class StableDiffusion(TorchModel): config: Optional[dict] = None, save_config_function: Callable = save_configuration, **kwargs): - # Save only the lora model, skip saving and copying the original weights - if self.lora_tune: + config['pipeline']['type'] = 'diffusers-stable-diffusion' + # Skip copying the original weights for lora and dreambooth method + if self.lora_tune or self.dreambooth_tune: pass else: super().save_pretrained(target_folder, save_checkpoint_names, diff --git a/modelscope/preprocessors/multi_modal.py b/modelscope/preprocessors/multi_modal.py index 82d44da8..a8867aef 100644 --- a/modelscope/preprocessors/multi_modal.py +++ b/modelscope/preprocessors/multi_modal.py @@ -11,6 +11,7 @@ import torch from PIL import Image from timm.data import create_transform from torchvision import transforms +from torchvision.datasets import ImageFolder from torchvision.transforms import Compose, Normalize, Resize, ToTensor from modelscope.hub.snapshot_download import snapshot_download @@ -55,7 +56,6 @@ class DiffusionImageGenerationPreprocessor(Preprocessor): transforms.Resize( self.preprocessor_resolution, interpolation=transforms.InterpolationMode.BILINEAR), - transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(self.preprocessor_mean, self.preprocessor_std), diff --git a/modelscope/trainers/multi_modal/dreambooth_diffusion/__init__.py b/modelscope/trainers/multi_modal/dreambooth_diffusion/__init__.py new file mode 100644 index 00000000..430d3c9e --- /dev/null +++ b/modelscope/trainers/multi_modal/dreambooth_diffusion/__init__.py @@ -0,0 +1,2 @@ +# Copyright © Alibaba, Inc. and its affiliates. +from .dreambooth_diffusion_trainer import DreamboothDiffusionTrainer diff --git a/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py b/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py new file mode 100644 index 00000000..65623ed8 --- /dev/null +++ b/modelscope/trainers/multi_modal/dreambooth_diffusion/dreambooth_diffusion_trainer.py @@ -0,0 +1,384 @@ +# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved. +import hashlib +import itertools +import shutil +import warnings +from collections.abc import Mapping +from pathlib import Path +from typing import Union + +import torch +import torch.nn.functional as F +from diffusers import DiffusionPipeline +from diffusers.loaders import AttnProcsLayers +from diffusers.models.attention_processor import LoRAAttnProcessor +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torchvision import transforms +from tqdm.auto import tqdm + +from modelscope.metainfo import Trainers +from modelscope.outputs import ModelOutputBase, OutputKeys +from modelscope.trainers.builder import TRAINERS +from modelscope.trainers.hooks.checkpoint.checkpoint_hook import CheckpointHook +from modelscope.trainers.hooks.checkpoint.checkpoint_processor import \ + CheckpointProcessor +from modelscope.trainers.optimizer.builder import build_optimizer +from modelscope.trainers.trainer import EpochBasedTrainer +from modelscope.utils.config import ConfigDict +from modelscope.utils.constant import ModeKeys +from modelscope.utils.file_utils import func_receive_dict_inputs +from modelscope.utils.torch_utils import is_dist + + +class DreamboothCheckpointProcessor(CheckpointProcessor): + + def __init__(self, model_dir): + self.model_dir = model_dir + + def save_checkpoints(self, + trainer, + checkpoint_path_prefix, + output_dir, + meta=None): + """Save the state dict for dreambooth model. + """ + pipeline_args = {} + if trainer.model.text_encoder is not None: + pipeline_args['text_encoder'] = trainer.model.text_encoder + pipeline = DiffusionPipeline.from_pretrained( + self.model_dir, + unet=trainer.model.unet, + **pipeline_args, + ) + scheduler_args = {} + pipeline.scheduler = pipeline.scheduler.from_config( + pipeline.scheduler.config, **scheduler_args) + pipeline.save_pretrained(output_dir) + + +class ClassDataset(Dataset): + + def __init__( + self, + tokenizer, + class_data_root=None, + class_prompt=None, + class_num_images=None, + size=512, + center_crop=False, + ): + """A dataset to prepare class images with the prompts for fine-tuning the model. + It pre-processes the images and the tokenizes prompts. + + Args: + tokenizer: The tokenizer to use for tokenization. + class_data_root: The saved class data path. + class_prompt: The prompt to use for class images. + class_num_images: The number of class images to use. + size: The size to resize the images. + center_crop: Whether to do center crop or random crop. + + """ + self.size = size + self.center_crop = center_crop + self.tokenizer = tokenizer + + if class_data_root is not None: + self.class_data_root = Path(class_data_root) + self.class_data_root.mkdir(parents=True, exist_ok=True) + self.class_images_path = list(self.class_data_root.iterdir()) + if class_num_images is not None: + self.num_class_images = min( + len(self.class_images_path), class_num_images) + else: + self.num_class_images = len(self.class_images_path) + self.class_prompt = class_prompt + else: + raise ValueError( + f"Class {self.class_data_root} class data root doesn't exists." + ) + + self.image_transforms = transforms.Compose([ + transforms.Resize( + size, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(size) + if center_crop else transforms.RandomCrop(size), + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), + ]) + + def __len__(self): + return self.num_class_images + + def __getitem__(self, index): + example = {} + + if self.class_data_root: + class_image = Image.open( + self.class_images_path[index % self.num_class_images]) + class_image = exif_transpose(class_image) + + if not class_image.mode == 'RGB': + class_image = class_image.convert('RGB') + example['pixel_values'] = self.image_transforms(class_image) + + class_text_inputs = self.tokenizer( + self.class_prompt, + max_length=self.tokenizer.model_max_length, + truncation=True, + padding='max_length', + return_tensors='pt') + input_ids = torch.squeeze(class_text_inputs.input_ids) + example['input_ids'] = input_ids + + return example + + +class PromptDataset(Dataset): + + def __init__(self, prompt, num_samples): + """Dataset to prepare the prompts to generate class images. + + Args: + prompt: Class prompt. + num_samples: The number sample for class images. + + """ + self.prompt = prompt + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, index): + example = {} + example['prompt'] = self.prompt + example['index'] = index + return example + + +@TRAINERS.register_module(module_name=Trainers.dreambooth_diffusion) +class DreamboothDiffusionTrainer(EpochBasedTrainer): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + """Dreambooth trainers for fine-tuning stable diffusion + + Args: + with_prior_preservation: a boolean indicating whether to enable prior loss. + instance_prompt: a string specifying the instance prompt. + class_prompt: a string specifying the class prompt. + class_data_dir: the path to the class data directory. + num_class_images: the number of class images to generate. + prior_loss_weight: the weight of the prior loss. + + """ + self.with_prior_preservation = kwargs.pop('with_prior_preservation', + False) + self.instance_prompt = kwargs.pop('instance_prompt', + 'a photo of sks dog') + self.class_prompt = kwargs.pop('class_prompt', 'a photo of dog') + self.class_data_dir = kwargs.pop('class_data_dir', '/tmp/class_data') + self.num_class_images = kwargs.pop('num_class_images', 200) + self.resolution = kwargs.pop('resolution', 512) + self.prior_loss_weight = kwargs.pop('prior_loss_weight', 1.0) + + # Save checkpoint and configurate files. + ckpt_hook = list( + filter(lambda hook: isinstance(hook, CheckpointHook), + self.hooks))[0] + ckpt_hook.set_processor(DreamboothCheckpointProcessor(self.model_dir)) + + # Check for conflicts and conflicts + if self.with_prior_preservation: + if self.class_data_dir is None: + raise ValueError( + 'You must specify a data directory for class images.') + if self.class_prompt is None: + raise ValueError('You must specify prompt for class images.') + else: + if self.class_data_dir is not None: + warnings.warn( + 'You need not use --class_data_dir without --with_prior_preservation.' + ) + if self.class_prompt is not None: + warnings.warn( + 'You need not use --class_prompt without --with_prior_preservation.' + ) + + # Generate class images if prior preservation is enabled. + if self.with_prior_preservation: + class_images_dir = Path(self.class_data_dir) + if not class_images_dir.exists(): + class_images_dir.mkdir(parents=True) + cur_class_images = len(list(class_images_dir.iterdir())) + + if cur_class_images < self.num_class_images: + if torch.cuda.device_count() > 1: + warnings.warn('Multiple GPU inference not yet supported.') + pipeline = DiffusionPipeline.from_pretrained( + self.model_dir, + torch_dtype=torch.float32, + safety_checker=None, + revision=None, + ) + pipeline.set_progress_bar_config(disable=True) + + num_new_images = self.num_class_images - cur_class_images + sample_dataset = PromptDataset(self.instance_prompt, + num_new_images) + sample_dataloader = torch.utils.data.DataLoader(sample_dataset) + + pipeline.to(self.device) + + for example in tqdm( + sample_dataloader, desc='Generating class images'): + images = pipeline(example['prompt']).images + for i, image in enumerate(images): + hash_image = hashlib.sha1(image.tobytes()).hexdigest() + image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" + image.save(image_filename) + + del pipeline + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Class Dataset and DataLoaders creation + class_dataset = ClassDataset( + class_data_root=self.class_data_dir + if self.with_prior_preservation else None, + class_prompt=self.class_prompt, + class_num_images=self.num_class_images, + tokenizer=self.model.tokenizer, + size=self.resolution, + center_crop=False, + ) + class_dataloader = torch.utils.data.DataLoader( + class_dataset, + batch_size=1, + shuffle=True, + ) + self.iter_class_dataloader = itertools.cycle(class_dataloader) + + def build_optimizer(self, cfg: ConfigDict, default_args: dict = None): + try: + return build_optimizer( + self.model.unet.parameters(), + cfg=cfg, + default_args=default_args) + except KeyError as e: + self.logger.error( + f'Build optimizer error, the optimizer {cfg} is a torch native component, ' + f'please check if your torch with version: {torch.__version__} matches the config.' + ) + raise e + + def train_step(self, model, inputs): + """ Perform a training step on a batch of inputs. + + Subclass and override to inject custom behavior. + + Args: + model (`TorchModel`): The model to train. + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + + Return: + `torch.Tensor`: The tensor with training loss on this batch. + """ + model.train() + self._mode = ModeKeys.TRAIN + # call model forward but not __call__ to skip postprocess + + receive_dict_inputs = func_receive_dict_inputs( + self.unwrap_module(self.model).forward) + + if isinstance(inputs, Mapping) and not receive_dict_inputs: + train_outputs = model.forward(**inputs) + else: + train_outputs = model.forward(inputs) + + if self.with_prior_preservation: + # Convert to latent space + batch = next(self.iter_class_dataloader) + target_prior = batch['pixel_values'].to(self.device) + 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() + latents = latents * self.model.vae.config.scaling_factor + + # Sample noise that we'll add to the latents + noise = torch.randn_like(latents) + bsz = latents.shape[0] + # Sample a random timestep for each image + timesteps = torch.randint( + 0, + self.model.noise_scheduler.num_train_timesteps, (bsz, ), + device=latents.device) + timesteps = timesteps.long() + + # Add noise to the latents according to the noise magnitude at each timestep + # (this is the forward diffusion process) + noisy_latents = self.model.noise_scheduler.add_noise( + latents, noise, timesteps) + + # Get the text embedding for conditioning + with torch.no_grad(): + encoder_hidden_states = self.model.text_encoder(input_ids)[0] + + # Get the target for loss depending on the prediction type + if self.model.noise_scheduler.config.prediction_type == 'epsilon': + target_prior = noise + elif self.model.noise_scheduler.config.prediction_type == 'v_prediction': + target_prior = self.model.noise_scheduler.get_velocity( + latents, noise, timesteps) + else: + raise ValueError( + f'Unknown prediction type {self.model.noise_scheduler.config.prediction_type}' + ) + + # Predict the noise residual and compute loss + model_pred_prior = self.model.unet(noisy_latents, timesteps, + encoder_hidden_states).sample + + # Compute prior loss + prior_loss = F.mse_loss( + model_pred_prior.float(), + target_prior.float(), + reduction='mean') + # Add the prior loss to the instance loss. + train_outputs[ + OutputKeys.LOSS] += self.prior_loss_weight * prior_loss + + if isinstance(train_outputs, ModelOutputBase): + train_outputs = train_outputs.to_dict() + if not isinstance(train_outputs, dict): + raise TypeError('"model.forward()" must return a dict') + + # add model output info to log + if 'log_vars' not in train_outputs: + default_keys_pattern = ['loss'] + match_keys = set([]) + for key_p in default_keys_pattern: + match_keys.update( + [key for key in train_outputs.keys() if key_p in key]) + + log_vars = {} + for key in match_keys: + value = train_outputs.get(key, None) + if value is not None: + if is_dist(): + value = value.data.clone().to('cuda') + dist.all_reduce(value.div_(dist.get_world_size())) + log_vars.update({key: value.item()}) + self.log_buffer.update(log_vars) + else: + self.log_buffer.update(train_outputs['log_vars']) + + self.train_outputs = train_outputs diff --git a/modelscope/trainers/multi_modal/lora_diffusion/__init__.py b/modelscope/trainers/multi_modal/lora_diffusion/__init__.py index 311d2789..ebddd00b 100644 --- a/modelscope/trainers/multi_modal/lora_diffusion/__init__.py +++ b/modelscope/trainers/multi_modal/lora_diffusion/__init__.py @@ -1 +1,2 @@ +# Copyright © Alibaba, Inc. and its affiliates. from .lora_diffusion_trainer import LoraDiffusionTrainer diff --git a/tests/pipelines/test_diffusers_stable_diffusion.py b/tests/pipelines/test_diffusers_stable_diffusion.py index 57eae4a3..432706c7 100644 --- a/tests/pipelines/test_diffusers_stable_diffusion.py +++ b/tests/pipelines/test_diffusers_stable_diffusion.py @@ -21,7 +21,7 @@ class DiffusersStableDiffusionTest(unittest.TestCase): def test_run(self): diffusers_pipeline = pipeline(task=self.task, model=self.model_id) output = diffusers_pipeline({ - 'prompt': self.test_input, + 'text': self.test_input, 'height': 512, 'width': 512 }) diff --git a/tests/trainers/test_dreambooth_diffusion_trainer.py b/tests/trainers/test_dreambooth_diffusion_trainer.py new file mode 100644 index 00000000..b57d8a90 --- /dev/null +++ b/tests/trainers/test_dreambooth_diffusion_trainer.py @@ -0,0 +1,98 @@ +# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved. +import os +import shutil +import tempfile +import unittest + +import cv2 + +from modelscope.metainfo import Trainers +from modelscope.msdatasets import MsDataset +from modelscope.pipelines import pipeline +from modelscope.trainers import build_trainer +from modelscope.utils.constant import DownloadMode +from modelscope.utils.test_utils import test_level + + +class TestDreamboothDiffusionTrainer(unittest.TestCase): + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + + self.train_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='train', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + self.eval_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='validation', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + + self.max_epochs = 5 + + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + super().tearDown() + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_dreambooth_diffusion_train(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.8' + prompt = 'a dog.' + + def cfg_modify_fn(cfg): + cfg.train.max_epochs = self.max_epochs + cfg.train.lr_scheduler = { + 'type': 'LambdaLR', + 'lr_lambda': lambda _: 1, + 'last_epoch': -1 + } + cfg.train.optimizer.lr = 5e-6 + return cfg + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=self.train_dataset, + eval_dataset=self.eval_dataset, + cfg_modify_fn=cfg_modify_fn) + + trainer = build_trainer( + name=Trainers.dreambooth_diffusion, default_args=kwargs) + trainer.train() + result = trainer.evaluate() + print(f'Dreambooth-diffusion train output: {result}.') + + results_files = os.listdir(self.tmp_dir) + self.assertIn(f'{trainer.timestamp}.log.json', results_files) + + pipe = pipeline( + task=Tasks.text_to_image_synthesis, model=f'{self.tmp_dir}/output') + output = pipe({'text': prompt}) + cv2.imwrite('./dreambooth_result.png', output['output_imgs'][0]) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_dreambooth_diffusion_eval(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.8' + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=None, + eval_dataset=self.eval_dataset) + + trainer = build_trainer( + name=Trainers.dreambooth_diffusion, default_args=kwargs) + result = trainer.evaluate() + print(f'Dreambooth-diffusion eval output: {result}.') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/trainers/test_lora_diffusion_trainer.py b/tests/trainers/test_lora_diffusion_trainer.py index 2ba89665..a9b9e299 100644 --- a/tests/trainers/test_lora_diffusion_trainer.py +++ b/tests/trainers/test_lora_diffusion_trainer.py @@ -38,7 +38,7 @@ class TestLoraDiffusionTrainer(unittest.TestCase): @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_lora_diffusion_train(self): model_id = 'AI-ModelScope/stable-diffusion-v1-5' - model_revision = 'v1.0.6' + model_revision = 'v1.0.9' def cfg_modify_fn(cfg): cfg.train.max_epochs = self.max_epochs @@ -70,7 +70,7 @@ class TestLoraDiffusionTrainer(unittest.TestCase): @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_lora_diffusion_eval(self): model_id = 'AI-ModelScope/stable-diffusion-v1-5' - model_revision = 'v1.0.6' + model_revision = 'v1.0.9' kwargs = dict( model=model_id,