fengshen/examples/finetune_taiyi_stable_diffusion/finetune.py

import os
import torch
import argparse
from pytorch_lightning import (
    LightningModule,
    Trainer,
)
from pytorch_lightning.callbacks import (
    LearningRateMonitor,
)
from fengshen.data.universal_datamodule import UniversalDataModule
from fengshen.models.model_utils import (
    add_module_args,
    configure_optimizers,
    get_total_steps,
)
from fengshen.utils.universal_checkpoint import UniversalCheckpoint
from transformers import BertTokenizer, BertModel
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
from torch.nn import functional as F
from fengshen.data.taiyi_stable_diffusion_datasets.taiyi_datasets import add_data_args, load_data
from torchvision import transforms
from PIL import Image
from torch.utils.data._utils.collate import default_collate


class Collator():
    def __init__(self, args, tokenizer):
        self.image_transforms = transforms.Compose(
            [
                transforms.Resize(
                    args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.CenterCrop(
                    args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )
        self.tokenizer = tokenizer

    def __call__(self, inputs):
        examples = []
        max_length = min(max([len(i['caption']) for i in inputs]), 512)
        for i in inputs:
            example = {}
            instance_image = Image.open(i['img_path'])
            if not instance_image.mode == "RGB":
                instance_image = instance_image.convert("RGB")
            example["pixel_values"] = self.image_transforms(instance_image)
            example["input_ids"] = self.tokenizer(
                i['caption'],
                padding="max_length",
                truncation=True,
                max_length=max_length,
                return_tensors='pt',
            )['input_ids'][0]
            examples.append(example)
        return default_collate(examples)


class StableDiffusion(LightningModule):
    @staticmethod
    def add_module_specific_args(parent_parser):
        parser = parent_parser.add_argument_group('Taiyi Stable Diffusion Module')
        parser.add_argument('--freeze_unet', action='store_true', default=False)
        parser.add_argument('--freeze_text_encoder', action='store_true', default=False)
        return parent_parser

    def __init__(self, args):
        super().__init__()
        self.tokenizer = BertTokenizer.from_pretrained(
            args.model_path, subfolder="tokenizer")
        self.text_encoder = BertModel.from_pretrained(
            args.model_path, subfolder="text_encoder")  # load from taiyi_finetune-v0
        self.vae = AutoencoderKL.from_pretrained(
            args.model_path, subfolder="vae")
        self.unet = UNet2DConditionModel.from_pretrained(
            args.model_path, subfolder="unet")
        # TODO: 使用xformers配合deepspeed速度反而有下降(待确认
        self.unet.set_use_memory_efficient_attention_xformers(False)

        self.noise_scheduler = DDPMScheduler(
            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000
        )

        for param in self.vae.parameters():
            param.requires_grad = False

        if args.freeze_text_encoder:
            for param in self.text_encoder.parameters():
                param.requires_grad = False

        if args.freeze_unet:
            for param in self.unet.parameters():
                param.requires_grad = False

        self.save_hyperparameters(args)

    def setup(self, stage) -> None:
        if stage == 'fit':
            self.total_steps = get_total_steps(self.trainer, self.hparams)
            print('Total steps: {}' .format(self.total_steps))

    def configure_optimizers(self):
        return configure_optimizers(self)

    def training_step(self, batch, batch_idx):
        self.text_encoder.train()

        latents = self.vae.encode(batch["pixel_values"]).latent_dist.sample()
        latents = latents * 0.18215

        # Sample noise that we'll add to the latents
        noise = torch.randn(latents.shape).to(latents.device)
        noise = noise.to(dtype=self.unet.dtype)
        bsz = latents.shape[0]
        # Sample a random timestep for each image
        timesteps = torch.randint(
            0, self.noise_scheduler.config.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.noise_scheduler.add_noise(latents, noise, timesteps)
        noisy_latents = noisy_latents.to(dtype=self.unet.dtype)

        # Get the text embedding for conditioning
        encoder_hidden_states = self.text_encoder(batch["input_ids"])[0]

        # Predict the noise residual
        noise_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states).sample

        loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
        self.log("train_loss", loss.item(),  on_epoch=False, prog_bar=True, logger=True)

        if self.trainer.global_rank == 0 and self.global_step == 100:
            # 打印显存占用
            from fengshen.utils.utils import report_memory
            report_memory('stable diffusion')

        return {"loss": loss}

    def on_save_checkpoint(self, checkpoint) -> None:
        if self.trainer.global_rank == 0:
            print('saving model...')
            pipeline = StableDiffusionPipeline.from_pretrained(
                self.hparams.model_path,
                text_encoder=self.text_encoder,
                tokenizer=self.tokenizer,
                unet=self.unet)
            self.trainer.current_epoch
            pipeline.save_pretrained(os.path.join(
                args.default_root_dir, f'hf_out_{self.trainer.current_epoch}_{self.trainer.global_step}'))

    def on_load_checkpoint(self, checkpoint) -> None:
        # 兼容低版本lightning，低版本lightning从ckpt起来时steps数会被重置为0
        global_step_offset = checkpoint["global_step"]
        if 'global_samples' in checkpoint:
            self.consumed_samples = checkpoint['global_samples']
        self.trainer.fit_loop.epoch_loop._batches_that_stepped = global_step_offset


if __name__ == '__main__':
    args_parser = argparse.ArgumentParser()
    args_parser = add_module_args(args_parser)
    args_parser = add_data_args(args_parser)
    args_parser = UniversalDataModule.add_data_specific_args(args_parser)
    args_parser = Trainer.add_argparse_args(args_parser)
    args_parser = StableDiffusion.add_module_specific_args(args_parser)
    args_parser = UniversalCheckpoint.add_argparse_args(args_parser)
    args = args_parser.parse_args()

    lr_monitor = LearningRateMonitor(logging_interval='step')
    checkpoint_callback = UniversalCheckpoint(args)
    trainer = Trainer.from_argparse_args(args,
                                         callbacks=[
                                             lr_monitor,
                                             checkpoint_callback])

    model = StableDiffusion(args)
    tokenizer = model.tokenizer
    datasets = load_data(args, global_rank=trainer.global_rank)
    collate_fn = Collator(args, tokenizer)

    datamoule = UniversalDataModule(
        tokenizer=tokenizer, collate_fn=collate_fn, args=args, datasets=datasets)

    trainer.fit(model, datamoule, ckpt_path=args.load_ckpt_path)