train_unconditional.py 29.7 KB
Newer Older
anton-l's avatar
anton-l committed
1
import argparse
2
import inspect
3
import logging
4
import math
anton-l's avatar
anton-l committed
5
import os
6
import shutil
7
from datetime import timedelta
8
9
from pathlib import Path
from typing import Optional
anton-l's avatar
anton-l committed
10

11
import accelerate
12
import datasets
13
14
import torch
import torch.nn.functional as F
15
from accelerate import Accelerator, InitProcessGroupKwargs
16
from accelerate.logging import get_logger
17
from accelerate.utils import ProjectConfiguration
anton-l's avatar
anton-l committed
18
from datasets import load_dataset
19
from huggingface_hub import HfFolder, Repository, create_repo, whoami
20
from packaging import version
21
from torchvision import transforms
anton-l's avatar
anton-l committed
22
from tqdm.auto import tqdm
anton-l's avatar
anton-l committed
23

24
25
26
27
import diffusers
from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel
28
from diffusers.utils import check_min_version, is_accelerate_version, is_tensorboard_available, is_wandb_available
29
from diffusers.utils.import_utils import is_xformers_available
30

anton-l's avatar
anton-l committed
31

32
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
33
check_min_version("0.22.0.dev0")
34

35
logger = get_logger(__name__, log_level="INFO")
anton-l's avatar
anton-l committed
36
37


38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def _extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.

    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    if not isinstance(arr, torch.Tensor):
        arr = torch.from_numpy(arr)
    res = arr[timesteps].float().to(timesteps.device)
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res.expand(broadcast_shape)


56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def parse_args():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default=None,
        help=(
            "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
            " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
            " or to a folder containing files that HF Datasets can understand."
        ),
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default=None,
        help="The config of the Dataset, leave as None if there's only one config.",
    )
74
75
76
77
78
79
    parser.add_argument(
        "--model_config_name_or_path",
        type=str,
        default=None,
        help="The config of the UNet model to train, leave as None to use standard DDPM configuration.",
    )
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    parser.add_argument(
        "--train_data_dir",
        type=str,
        default=None,
        help=(
            "A folder containing the training data. Folder contents must follow the structure described in"
            " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
            " must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
        ),
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="ddpm-model-64",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument("--overwrite_output_dir", action="store_true")
    parser.add_argument(
        "--cache_dir",
        type=str,
        default=None,
        help="The directory where the downloaded models and datasets will be stored.",
    )
    parser.add_argument(
        "--resolution",
        type=int,
        default=64,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
    parser.add_argument(
        "--center_crop",
        default=False,
        action="store_true",
        help=(
            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
            " cropped. The images will be resized to the resolution first before cropping."
        ),
    )
    parser.add_argument(
        "--random_flip",
        default=False,
        action="store_true",
        help="whether to randomly flip images horizontally",
    )
127
128
129
130
    parser.add_argument(
        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
    )
    parser.add_argument(
131
132
133
134
135
136
137
138
139
140
        "--eval_batch_size", type=int, default=16, help="The number of images to generate for evaluation."
    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=0,
        help=(
            "The number of subprocesses to use for data loading. 0 means that the data will be loaded in the main"
            " process."
        ),
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
    )
    parser.add_argument("--num_epochs", type=int, default=100)
    parser.add_argument("--save_images_epochs", type=int, default=10, help="How often to save images during training.")
    parser.add_argument(
        "--save_model_epochs", type=int, default=10, help="How often to save the model during training."
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        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(
        "--lr_scheduler",
        type=str,
        default="cosine",
        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=500, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--adam_beta1", type=float, default=0.95, 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-6, help="Weight decay magnitude for the Adam optimizer."
    )
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer.")
    parser.add_argument(
        "--use_ema",
        action="store_true",
        help="Whether to use Exponential Moving Average for the final model weights.",
    )
    parser.add_argument("--ema_inv_gamma", type=float, default=1.0, help="The inverse gamma value for the EMA decay.")
    parser.add_argument("--ema_power", type=float, default=3 / 4, help="The power value for the EMA decay.")
    parser.add_argument("--ema_max_decay", type=float, default=0.9999, help="The maximum decay magnitude for EMA.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--hub_model_id",
        type=str,
        default=None,
        help="The name of the repository to keep in sync with the local `output_dir`.",
    )
    parser.add_argument(
        "--hub_private_repo", action="store_true", help="Whether or not to create a private repository."
    )
196
197
198
199
200
201
202
203
204
205
    parser.add_argument(
        "--logger",
        type=str,
        default="tensorboard",
        choices=["tensorboard", "wandb"],
        help=(
            "Whether to use [tensorboard](https://www.tensorflow.org/tensorboard) or [wandb](https://www.wandb.ai)"
            " for experiment tracking and logging of model metrics and model checkpoints"
        ),
    )
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="no",
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose"
            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
            "and an Nvidia Ampere GPU."
        ),
    )
227
    parser.add_argument(
228
229
230
231
        "--prediction_type",
        type=str,
        default="epsilon",
        choices=["epsilon", "sample"],
232
        help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.",
233
234
    )
    parser.add_argument("--ddpm_num_steps", type=int, default=1000)
235
    parser.add_argument("--ddpm_num_inference_steps", type=int, default=1000)
236
    parser.add_argument("--ddpm_beta_schedule", type=str, default="linear")
237
238
239
240
241
242
243
244
245
    parser.add_argument(
        "--checkpointing_steps",
        type=int,
        default=500,
        help=(
            "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
            " training using `--resume_from_checkpoint`."
        ),
    )
246
    parser.add_argument(
247
        "--checkpoints_total_limit",
248
249
        type=int,
        default=None,
250
        help=("Max number of checkpoints to store."),
251
    )
252
253
254
255
256
257
258
259
260
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help=(
            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
        ),
    )
261
262
263
    parser.add_argument(
        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
    )
264

265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
        args.local_rank = env_local_rank

    if args.dataset_name is None and args.train_data_dir is None:
        raise ValueError("You must specify either a dataset name from the hub or a train data directory.")

    return args


def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
    if token is None:
        token = HfFolder.get_token()
    if organization is None:
        username = whoami(token)["name"]
        return f"{username}/{model_id}"
    else:
        return f"{organization}/{model_id}"


anton-l's avatar
anton-l committed
286
def main(args):
287
    logging_dir = os.path.join(args.output_dir, args.logging_dir)
288
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
289

Patrick von Platen's avatar
Patrick von Platen committed
290
    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=7200))  # a big number for high resolution or big dataset
291
    accelerator = Accelerator(
292
        gradient_accumulation_steps=args.gradient_accumulation_steps,
293
        mixed_precision=args.mixed_precision,
294
        log_with=args.logger,
295
        project_config=accelerator_project_config,
296
        kwargs_handlers=[kwargs],
297
    )
anton-l's avatar
anton-l committed
298

299
300
301
302
303
304
305
306
307
    if args.logger == "tensorboard":
        if not is_tensorboard_available():
            raise ImportError("Make sure to install tensorboard if you want to use it for logging during training.")

    elif args.logger == "wandb":
        if not is_wandb_available():
            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
        import wandb

308
309
310
311
    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
        def save_model_hook(models, weights, output_dir):
312
313
314
            if accelerator.is_main_process:
                if args.use_ema:
                    ema_model.save_pretrained(os.path.join(output_dir, "unet_ema"))
315

316
317
                for i, model in enumerate(models):
                    model.save_pretrained(os.path.join(output_dir, "unet"))
318

319
320
                    # make sure to pop weight so that corresponding model is not saved again
                    weights.pop()
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342

        def load_model_hook(models, input_dir):
            if args.use_ema:
                load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), UNet2DModel)
                ema_model.load_state_dict(load_model.state_dict())
                ema_model.to(accelerator.device)
                del load_model

            for i in range(len(models)):
                # pop models so that they are not loaded again
                model = models.pop()

                # load diffusers style into model
                load_model = UNet2DModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

                model.load_state_dict(load_model.state_dict())
                del load_model

        accelerator.register_save_state_pre_hook(save_model_hook)
        accelerator.register_load_state_pre_hook(load_model_hook)

343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.push_to_hub:
            if args.hub_model_id is None:
                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
            else:
                repo_name = args.hub_model_id
364
365
            create_repo(repo_name, exist_ok=True, token=args.hub_token)
            repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
366
367
368
369
370
371
372
373
374
375

            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
                if "step_*" not in gitignore:
                    gitignore.write("step_*\n")
                if "epoch_*" not in gitignore:
                    gitignore.write("epoch_*\n")
        elif args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)

    # Initialize the model
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
    if args.model_config_name_or_path is None:
        model = UNet2DModel(
            sample_size=args.resolution,
            in_channels=3,
            out_channels=3,
            layers_per_block=2,
            block_out_channels=(128, 128, 256, 256, 512, 512),
            down_block_types=(
                "DownBlock2D",
                "DownBlock2D",
                "DownBlock2D",
                "DownBlock2D",
                "AttnDownBlock2D",
                "DownBlock2D",
            ),
            up_block_types=(
                "UpBlock2D",
                "AttnUpBlock2D",
                "UpBlock2D",
                "UpBlock2D",
                "UpBlock2D",
                "UpBlock2D",
            ),
        )
    else:
        config = UNet2DModel.load_config(args.model_config_name_or_path)
        model = UNet2DModel.from_config(config)
403

404
405
406
407
408
409
410
411
    # Create EMA for the model.
    if args.use_ema:
        ema_model = EMAModel(
            model.parameters(),
            decay=args.ema_max_decay,
            use_ema_warmup=True,
            inv_gamma=args.ema_inv_gamma,
            power=args.ema_power,
412
413
            model_cls=UNet2DModel,
            model_config=model.config,
414
415
        )

416
417
418
419
420
421
422
423
    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
        args.mixed_precision = accelerator.mixed_precision
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
        args.mixed_precision = accelerator.mixed_precision

424
425
426
427
428
429
430
431
432
433
434
435
436
    if args.enable_xformers_memory_efficient_attention:
        if is_xformers_available():
            import xformers

            xformers_version = version.parse(xformers.__version__)
            if xformers_version == version.parse("0.0.16"):
                logger.warn(
                    "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
                )
            model.enable_xformers_memory_efficient_attention()
        else:
            raise ValueError("xformers is not available. Make sure it is installed correctly")

437
438
    # Initialize the scheduler
    accepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())
439
    if accepts_prediction_type:
440
441
442
        noise_scheduler = DDPMScheduler(
            num_train_timesteps=args.ddpm_num_steps,
            beta_schedule=args.ddpm_beta_schedule,
443
            prediction_type=args.prediction_type,
444
445
446
447
        )
    else:
        noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

448
    # Initialize the optimizer
449
450
451
452
453
454
455
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )
anton-l's avatar
anton-l committed
456

457
458
    # Get the datasets: you can either provide your own training and evaluation files (see below)
    # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub).
459

460
461
    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
    # download the dataset.
462
463
464
465
466
467
468
469
470
    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            cache_dir=args.cache_dir,
            split="train",
        )
    else:
        dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")
471
472
473
474
        # See more about loading custom images at
        # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder

    # Preprocessing the datasets and DataLoaders creation.
475
    augmentations = transforms.Compose(
476
        [
477
478
479
480
481
            transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
            transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
482
483
        ]
    )
anton-l's avatar
anton-l committed
484

485
    def transform_images(examples):
anton-l's avatar
anton-l committed
486
487
488
        images = [augmentations(image.convert("RGB")) for image in examples["image"]]
        return {"input": images}

489
490
    logger.info(f"Dataset size: {len(dataset)}")

491
    dataset.set_transform(transform_images)
492
493
494
    train_dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
    )
anton-l's avatar
anton-l committed
495

496
    # Initialize the learning rate scheduler
anton-l's avatar
anton-l committed
497
    lr_scheduler = get_scheduler(
498
        args.lr_scheduler,
anton-l's avatar
anton-l committed
499
        optimizer=optimizer,
500
501
        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=(len(train_dataloader) * args.num_epochs),
anton-l's avatar
anton-l committed
502
503
    )

504
    # Prepare everything with our `accelerator`.
anton-l's avatar
anton-l committed
505
506
507
    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, lr_scheduler
    )
508

509
510
    if args.use_ema:
        ema_model.to(accelerator.device)
anton-l's avatar
anton-l committed
511

512
513
    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
514
515
516
517
    if accelerator.is_main_process:
        run = os.path.split(__file__)[-1].split(".")[0]
        accelerator.init_trackers(run)

518
519
520
521
522
523
524
525
526
527
528
529
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    max_train_steps = args.num_epochs * num_update_steps_per_epoch

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(dataset)}")
    logger.info(f"  Num Epochs = {args.num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

anton-l's avatar
anton-l committed
530
    global_step = 0
531
532
    first_epoch = 0

533
    # Potentially load in the weights and states from a previous save
534
535
536
537
538
539
540
541
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint != "latest":
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
            path = dirs[-1] if len(dirs) > 0 else None

        if path is None:
            accelerator.print(
                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            args.resume_from_checkpoint = None
        else:
            accelerator.print(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(args.output_dir, path))
            global_step = int(path.split("-")[1])

            resume_global_step = global_step * args.gradient_accumulation_steps
            first_epoch = global_step // num_update_steps_per_epoch
            resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
557

558
    # Train!
559
    for epoch in range(first_epoch, args.num_epochs):
anton-l's avatar
anton-l committed
560
        model.train()
561
        progress_bar = tqdm(total=num_update_steps_per_epoch, disable=not accelerator.is_local_main_process)
562
563
        progress_bar.set_description(f"Epoch {epoch}")
        for step, batch in enumerate(train_dataloader):
564
565
566
567
568
569
            # Skip steps until we reach the resumed step
            if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
                if step % args.gradient_accumulation_steps == 0:
                    progress_bar.update(1)
                continue

570
            clean_images = batch["input"].to(weight_dtype)
571
            # Sample noise that we'll add to the images
572
            noise = torch.randn(clean_images.shape, dtype=weight_dtype, device=clean_images.device)
573
            bsz = clean_images.shape[0]
574
575
            # Sample a random timestep for each image
            timesteps = torch.randint(
576
                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=clean_images.device
577
            ).long()
578

579
            # Add noise to the clean images according to the noise magnitude at each timestep
580
            # (this is the forward diffusion process)
581
582
583
584
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

            with accelerator.accumulate(model):
                # Predict the noise residual
585
586
                model_output = model(noisy_images, timesteps).sample

587
                if args.prediction_type == "epsilon":
588
                    loss = F.mse_loss(model_output.float(), noise.float())  # this could have different weights!
589
                elif args.prediction_type == "sample":
590
591
592
593
                    alpha_t = _extract_into_tensor(
                        noise_scheduler.alphas_cumprod, timesteps, (clean_images.shape[0], 1, 1, 1)
                    )
                    snr_weights = alpha_t / (1 - alpha_t)
594
595
                    # use SNR weighting from distillation paper
                    loss = snr_weights * F.mse_loss(model_output.float(), clean_images.float(), reduction="none")
596
                    loss = loss.mean()
597
598
                else:
                    raise ValueError(f"Unsupported prediction type: {args.prediction_type}")
599

600
                accelerator.backward(loss)
601

602
603
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
604
605
606
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
607

608
609
            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
610
611
                if args.use_ema:
                    ema_model.step(model.parameters())
612
613
614
                progress_bar.update(1)
                global_step += 1

615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
                if accelerator.is_main_process:
                    if global_step % args.checkpointing_steps == 0:
                        # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                        if args.checkpoints_total_limit is not None:
                            checkpoints = os.listdir(args.output_dir)
                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))

                            # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
                            if len(checkpoints) >= args.checkpoints_total_limit:
                                num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
                                removing_checkpoints = checkpoints[0:num_to_remove]

                                logger.info(
                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
                                )
                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")

                                for removing_checkpoint in removing_checkpoints:
                                    removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
                                    shutil.rmtree(removing_checkpoint)

637
638
639
640
                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)
                        logger.info(f"Saved state to {save_path}")

641
642
            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
            if args.use_ema:
643
                logs["ema_decay"] = ema_model.cur_decay_value
644
645
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)
646
        progress_bar.close()
anton-l's avatar
anton-l committed
647

anton-l's avatar
anton-l committed
648
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
649

anton-l's avatar
anton-l committed
650
        # Generate sample images for visual inspection
anton-l's avatar
anton-l committed
651
        if accelerator.is_main_process:
anton-l's avatar
anton-l committed
652
            if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
653
                unet = accelerator.unwrap_model(model)
654

655
                if args.use_ema:
656
                    ema_model.store(unet.parameters())
657
                    ema_model.copy_to(unet.parameters())
658

659
                pipeline = DDPMPipeline(
660
                    unet=unet,
661
                    scheduler=noise_scheduler,
anton-l's avatar
anton-l committed
662
                )
anton-l's avatar
anton-l committed
663

664
                generator = torch.Generator(device=pipeline.device).manual_seed(0)
anton-l's avatar
anton-l committed
665
                # run pipeline in inference (sample random noise and denoise)
666
667
668
                images = pipeline(
                    generator=generator,
                    batch_size=args.eval_batch_size,
669
                    num_inference_steps=args.ddpm_num_inference_steps,
670
671
                    output_type="numpy",
                ).images
anton-l's avatar
anton-l committed
672

673
674
675
                if args.use_ema:
                    ema_model.restore(unet.parameters())

anton-l's avatar
anton-l committed
676
677
                # denormalize the images and save to tensorboard
                images_processed = (images * 255).round().astype("uint8")
678

679
                if args.logger == "tensorboard":
680
681
682
                    if is_accelerate_version(">=", "0.17.0.dev0"):
                        tracker = accelerator.get_tracker("tensorboard", unwrap=True)
                    else:
683
                        tracker = accelerator.get_tracker("tensorboard")
684
                    tracker.add_images("test_samples", images_processed.transpose(0, 3, 1, 2), epoch)
685
                elif args.logger == "wandb":
686
                    # Upcoming `log_images` helper coming in https://github.com/huggingface/accelerate/pull/962/files
687
688
689
690
                    accelerator.get_tracker("wandb").log(
                        {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch},
                        step=global_step,
                    )
anton-l's avatar
anton-l committed
691

692
693
            if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
                # save the model
694
695
696
697
698
699
700
701
702
703
704
                unet = accelerator.unwrap_model(model)

                if args.use_ema:
                    ema_model.store(unet.parameters())
                    ema_model.copy_to(unet.parameters())

                pipeline = DDPMPipeline(
                    unet=unet,
                    scheduler=noise_scheduler,
                )

705
                pipeline.save_pretrained(args.output_dir)
706
707
708
709

                if args.use_ema:
                    ema_model.restore(unet.parameters())

710
                if args.push_to_hub:
711
                    repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False)
anton-l's avatar
anton-l committed
712

713
714
    accelerator.end_training()

anton-l's avatar
anton-l committed
715
716

if __name__ == "__main__":
717
    args = parse_args()
anton-l's avatar
anton-l committed
718
    main(args)