machinelearnear/dreambooth-training

1import argparse2import itertools3import math4import os5from pathlib import Path6from typing import Optional7import subprocess8import sys9import gc10import random11 12import torch13import torch.nn.functional as F14import torch.utils.checkpoint15from torch.utils.data import Dataset16 17from accelerate import Accelerator18from accelerate.logging import get_logger19from accelerate.utils import set_seed20from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel21from diffusers.optimization import get_scheduler22from huggingface_hub import HfFolder, Repository, whoami23from PIL import Image24from torchvision import transforms25from tqdm.auto import tqdm26from transformers import CLIPTextModel, CLIPTokenizer27 28 29logger = get_logger(__name__)30 31 32def parse_args():33    parser = argparse.ArgumentParser(description="Simple example of a training script.")34    parser.add_argument(35        "--pretrained_model_name_or_path",36        type=str,37        default=None,38        #required=True,39        help="Path to pretrained model or model identifier from huggingface.co/models.",40    )41    parser.add_argument(42        "--tokenizer_name",43        type=str,44        default=None,45        help="Pretrained tokenizer name or path if not the same as model_name",46    )47    parser.add_argument(48        "--instance_data_dir",49        type=str,50        default=None,51        #required=True,52        help="A folder containing the training data of instance images.",53    )54    parser.add_argument(55        "--class_data_dir",56        type=str,57        default=None,58        #required=False,59        help="A folder containing the training data of class images.",60    )61    parser.add_argument(62        "--instance_prompt",63        type=str,64        default=None,65        help="The prompt with identifier specifying the instance",66    )67    parser.add_argument(68        "--class_prompt",69        type=str,70        default="",71        help="The prompt to specify images in the same class as provided instance images.",72    )73    parser.add_argument(74        "--with_prior_preservation",75        default=False,76        action="store_true",77        help="Flag to add prior preservation loss.",78    )79    parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.")80    parser.add_argument(81        "--num_class_images",82        type=int,83        default=100,84        help=(85            "Minimal class images for prior preservation loss. If not have enough images, additional images will be"86            " sampled with class_prompt."87        ),88    )89    parser.add_argument(90        "--output_dir",91        type=str,92        default="",93        help="The output directory where the model predictions and checkpoints will be written.",94    )95    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")96    parser.add_argument(97        "--resolution",98        type=int,99        default=512,100        help=(101            "The resolution for input images, all the images in the train/validation dataset will be resized to this"102            " resolution"103        ),104    )105    parser.add_argument(106        "--center_crop", action="store_true", help="Whether to center crop images before resizing to resolution"107    )108    parser.add_argument("--train_text_encoder", action="store_true", help="Whether to train the text encoder")109    parser.add_argument(110        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."111    )112    parser.add_argument(113        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."114    )115    parser.add_argument("--num_train_epochs", type=int, default=1)116    parser.add_argument(117        "--max_train_steps",118        type=int,119        default=None,120        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",121    )122    parser.add_argument(123        "--gradient_accumulation_steps",124        type=int,125        default=1,126        help="Number of updates steps to accumulate before performing a backward/update pass.",127    )128    parser.add_argument(129        "--gradient_checkpointing",130        action="store_true",131        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",132    )133    parser.add_argument(134        "--learning_rate",135        type=float,136        default=5e-6,137        help="Initial learning rate (after the potential warmup period) to use.",138    )139    parser.add_argument(140        "--scale_lr",141        action="store_true",142        default=False,143        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",144    )145    parser.add_argument(146        "--lr_scheduler",147        type=str,148        default="constant",149        help=(150            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'151            ' "constant", "constant_with_warmup"]'152        ),153    )154    parser.add_argument(155        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."156    )157    parser.add_argument(158        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."159    )160    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")161    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")162    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")163    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")164    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")165    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")166    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")167    parser.add_argument(168        "--hub_model_id",169        type=str,170        default=None,171        help="The name of the repository to keep in sync with the local `output_dir`.",172    )173    parser.add_argument(174        "--logging_dir",175        type=str,176        default="logs",177        help=(178            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"179            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."180        ),181    )182    parser.add_argument(183        "--mixed_precision",184        type=str,185        default="no",186        choices=["no", "fp16", "bf16"],187        help=(188            "Whether to use mixed precision. Choose"189            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."190            "and an Nvidia Ampere GPU."191        ),192    )193 194    parser.add_argument(195        "--save_n_steps",196        type=int,197        default=1,198        help=("Save the model every n global_steps"),199    )200    201    202    parser.add_argument(203        "--save_starting_step",204        type=int,205        default=1,206        help=("The step from which it starts saving intermediary checkpoints"),207    )208    209    parser.add_argument(210        "--stop_text_encoder_training",211        type=int,212        default=1000000,213        help=("The step at which the text_encoder is no longer trained"),214    )215 216 217    parser.add_argument(218        "--image_captions_filename",219        action="store_true",220        help="Get captions from filename",221    )    222    223    224    parser.add_argument(225        "--dump_only_text_encoder",226        action="store_true",227        default=False,        228        help="Dump only text encoder",229    )230 231    parser.add_argument(232        "--train_only_unet",233        action="store_true",234        default=False,        235        help="Train only the unet",236    )237    238    parser.add_argument(239        "--cache_latents",240        action="store_true",241        default=False,        242        help="Train only the unet",243    )244    245    parser.add_argument(246        "--Session_dir",247        type=str,248        default="",     249        help="Current session directory",250    )    251 252    253    254 255    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")256 257    args = parser.parse_args()258    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))259    if env_local_rank != -1 and env_local_rank != args.local_rank:260        args.local_rank = env_local_rank261 262    #if args.instance_data_dir is None:263    #    raise ValueError("You must specify a train data directory.")264 265    #if args.with_prior_preservation:266    #    if args.class_data_dir is None:267    #        raise ValueError("You must specify a data directory for class images.")268    #    if args.class_prompt is None:269    #        raise ValueError("You must specify prompt for class images.")270 271    return args272 273 274class DreamBoothDataset(Dataset):275    """276    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.277    It pre-processes the images and the tokenizes prompts.278    """279 280    def __init__(281        self,282        instance_data_root,283        instance_prompt,284        tokenizer,285        args,286        class_data_root=None,287        class_prompt=None,288        size=512,289        center_crop=False,290    ):291        self.size = size292        self.center_crop = center_crop293        self.tokenizer = tokenizer294        self.image_captions_filename = None295 296        self.instance_data_root = Path(instance_data_root)297        if not self.instance_data_root.exists():298            raise ValueError("Instance images root doesn't exists.")299 300        self.instance_images_path = list(Path(instance_data_root).iterdir())301        self.num_instance_images = len(self.instance_images_path)302        self.instance_prompt = instance_prompt303        self._length = self.num_instance_images304 305        if args.image_captions_filename:306            self.image_captions_filename = True307        308        if class_data_root is not None:309            self.class_data_root = Path(class_data_root)310            self.class_data_root.mkdir(parents=True, exist_ok=True)311            self.class_images_path = list(self.class_data_root.iterdir())312            random.shuffle(self.class_images_path)313            self.num_class_images = len(self.class_images_path)314            self._length = max(self.num_class_images, self.num_instance_images)315            self.class_prompt = class_prompt316        else:317            self.class_data_root = None318 319        self.image_transforms = transforms.Compose(320            [321                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),322                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),323                transforms.ToTensor(),324                transforms.Normalize([0.5], [0.5]),325            ]326        )327 328    def __len__(self):329        return self._length330 331    def __getitem__(self, index):332        example = {}333        path = self.instance_images_path[index % self.num_instance_images]334        instance_image = Image.open(path)335        if not instance_image.mode == "RGB":336            instance_image = instance_image.convert("RGB")337            338        instance_prompt = self.instance_prompt339        340        if self.image_captions_filename:341            filename = Path(path).stem342            pt=''.join([i for i in filename if not i.isdigit()])343            pt=pt.replace("_"," ")344            pt=pt.replace("(","")345            pt=pt.replace(")","")346            pt=pt.replace("-","")347            instance_prompt = pt348            sys.stdout.write(" " +instance_prompt+" ")349            sys.stdout.flush()350 351 352        example["instance_images"] = self.image_transforms(instance_image)353        example["instance_prompt_ids"] = self.tokenizer(354            instance_prompt,355            padding="do_not_pad",356            truncation=True,357            max_length=self.tokenizer.model_max_length,358        ).input_ids359 360        if self.class_data_root:361            class_image = Image.open(self.class_images_path[index % self.num_class_images])362            if not class_image.mode == "RGB":363                class_image = class_image.convert("RGB")364            example["class_images"] = self.image_transforms(class_image)365            example["class_prompt_ids"] = self.tokenizer(366                self.class_prompt,367                padding="do_not_pad",368                truncation=True,369                max_length=self.tokenizer.model_max_length,370            ).input_ids371 372        return example373 374 375 376class PromptDataset(Dataset):377    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."378 379    def __init__(self, prompt, num_samples):380        self.prompt = prompt381        self.num_samples = num_samples382 383    def __len__(self):384        return self.num_samples385 386    def __getitem__(self, index):387        example = {}388        example["prompt"] = self.prompt389        example["index"] = index390        return example391 392class LatentsDataset(Dataset):393    def __init__(self, latents_cache, text_encoder_cache):394        self.latents_cache = latents_cache395        self.text_encoder_cache = text_encoder_cache396 397    def __len__(self):398        return len(self.latents_cache)399 400    def __getitem__(self, index):401        return self.latents_cache[index], self.text_encoder_cache[index]402 403def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):404    if token is None:405        token = HfFolder.get_token()406    if organization is None:407        username = whoami(token)["name"]408        return f"{username}/{model_id}"409    else:410        return f"{organization}/{model_id}"411 412def merge_two_dicts(starting_dict: dict, updater_dict: dict) -> dict:413    """414    Starts from base starting dict and then adds the remaining key values from updater replacing the values from415    the first starting/base dict with the second updater dict.416 417    For later: how does d = {**d1, **d2} replace collision?418 419    :param starting_dict:420    :param updater_dict:421    :return:422    """423    new_dict: dict = starting_dict.copy()   # start with keys and values of starting_dict424    new_dict.update(updater_dict)    # modifies starting_dict with keys and values of updater_dict425    return new_dict426 427def merge_args(args1: argparse.Namespace, args2: argparse.Namespace) -> argparse.Namespace:428    """429 430    ref: https://stackoverflow.com/questions/56136549/how-can-i-merge-two-argparse-namespaces-in-python-2-x431    :param args1:432    :param args2:433    :return:434    """435    # - the merged args436    # The vars() function returns the __dict__ attribute to values of the given object e.g {field:value}.437    merged_key_values_for_namespace: dict = merge_two_dicts(vars(args1), vars(args2))438    args = argparse.Namespace(**merged_key_values_for_namespace)439    return args440 441def run_training(args_imported):442    args_default = parse_args()443    args = merge_args(args_default, args_imported)444    print(args)445    logging_dir = Path(args.output_dir, args.logging_dir)446    i=args.save_starting_step447    accelerator = Accelerator(448        gradient_accumulation_steps=args.gradient_accumulation_steps,449        mixed_precision=args.mixed_precision,450        log_with="tensorboard",451        logging_dir=logging_dir,452    )453 454    # Currently, it's not possible to do gradient accumulation when training two models with accelerate.accumulate455    # This will be enabled soon in accelerate. For now, we don't allow gradient accumulation when training two models.456    # TODO (patil-suraj): Remove this check when gradient accumulation with two models is enabled in accelerate.457    if args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1:458        raise ValueError(459            "Gradient accumulation is not supported when training the text encoder in distributed training. "460            "Please set gradient_accumulation_steps to 1. This feature will be supported in the future."461        )462 463    if args.seed is not None:464        set_seed(args.seed)465 466    if args.with_prior_preservation:467        class_images_dir = Path(args.class_data_dir)468        if not class_images_dir.exists():469            class_images_dir.mkdir(parents=True)470        cur_class_images = len(list(class_images_dir.iterdir()))471 472        if cur_class_images < args.num_class_images:473            torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32474            pipeline = StableDiffusionPipeline.from_pretrained(475                args.pretrained_model_name_or_path, torch_dtype=torch_dtype476            )477            pipeline.set_progress_bar_config(disable=True)478 479            num_new_images = args.num_class_images - cur_class_images480            logger.info(f"Number of class images to sample: {num_new_images}.")481 482            sample_dataset = PromptDataset(args.class_prompt, num_new_images)483            sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)484 485            sample_dataloader = accelerator.prepare(sample_dataloader)486            pipeline.to(accelerator.device)487 488            for example in tqdm(489                sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process490            ):491                with torch.autocast("cuda"):                492                    images = pipeline(example["prompt"]).images493 494                for i, image in enumerate(images):495                    image.save(class_images_dir / f"{example['index'][i] + cur_class_images}.jpg")496 497            del pipeline498            if torch.cuda.is_available():499                torch.cuda.empty_cache()500 501    # Handle the repository creation502    if accelerator.is_main_process:503        if args.push_to_hub:504            if args.hub_model_id is None:505                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)506            else:507                repo_name = args.hub_model_id508            repo = Repository(args.output_dir, clone_from=repo_name)509 510            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:511                if "step_*" not in gitignore:512                    gitignore.write("step_*\n")513                if "epoch_*" not in gitignore:514                    gitignore.write("epoch_*\n")515        elif args.output_dir is not None:516            os.makedirs(args.output_dir, exist_ok=True)517 518    # Load the tokenizer519    if args.tokenizer_name:520        tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)521    elif args.pretrained_model_name_or_path:522        tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")523 524    # Load models and create wrapper for stable diffusion525    if args.train_only_unet:526      if os.path.exists(str(args.output_dir+"/text_encoder_trained")):527        text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder_trained")528      elif os.path.exists(str(args.output_dir+"/text_encoder")):529        text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder")530      else:531        text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")532    else:533      text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")534    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")535    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")536 537    vae.requires_grad_(False)538    if not args.train_text_encoder:539        text_encoder.requires_grad_(False)540 541    if args.gradient_checkpointing:542        unet.enable_gradient_checkpointing()543        if args.train_text_encoder:544            text_encoder.gradient_checkpointing_enable()545 546    if args.scale_lr:547        args.learning_rate = (548            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes549        )550 551    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs552    if args.use_8bit_adam:553        try:554            import bitsandbytes as bnb555        except ImportError:556            raise ImportError(557                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."558            )559 560        optimizer_class = bnb.optim.AdamW8bit561    else:562        optimizer_class = torch.optim.AdamW563 564    params_to_optimize = (565        itertools.chain(unet.parameters(), text_encoder.parameters()) if args.train_text_encoder else unet.parameters()566    )567    optimizer = optimizer_class(568        params_to_optimize,569        lr=args.learning_rate,570        betas=(args.adam_beta1, args.adam_beta2),571        weight_decay=args.adam_weight_decay,572        eps=args.adam_epsilon,573    )574 575    noise_scheduler = DDPMScheduler.from_config(args.pretrained_model_name_or_path, subfolder="scheduler")576 577    train_dataset = DreamBoothDataset(578        instance_data_root=args.instance_data_dir,579        instance_prompt=args.instance_prompt,580        class_data_root=args.class_data_dir if args.with_prior_preservation else None,581        class_prompt=args.class_prompt,582        tokenizer=tokenizer,583        size=args.resolution,584        center_crop=args.center_crop,585        args=args,586    )587 588    def collate_fn(examples):589        input_ids = [example["instance_prompt_ids"] for example in examples]590        pixel_values = [example["instance_images"] for example in examples]591 592        # Concat class and instance examples for prior preservation.593        # We do this to avoid doing two forward passes.594        if args.with_prior_preservation:595            input_ids += [example["class_prompt_ids"] for example in examples]596            pixel_values += [example["class_images"] for example in examples]597 598        pixel_values = torch.stack(pixel_values)599        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()600 601        input_ids = tokenizer.pad({"input_ids": input_ids}, padding=True, return_tensors="pt").input_ids602 603        batch = {604            "input_ids": input_ids,605            "pixel_values": pixel_values,606        }607        return batch608 609    train_dataloader = torch.utils.data.DataLoader(610        train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn611    )612 613    # Scheduler and math around the number of training steps.614    overrode_max_train_steps = False615    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)616    if args.max_train_steps is None:617        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch618        overrode_max_train_steps = True619 620    lr_scheduler = get_scheduler(621        args.lr_scheduler,622        optimizer=optimizer,623        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,624        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,625    )626 627    if args.train_text_encoder:628        unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(629            unet, text_encoder, optimizer, train_dataloader, lr_scheduler630        )631    else:632        unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(633            unet, optimizer, train_dataloader, lr_scheduler634        )635 636    weight_dtype = torch.float32637    if args.mixed_precision == "fp16":638        weight_dtype = torch.float16639    elif args.mixed_precision == "bf16":640        weight_dtype = torch.bfloat16641 642    # Move text_encode and vae to gpu.643    # For mixed precision training we cast the text_encoder and vae weights to half-precision644    # as these models are only used for inference, keeping weights in full precision is not required.645    vae.to(accelerator.device, dtype=weight_dtype)646    if not args.train_text_encoder:647        text_encoder.to(accelerator.device, dtype=weight_dtype)648 649 650    if args.cache_latents:651        latents_cache = []652        text_encoder_cache = []653        for batch in tqdm(train_dataloader, desc="Caching latents"):654            with torch.no_grad():655                batch["pixel_values"] = batch["pixel_values"].to(accelerator.device, non_blocking=True, dtype=weight_dtype)656                batch["input_ids"] = batch["input_ids"].to(accelerator.device, non_blocking=True)657                latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist)658                if args.train_text_encoder:659                    text_encoder_cache.append(batch["input_ids"])660                else:661                    text_encoder_cache.append(text_encoder(batch["input_ids"])[0])662        train_dataset = LatentsDataset(latents_cache, text_encoder_cache)663        train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, collate_fn=lambda x: x, shuffle=True)664 665        del vae666        #if not args.train_text_encoder:667        #    del text_encoder668        if torch.cuda.is_available():669            torch.cuda.empty_cache()670 671    # We need to recalculate our total training steps as the size of the training dataloader may have changed.672    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)673    if overrode_max_train_steps:674        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch675    # Afterwards we recalculate our number of training epochs676    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)677 678    # We need to initialize the trackers we use, and also store our configuration.679    # The trackers initializes automatically on the main process.680    if accelerator.is_main_process:681        accelerator.init_trackers("dreambooth", config=vars(args))682 683    def bar(prg):684       br='|'+'█' * prg + ' ' * (25-prg)+'|'685       return br686 687    # Train!688    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps689 690    logger.info("***** Running training *****")691    logger.info(f"  Num examples = {len(train_dataset)}")692    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")693    logger.info(f"  Num Epochs = {args.num_train_epochs}")694    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")695    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")696    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")697    logger.info(f"  Total optimization steps = {args.max_train_steps}")698    # Only show the progress bar once on each machine.699    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)700    global_step = 0701 702    for epoch in range(args.num_train_epochs):703        unet.train()704        if args.train_text_encoder:705            text_encoder.train()706        for step, batch in enumerate(train_dataloader):707            with accelerator.accumulate(unet):708                # Convert images to latent space709                with torch.no_grad():710                    if args.cache_latents:711                        latents_dist = batch[0][0]712                    else:713                        latents_dist = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist714                    latents = latents_dist.sample() * 0.18215715 716                # Sample noise that we'll add to the latents717                noise = torch.randn_like(latents)718                bsz = latents.shape[0]719                # Sample a random timestep for each image720                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)721                timesteps = timesteps.long()722 723                # Add noise to the latents according to the noise magnitude at each timestep724                # (this is the forward diffusion process)725                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)726 727                # Get the text embedding for conditioning728                if(args.cache_latents):729                    if args.train_text_encoder:730                        encoder_hidden_states = text_encoder(batch[0][1])[0]731                    else:732                        encoder_hidden_states = batch[0][1]733                else:734                    encoder_hidden_states = text_encoder(batch["input_ids"])[0]735 736                # Predict the noise residual737                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample738                739                # Get the target for loss depending on the prediction type740                if noise_scheduler.config.prediction_type == "epsilon":741                    target = noise742                elif noise_scheduler.config.prediction_type == "v_prediction":743                    target = noise_scheduler.get_velocity(latents, noise, timesteps)744                else:745                    raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")746                747                if args.with_prior_preservation:748                    # Chunk the noise and model_pred into two parts and compute the loss on each part separately.749                    model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)750                    target, target_prior = torch.chunk(target, 2, dim=0)751 752                    # Compute instance loss753                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()754 755                    # Compute prior loss756                    prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")757 758                    # Add the prior loss to the instance loss.759                    loss = loss + args.prior_loss_weight * prior_loss760                else:761                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")762 763                accelerator.backward(loss)764                if accelerator.sync_gradients:765                    params_to_clip = (766                        itertools.chain(unet.parameters(), text_encoder.parameters())767                        if args.train_text_encoder768                        else unet.parameters()769                    )770                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)771                optimizer.step()772                lr_scheduler.step()773                optimizer.zero_grad()774 775            # Checks if the accelerator has performed an optimization step behind the scenes776            if accelerator.sync_gradients:777                progress_bar.update(1)778                global_step += 1779 780            fll=round((global_step*100)/args.max_train_steps)781            fll=round(fll/4)782            pr=bar(fll)783            784            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}785            progress_bar.set_postfix(**logs)786            progress_bar.set_description_str("Progress:"+pr)787            accelerator.log(logs, step=global_step)788 789            if global_step >= args.max_train_steps:790                break791 792            if args.train_text_encoder and global_step == args.stop_text_encoder_training and global_step >= 30:793              if accelerator.is_main_process:794                print(" " +" Freezing the text_encoder ..."+" ")                795                frz_dir=args.output_dir + "/text_encoder_frozen"796                if os.path.exists(frz_dir):797                  subprocess.call('rm -r '+ frz_dir, shell=True)798                os.mkdir(frz_dir)799                pipeline = StableDiffusionPipeline.from_pretrained(800                    args.pretrained_model_name_or_path,801                    unet=accelerator.unwrap_model(unet),802                    text_encoder=accelerator.unwrap_model(text_encoder),803                )804                pipeline.text_encoder.save_pretrained(frz_dir)805                         806            if args.save_n_steps >= 200:807               if global_step < args.max_train_steps and global_step+1==i:808                  ckpt_name = "_step_" + str(global_step+1)809                  save_dir = Path(args.output_dir+ckpt_name)810                  save_dir=str(save_dir)811                  save_dir=save_dir.replace(" ", "_")                    812                  if not os.path.exists(save_dir):813                     os.mkdir(save_dir)814                  inst=save_dir[16:]815                  inst=inst.replace(" ", "_")816                  print(" SAVING CHECKPOINT: "+args.Session_dir+"/"+inst+".ckpt")817                  # Create the pipeline using the trained modules and save it.818                  if accelerator.is_main_process:819                     pipeline = StableDiffusionPipeline.from_pretrained(820                           args.pretrained_model_name_or_path,821                           unet=accelerator.unwrap_model(unet),822                           text_encoder=accelerator.unwrap_model(text_encoder),823                     )824                     pipeline.save_pretrained(save_dir)825                     frz_dir=args.output_dir + "/text_encoder_frozen"                    826                     if args.train_text_encoder and os.path.exists(frz_dir):827                        subprocess.call('rm -r '+save_dir+'/text_encoder/*.*', shell=True)828                        subprocess.call('cp -f '+frz_dir +'/*.* '+ save_dir+'/text_encoder', shell=True)                     829                     chkpth=args.Session_dir+"/"+inst+".ckpt"830                     subprocess.call('python /content/diffusers/scripts/convert_diffusers_to_original_stable_diffusion.py --model_path ' + save_dir + ' --checkpoint_path ' + chkpth + ' --half', shell=True)831                     subprocess.call('rm -r '+ save_dir, shell=True)832                     i=i+args.save_n_steps833            834        accelerator.wait_for_everyone()835 836    # Create the pipeline using using the trained modules and save it.837    if accelerator.is_main_process:838      if args.dump_only_text_encoder:839         txt_dir=args.output_dir + "/text_encoder_trained"840         if not os.path.exists(txt_dir):841           os.mkdir(txt_dir)842         pipeline = StableDiffusionPipeline.from_pretrained(843             args.pretrained_model_name_or_path,844             unet=accelerator.unwrap_model(unet),845             text_encoder=accelerator.unwrap_model(text_encoder),846         )847         pipeline.text_encoder.save_pretrained(txt_dir)       848 849      elif args.train_only_unet:850        pipeline = StableDiffusionPipeline.from_pretrained(851            args.pretrained_model_name_or_path,852            unet=accelerator.unwrap_model(unet),853            text_encoder=accelerator.unwrap_model(text_encoder),854        )855        pipeline.save_pretrained(args.output_dir)856        txt_dir=args.output_dir + "/text_encoder_trained"857        subprocess.call('rm -r '+txt_dir, shell=True)858     859      else:860        pipeline = StableDiffusionPipeline.from_pretrained(861            args.pretrained_model_name_or_path,862            unet=accelerator.unwrap_model(unet),863            text_encoder=accelerator.unwrap_model(text_encoder),864        )865        frz_dir=args.output_dir + "/text_encoder_frozen"866        pipeline.save_pretrained(args.output_dir)867        if args.train_text_encoder and os.path.exists(frz_dir):868           subprocess.call('mv -f '+frz_dir +'/*.* '+ args.output_dir+'/text_encoder', shell=True)869           subprocess.call('rm -r '+ frz_dir, shell=True) 870 871        if args.push_to_hub:872            repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)873 874    accelerator.end_training()875    del pipeline876    torch.cuda.empty_cache()877    gc.collect()878if __name__ == "__main__":879    pass880    #main()881 882