textual_inversion.py

import paddle
from PIL import Image
import gc
def image_grid(imgs, rows=2, cols=2):
    assert len(imgs) == rows * cols
    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols * w, rows * h))
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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 argparse
import glob
import itertools
import math
import os
import random
from pathlib import Path

import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.io import BatchSampler, DataLoader, Dataset, DistributedBatchSampler
from paddle.optimizer import AdamW
from paddle.vision.transforms import RandomHorizontalFlip
from PIL import Image
from tqdm.auto import tqdm

from paddlenlp.trainer import set_seed
from paddlenlp.transformers import AutoTokenizer, BertModel, CLIPTextModel
from paddlenlp.utils.log import logger
try:
    from ppdiffusers import (
        AutoencoderKL,
        DDPMScheduler,
        PNDMScheduler,
        StableDiffusionPipeline,
        UNet2DConditionModel,
        patch_to,
    )
    from ppdiffusers.modeling_utils import freeze_params, unwrap_model
    from ppdiffusers.optimization import get_scheduler
    from ppdiffusers.pipelines.alt_diffusion import RobertaSeriesModelWithTransformation
    from ppdiffusers.ppnlp_patch_utils import XLMRobertaTokenizer
    from ppdiffusers.utils import PIL_INTERPOLATION

    # patch
    @patch_to(RobertaSeriesModelWithTransformation)
    def get_input_embeddings(self):
        return self.roberta.embeddings.word_embeddings


    @patch_to(RobertaSeriesModelWithTransformation)
    def set_input_embeddings(self, value):
        self.roberta.embeddings.word_embeddings = value
except:
    AutoencoderKL = None
    DDPMScheduler = None
    PNDMScheduler = None
    StableDiffusionPipeline = None
    UNet2DConditionModel = None
    patch_to = None
    freeze_params, unwrap_model = None, None
    get_scheduler = None
    RobertaSeriesModelWithTransformation = None
    XLMRobertaTokenizer = None
    PIL_INTERPOLATION = None

def get_writer(args):
    if args.writer_type == "visualdl":
        from visualdl import LogWriter

        writer = LogWriter(logdir=args.logging_dir)
    elif args.writer_type == "tensorboard":
        from tensorboardX import SummaryWriter

        writer = SummaryWriter(logdir=args.logging_dir)
    else:
        raise ValueError("writer_type must be in ['visualdl', 'tensorboard']")
    return writer


def save_progress(text_encoder, placeholder_token_id, args, global_step=-1):
    learned_embeds = unwrap_model(
        text_encoder).get_input_embeddings().weight[placeholder_token_id]
    learned_embeds_dict = {
        args.placeholder_token: learned_embeds.detach().cpu()
    }

    # remove \/"*:?<>| in filename
    name = args.placeholder_token
    name = name.translate({
        92: 95,
        47: 95,
        42: 95,
        34: 95,
        58: 95,
        63: 95,
        60: 95,
        62: 95,
        124: 95
    })
    path = os.path.join(args.output_dir, "step-"+str(global_step))
    os.makedirs(path, exist_ok=True)
    paddle.save(learned_embeds_dict,
                os.path.join(args.output_dir, "step-"+str(global_step), f"{name}.pdparams"))
    print(
        f"Global_step: {global_step} 程序没有卡住，目前正在生成评估图片，请耐心等待！训练好的权重和评估图片将会自动保存到 {path} 目录下。")

def generate_image(text_encoder, unet, vae, tokenizer, eval_scheduler, args):
    text_encoder.eval()
    temp_pipeline = StableDiffusionPipeline(
        text_encoder=unwrap_model(text_encoder),
        unet=unet,
        vae=vae,
        tokenizer=tokenizer,
        scheduler=eval_scheduler,
        safety_checker=None,
        feature_extractor=None,
        requires_safety_checker=False,
    )
    temp_pipeline.set_progress_bar_config(disable=True)
    all_images = []
    for _ in range(4):
        all_images.append(
            temp_pipeline(args.image_logging_prompt, height=args.height, width=args.width, output_type="numpy").images[0]
        )
    all_images = np.stack(all_images, axis=0)
    text_encoder.train()
    return all_images, temp_pipeline.numpy_to_pil(all_images)

def parse_args():
    parser = argparse.ArgumentParser(
        description="Simple example of a training script.")
    parser.add_argument(
        "--save_steps",
        type=int,
        default=10,
        help="Save learned_embeds.pdparams every X updates steps.",
    )
    parser.add_argument(
        "--image_logging_prompt",
        type=str,
        default=None,
        help="Logging image use which prompt.",
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="CompVis/stable-diffusion-v1-4",
        required=False,
        help="Path to pretrained model or model identifier from local models.",
    )
    parser.add_argument(
        "--tokenizer_name",
        type=str,
        default=None,
        help="Pretrained tokenizer name or path if not the same as model_name",
    )
    parser.add_argument("--train_data_dir",
                        type=str,
                        default=None,
                        required=False,
                        help="A folder containing the training data.")
    parser.add_argument(
        "--placeholder_token",
        type=str,
        default=None,
        required=False,
        help="A token to use as a placeholder for the concept.",
    )
    parser.add_argument("--initializer_token",
                        type=str,
                        default=None,
                        required=False,
                        help="A token to use as initializer word.")
    parser.add_argument("--learnable_property",
                        type=str,
                        default="object",
                        help="Choose between 'object' and 'style'")
    parser.add_argument("--repeats",
                        type=int,
                        default=100,
                        help="How many times to repeat the training data.")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="text-inversion-model",
        help=
        "The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument("--seed",
                        type=int,
                        default=None,
                        help="A seed for reproducible training.")
    parser.add_argument(
        "--height",
        type=int,
        default=512,
        help=
        ("The height for input images, all the images in the train/validation dataset will be resized to this"
         " height"),
    )
    parser.add_argument(
        "--width",
        type=int,
        default=512,
        help=
        ("The width for input images, all the images in the train/validation dataset will be resized to this"
         " width"),
    )
    parser.add_argument(
        "--center_crop",
        action="store_true",
        help="Whether to center crop images before resizing to resolution")
    parser.add_argument(
        "--train_batch_size",
        type=int,
        default=1,
        help="Batch size (per device) for the training dataloader.")
    parser.add_argument("--num_train_epochs", type=int, default=100)
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=5000,
        help=
        "Total number of training steps to perform.  If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=4,
        help=
        "Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-4,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--scale_lr",
        action="store_true",
        default=False,
        help=
        "Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
    )
    parser.add_argument(
        "--lr_scheduler",
        type=str,
        default="constant",
        help=
        ('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
         ' "constant", "constant_with_warmup"]'),
    )
    parser.add_argument(
        "--lr_warmup_steps",
        type=int,
        default=0,
        help="Number of steps for the warmup in the lr scheduler.")
    parser.add_argument("--adam_beta1",
                        type=float,
                        default=0.9,
                        help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2",
                        type=float,
                        default=0.999,
                        help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam_weight_decay",
                        type=float,
                        default=1e-2,
                        help="Weight decay to use.")
    parser.add_argument("--adam_epsilon",
                        type=float,
                        default=1e-08,
                        help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm",
                        default=-1,
                        type=float,
                        help="Max gradient norm.")
    parser.add_argument("--language", default="en", choices=["en", "zh", "zh_en"], help="Model language.")

    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=
        ("[TensorBoard](https://www.tensorflow.org/tensorboard) or [VisualDL](https://www.paddlepaddle.org.cn/paddle/visualdl) log directory. Will default to"
         "*output_dir/logs"),
    )
    parser.add_argument("--writer_type",
                        type=str,
                        default="visualdl",
                        choices=["tensorboard", "visualdl"],
                        help="Log writer type.")

    args = parser.parse_args(args=[])
    return args


imagenet_templates_small = [
    "a photo of a {}",
    "a rendering of a {}",
    "a cropped photo of the {}",
    "the photo of a {}",
    "a photo of a clean {}",
    "a photo of a dirty {}",
    "a dark photo of the {}",
    "a photo of my {}",
    "a photo of the cool {}",
    "a close-up photo of a {}",
    "a bright photo of the {}",
    "a cropped photo of a {}",
    "a photo of the {}",
    "a good photo of the {}",
    "a photo of one {}",
    "a close-up photo of the {}",
    "a rendition of the {}",
    "a photo of the clean {}",
    "a rendition of a {}",
    "a photo of a nice {}",
    "a good photo of a {}",
    "a photo of the nice {}",
    "a photo of the small {}",
    "a photo of the weird {}",
    "a photo of the large {}",
    "a photo of a cool {}",
    "a photo of a small {}",
]

imagenet_style_templates_small = [
    "a painting in the style of {}",
    "a rendering in the style of {}",
    "a cropped painting in the style of {}",
    "the painting in the style of {}",
    "a clean painting in the style of {}",
    "a dirty painting in the style of {}",
    "a dark painting in the style of {}",
    "a picture in the style of {}",
    "a cool painting in the style of {}",
    "a close-up painting in the style of {}",
    "a bright painting in the style of {}",
    "a cropped painting in the style of {}",
    "a good painting in the style of {}",
    "a close-up painting in the style of {}",
    "a rendition in the style of {}",
    "a nice painting in the style of {}",
    "a small painting in the style of {}",
    "a weird painting in the style of {}",
    "a large painting in the style of {}",
]

zh_imagenet_templates_small = [
    "一张{}的照片",
    "{}的渲染",
    "{}裁剪过的照片",
    "一张干净的{}的照片",
    "{}的黑暗照片",
    "我的{}的照片",
    "酷的{}的照片",
    "{}的特写照片",
    "{}的明亮照片",
    "{}的裁剪照片",
    "{}的照片",
    "{}的好照片",
    "一张{}的照片",
    "干净的照片{}",
    "一张漂亮的{}的照片",
    "漂亮的照片{}",
    "一张很酷的照片{}",
    "一张奇怪的照片{}",
]

zh_imagenet_style_templates_small = [
    "一幅{}风格的画",
    "{}风格的渲染",
    "{}风格的裁剪画",
    "{}风格的绘画",
    "{}风格的一幅干净的画",
    "{}风格的黑暗画作",
    "{}风格的图片",
    "{}风格的一幅很酷的画",
    "{}风格的特写画",
    "一幅{}风格的明亮画作",
    "{}风格的一幅好画",
    "{}风格的特写画",
    "{}风格的艺术画",
    "一幅{}风格的漂亮画",
    "一幅{}风格的奇怪的画",
]

class TextualInversionDataset(Dataset):
    def __init__(
        self,
        data_root,
        tokenizer,
        learnable_property="object",  # [object, style]
        height=512,
        width=512,
        repeats=100,
        interpolation="bicubic",
        flip_p=0.5,
        set="train",
        placeholder_token="*",
        center_crop=False,
        language="en",
    ):
        self.data_root = data_root
        self.tokenizer = tokenizer
        self.learnable_property = learnable_property
        self.height = height
        self.width = width
        self.placeholder_token = placeholder_token
        self.center_crop = center_crop
        self.flip_p = flip_p

        if not Path(data_root).exists():
            raise ValueError(f"{data_root} dir doesn't exists.")

        ext = ["png", "jpg", "jpeg", "bmp", "PNG", "JPG", "JPEG", "BMP"]
        self.image_paths = []
        for e in ext:
            self.image_paths.extend(glob.glob(os.path.join(data_root, "*." + e)))

        self.num_images = len(self.image_paths)
        self._length = self.num_images

        if set == "train":
            self._length = self.num_images * repeats

        self.interpolation = {
            "linear": PIL_INTERPOLATION["linear"],
            "bilinear": PIL_INTERPOLATION["bilinear"],
            "bicubic": PIL_INTERPOLATION["bicubic"],
            "lanczos": PIL_INTERPOLATION["lanczos"],
        }[interpolation]

        self.templates = []
        if learnable_property == "style":
            if "en" in language:
                self.templates.extend(imagenet_style_templates_small)
            if "zh" in language:
                self.templates.extend(zh_imagenet_style_templates_small)
        else:
            if "en" in language:
                self.templates.extend(imagenet_templates_small)
            if "zh" in language:
                self.templates.extend(zh_imagenet_templates_small)
        self.flip_transform = RandomHorizontalFlip(prob=self.flip_p)

    def __len__(self):
        return self._length

    def __getitem__(self, i):
        example = {}
        image = Image.open(self.image_paths[i % self.num_images])

        if not image.mode == "RGB":
            image = image.convert("RGB")

        placeholder_string = self.placeholder_token
        text = random.choice(self.templates).format(placeholder_string)

        example["input_ids"] = self.tokenizer(
            text,
            padding="do_not_pad",
            truncation=True,
            max_length=self.tokenizer.model_max_length,
        ).input_ids

        # default to score-sde preprocessing
        img = np.array(image).astype(np.uint8)

        if self.center_crop:
            crop = min(img.shape[0], img.shape[1])
            h, w, = (
                img.shape[0],
                img.shape[1],
            )
            img = img[(h - crop) // 2 : (h + crop) // 2, (w - crop) // 2 : (w + crop) // 2]

        image = Image.fromarray(img)
        image = image.resize((self.width, self.height), resample=self.interpolation)

        image = self.flip_transform(image)
        image = np.array(image).astype(np.uint8)
        image = (image / 127.5 - 1.0).astype(np.float32).transpose([2, 0, 1])

        example["pixel_values"] = image
        return example


def main(args):
    rank = paddle.distributed.get_rank()
    num_processes = paddle.distributed.get_world_size()
    if num_processes > 1:
        paddle.distributed.init_parallel_env()

    # If passed along, set the training seed now.
    if args.seed is not None:
        seed = args.seed + rank
        set_seed(seed)

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    # Load the tokenizer and add the placeholder token as a additional special token
    try:
        if args.tokenizer_name:
            tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
        elif args.pretrained_model_name_or_path:
            tokenizer = AutoTokenizer.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "tokenizer"))
    except:
        tokenizer = XLMRobertaTokenizer.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "tokenizer"))
    # Add the placeholder token in tokenizer
    num_added_tokens = tokenizer.add_tokens(args.placeholder_token)
    if num_added_tokens == 0:
        raise ValueError(f"单词 {args.placeholder_token} 原本就已经存在了哦. 请用一个新的词汇.")

    # Convert the initializer_token, placeholder_token to ids
    token_ids = tokenizer.encode(args.initializer_token,
                                 add_special_tokens=False)["input_ids"]
    # Check if initializer_token is a single token or a sequence of tokens
    if len(token_ids) > 1:
        # raise ValueError("The initializer token must be a single token.")
        print(
            f"用来初始化的 ‘最接近的单词’ 只能是一个简单词, {args.initializer_token} 不可以哟, 因此我们使用随机生成的单词！")

    initializer_token_id = token_ids[0]

    placeholder_token_id = tokenizer.convert_tokens_to_ids(
        args.placeholder_token)

    # Load models and create wrapper for stable diffusion
    if args.text_encoder is None:
        # Load models and create wrapper for stable diffusion
        if "Taiyi-Stable-Diffusion-1B-Chinese-v0.1" in args.pretrained_model_name_or_path:
            model_cls = BertModel
        if "AltDiffusion" in args.pretrained_model_name_or_path:
            model_cls = RobertaSeriesModelWithTransformation
        else:
            model_cls = CLIPTextModel
        text_encoder = model_cls.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "text_encoder"))
    else:
        text_encoder = args.text_encoder

    if args.vae is None:
        vae = AutoencoderKL.from_pretrained(args.model_name, subfolder="vae")
    else:
        vae = args.vae

    if args.unet is None:
        unet = UNet2DConditionModel.from_pretrained(args.model_name,
                                                    subfolder="unet")
    else:
        unet = args.unet

    # Resize the token embeddings as we are adding new special tokens to the tokenizer
    text_encoder.resize_token_embeddings(len(tokenizer))

    noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
    eval_scheduler = PNDMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")

    # Initialise the newly added placeholder token with the embeddings of the initializer token
    with paddle.no_grad():
        token_embeds = text_encoder.get_input_embeddings()
        if len(token_ids) == 1:
            token_embeds.weight[placeholder_token_id] = token_embeds.weight[
                initializer_token_id]

    # Freeze vae and unet
    freeze_params(vae.parameters())
    freeze_params(unet.parameters())

    # Freeze all parameters except for the token embeddings in text encoder
    if isinstance(text_encoder, BertModel):
        # bert text_encoder
        params_to_freeze = itertools.chain(
            text_encoder.encoder.parameters(),
            text_encoder.pooler.parameters(),
            text_encoder.embeddings.position_embeddings.parameters(),
            text_encoder.embeddings.token_type_embeddings.parameters(),
            text_encoder.embeddings.layer_norm.parameters(),
        )
    # Freeze all parameters except for the token embeddings in text encoder
    elif isinstance(text_encoder, RobertaSeriesModelWithTransformation):
        # roberta text_encoder
        params_to_freeze = itertools.chain(
            text_encoder.transformation.parameters(),
            text_encoder.roberta.encoder.parameters(),
            text_encoder.roberta.pooler.parameters(),
            text_encoder.roberta.embeddings.position_embeddings.parameters(),
            text_encoder.roberta.embeddings.token_type_embeddings.parameters(),
            text_encoder.roberta.embeddings.layer_norm.parameters(),
        )
    else:
        # clip text_encoder
        params_to_freeze = itertools.chain(
            text_encoder.text_model.transformer.parameters(),
            text_encoder.text_model.ln_final.parameters(),
            text_encoder.text_model.positional_embedding.parameters(),
        )
    freeze_params(params_to_freeze)

    if args.scale_lr:
        args.learning_rate = (args.learning_rate *
                              args.gradient_accumulation_steps *
                              args.train_batch_size * num_processes)

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        learning_rate=args.learning_rate,
        num_warmup_steps=args.lr_warmup_steps *
        args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps *
        args.gradient_accumulation_steps,
    )

    # Initialize the optimizer
    optimizer = AdamW(learning_rate=lr_scheduler,
                      parameters=text_encoder.get_input_embeddings().parameters(),
                      beta1=args.adam_beta1,
                      beta2=args.adam_beta2,
                      weight_decay=args.adam_weight_decay,
                      epsilon=args.adam_epsilon,
                      grad_clip=nn.ClipGradByGlobalNorm(args.max_grad_norm) if args.max_grad_norm > 0 else None)

    if num_processes > 1:
        text_encoder = paddle.DataParallel(text_encoder)

    train_dataset = TextualInversionDataset(
        data_root=args.train_data_dir,
        tokenizer=tokenizer,
        height=args.height,
        width=args.width,
        placeholder_token=args.placeholder_token,
        repeats=args.repeats,
        learnable_property=args.learnable_property,
        center_crop=args.center_crop,
        set="train",
        language=args.language,
    )


    def collate_fn(examples):
        input_ids = [example["input_ids"] for example in examples]
        pixel_values = paddle.to_tensor([example["pixel_values"] for example in examples], dtype="float32")
        input_ids = tokenizer.pad(
            {"input_ids": input_ids}, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pd"
        ).input_ids
        batch = {
            "input_ids": input_ids,
            "pixel_values": pixel_values,
        }
        return batch

    train_sampler = DistributedBatchSampler(
        train_dataset, batch_size=args.train_batch_size,
        shuffle=True) if num_processes > 1 else BatchSampler(
            train_dataset, batch_size=args.train_batch_size, shuffle=True)
    train_dataloader = DataLoader(train_dataset,
                                  batch_sampler=train_sampler,
                                  collate_fn=collate_fn)

    # Scheduler and math around the number of training steps.
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps /
                                      num_update_steps_per_epoch)

    if rank == 0:
        writer = get_writer(args)

    # Train!
    total_batch_size = args.train_batch_size * args.gradient_accumulation_steps * num_processes
    progress_bar = tqdm(range(args.max_train_steps), disable=rank > 0)
    progress_bar.set_description("Train Steps")
    global_step = 0

    text_encoder_embedding_clone = unwrap_model(
        text_encoder).get_input_embeddings().weight.clone()

    # Keep vae and unet in eval model as we don't train these
    vae.eval()
    unet.eval()
    text_encoder.train()
    try:
        for epoch in range(args.num_train_epochs):
            for step, batch in enumerate(train_dataloader):
                # Convert images to latent space
                latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
                latents = latents * 0.18215

                # Sample noise that we'll add to the latents
                noise = paddle.randn(latents.shape)
                batch_size = latents.shape[0]
                # Sample a random timestep for each image
                timesteps = paddle.randint(0, noise_scheduler.config.num_train_timesteps, (batch_size,), dtype="int64")

                # Add noise to the latents according to the noise magnitude at each timestep
                # (this is the forward diffusion process)
                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

                # Get the text embedding for conditioning
                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
                # Predict the unet output
                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample

                # Get the target for loss depending on the prediction type
                if noise_scheduler.config.prediction_type == "epsilon":
                    target = noise
                elif noise_scheduler.config.prediction_type == "v_prediction":
                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
                else:
                    raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")

                loss = F.mse_loss(model_pred, target, reduction="none").mean([1, 2, 3]).mean()
                if args.gradient_accumulation_steps > 1:
                    loss = loss / args.gradient_accumulation_steps
                loss.backward()

                with paddle.no_grad():
                    # Get the index for tokens that we want to zero the grads for
                    index_grads_to_zero = (paddle.arange(
                        len(tokenizer)) == placeholder_token_id
                                           ).astype("float32").unsqueeze(-1)
                    unwrap_model(text_encoder).get_input_embeddings(
                    ).weight.grad = unwrap_model(
                        text_encoder).get_input_embeddings(
                        ).weight.grad * index_grads_to_zero

                if (step + 1) % args.gradient_accumulation_steps == 0:
                    optimizer.step()
                    with paddle.no_grad():
                        unwrap_model(text_encoder).get_input_embeddings(
                        ).weight[:-1] = text_encoder_embedding_clone[:-1]

                    lr_scheduler.step()
                    optimizer.clear_grad()
                    progress_bar.update(1)
                    global_step += 1
                    logs = {
                        "epoch":
                        str(epoch).zfill(4),
                        "step_loss":
                        round(loss.item() * args.gradient_accumulation_steps,
                              10),
                        "lr":
                        lr_scheduler.get_lr()
                    }
                    progress_bar.set_postfix(**logs)
                    if rank == 0:
                        for name, val in logs.items():
                            if name == "epoch": continue
                            writer.add_scalar(f"train/{name}",
                                              val,
                                              step=global_step)
                        if global_step % args.save_steps == 0:
                            save_progress(text_encoder, placeholder_token_id,
                                          args, global_step)
                            images, pil_images = generate_image(text_encoder, unet, vae, tokenizer, eval_scheduler, args)
                            writer.add_image("images", images, step=global_step, dataformats="NHWC")
                            name = args.placeholder_token
                            name = name.translate({
                                92: 95,
                                47: 95,
                                42: 95,
                                34: 95,
                                58: 95,
                                63: 95,
                                60: 95,
                                62: 95,
                                124: 95
                            })
                            image_grid(pil_images).save(os.path.join(args.output_dir, "step-"+str(global_step), f"{name}.jpg"))


                if global_step >= args.max_train_steps:
                    break

        if rank == 0:
            writer.close()
            save_progress(text_encoder, placeholder_token_id, args, global_step)
            print(f'训练完毕, 可以用新词 {args.placeholder_token} 去生成图片了.')
            del text_encoder
            del optimizer
            del vae
            del unet
            del text_encoder_embedding_clone
            gc.collect()
    except:
        save_progress(text_encoder, placeholder_token_id, args, global_step)
        del text_encoder
        del optimizer
        del vae
        del unet
        del text_encoder_embedding_clone
        gc.collect()