@@ -4,14 +4,13 @@
 import torch
 import torch.nn.functional as F
 
-import PIL.Image
 from accelerate import Accelerator
+from accelerate.logging import get_logger
 from datasets import load_dataset
-from diffusers import DDPMPipeline, DDPMScheduler, UNetModel
+from diffusers import DDIMPipeline, DDIMScheduler, UNetModel
 from diffusers.hub_utils import init_git_repo, push_to_hub
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import EMAModel
-from diffusers.utils import logging
 from torchvision.transforms import (
     CenterCrop,
     Compose,
@@ -24,11 +23,12 @@
 from tqdm.auto import tqdm
 
 
-logger = logging.get_logger(__name__)
+logger = get_logger(__name__)
 
 
 def main(args):
-    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+    logging_dir = os.path.join(args.output_dir, args.logging_dir)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, log_with="tensorboard", logging_dir=logging_dir)
 
     model = UNetModel(
         attn_resolutions=(16,),
@@ -39,8 +39,14 @@ def main(args):
         resamp_with_conv=True,
         resolution=args.resolution,
     )
-    noise_scheduler = DDPMScheduler(timesteps=1000, tensor_format="pt")
-    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    noise_scheduler = DDIMScheduler(timesteps=1000, tensor_format="pt")
+    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,
+    )
 
     augmentations = Compose(
         [
@@ -58,12 +64,12 @@ def transforms(examples):
         return {"input": images}
 
     dataset.set_transform(transforms)
-    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
+    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.train_batch_size, shuffle=True)
 
     lr_scheduler = get_scheduler(
-        "linear",
+        args.lr_scheduler,
         optimizer=optimizer,
-        num_warmup_steps=args.warmup_steps,
+        num_warmup_steps=args.lr_warmup_steps,
         num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
     )
 
@@ -76,81 +82,91 @@ def transforms(examples):
     if args.push_to_hub:
         repo = init_git_repo(args, at_init=True)
 
+    if accelerator.is_main_process:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        accelerator.init_trackers(run)
+
     # Train!
     is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
     world_size = torch.distributed.get_world_size() if is_distributed else 1
-    total_train_batch_size = args.batch_size * args.gradient_accumulation_steps * world_size
+    total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * world_size
     max_steps = len(train_dataloader) // args.gradient_accumulation_steps * args.num_epochs
     logger.info("***** Running training *****")
     logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
     logger.info(f"  Num Epochs = {args.num_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {args.batch_size}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
     logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
     logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     logger.info(f"  Total optimization steps = {max_steps}")
 
     global_step = 0
     for epoch in range(args.num_epochs):
         model.train()
-        with tqdm(total=len(train_dataloader), unit="ba") as pbar:
-            pbar.set_description(f"Epoch {epoch}")
-            for step, batch in enumerate(train_dataloader):
-                clean_images = batch["input"]
-                noise_samples = torch.randn(clean_images.shape).to(clean_images.device)
-                bsz = clean_images.shape[0]
-                timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_images.device).long()
-
-                # add noise onto the clean images according to the noise magnitude at each timestep
-                # (this is the forward diffusion process)
-                noisy_images = noise_scheduler.training_step(clean_images, noise_samples, timesteps)
-
-                if step % args.gradient_accumulation_steps != 0:
-                    with accelerator.no_sync(model):
-                        output = model(noisy_images, timesteps)
-                        # predict the noise residual
-                        loss = F.mse_loss(output, noise_samples)
-                        loss = loss / args.gradient_accumulation_steps
-                        accelerator.backward(loss)
-                else:
+        progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
+        progress_bar.set_description(f"Epoch {epoch}")
+        for step, batch in enumerate(train_dataloader):
+            clean_images = batch["input"]
+            noise_samples = torch.randn(clean_images.shape).to(clean_images.device)
+            bsz = clean_images.shape[0]
+            timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_images.device).long()
+
+            # add noise onto the clean images according to the noise magnitude at each timestep
+            # (this is the forward diffusion process)
+            noisy_images = noise_scheduler.add_noise(clean_images, noise_samples, timesteps)
+
+            if step % args.gradient_accumulation_steps != 0:
+                with accelerator.no_sync(model):
                     output = model(noisy_images, timesteps)
                     # predict the noise residual
                     loss = F.mse_loss(output, noise_samples)
                     loss = loss / args.gradient_accumulation_steps
                     accelerator.backward(loss)
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-                    optimizer.step()
-                    lr_scheduler.step()
-                    ema_model.step(model, global_step)
-                    optimizer.zero_grad()
-                pbar.update(1)
-                pbar.set_postfix(
-                    loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"], ema_decay=ema_model.decay
-                )
-                global_step += 1
+            else:
+                output = model(noisy_images, timesteps)
+                # predict the noise residual
+                loss = F.mse_loss(output, noise_samples)
+                loss = loss / args.gradient_accumulation_steps
+                accelerator.backward(loss)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+                lr_scheduler.step()
+                ema_model.step(model, global_step)
+                optimizer.zero_grad()
+            progress_bar.update(1)
+            progress_bar.set_postfix(
+                loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"], ema_decay=ema_model.decay
+            )
+            accelerator.log(
+                {
+                    "train_loss": loss.detach().item(),
+                    "epoch": epoch,
+                    "ema_decay": ema_model.decay,
+                    "step": global_step,
+                },
+                step=global_step,
+            )
+            global_step += 1
+        progress_bar.close()
 
         accelerator.wait_for_everyone()
 
         # Generate a sample image for visual inspection
         if accelerator.is_main_process:
             with torch.no_grad():
-                pipeline = DDPMPipeline(
-                    unet=accelerator.unwrap_model(ema_model.averaged_model), noise_scheduler=noise_scheduler
+                pipeline = DDIMPipeline(
+                    unet=accelerator.unwrap_model(ema_model.averaged_model),
+                    noise_scheduler=noise_scheduler,
                 )
 
                 generator = torch.manual_seed(0)
                 # run pipeline in inference (sample random noise and denoise)
-                image = pipeline(generator=generator)
+                images = pipeline(generator=generator, batch_size=args.eval_batch_size, num_inference_steps=50)
 
-            # process image to PIL
-            image_processed = image.cpu().permute(0, 2, 3, 1)
-            image_processed = (image_processed + 1.0) * 127.5
-            image_processed = image_processed.type(torch.uint8).numpy()
-            image_pil = PIL.Image.fromarray(image_processed[0])
+            # denormalize the images and save to tensorboard
+            images_processed = (images.cpu() + 1.0) * 127.5
+            images_processed = images_processed.clamp(0, 255).type(torch.uint8).numpy()
 
-            # save image
-            test_dir = os.path.join(args.output_dir, "test_samples")
-            os.makedirs(test_dir, exist_ok=True)
-            image_pil.save(f"{test_dir}/{epoch:04d}.png")
+            accelerator.trackers[0].writer.add_images("test_samples", images_processed, epoch)
 
             # save the model
             if args.push_to_hub:
@@ -159,6 +175,8 @@ def transforms(examples):
                 pipeline.save_pretrained(args.output_dir)
         accelerator.wait_for_everyone()
 
+    accelerator.end_training()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -167,18 +185,25 @@ def transforms(examples):
     parser.add_argument("--output_dir", type=str, default="ddpm-model")
     parser.add_argument("--overwrite_output_dir", action="store_true")
     parser.add_argument("--resolution", type=int, default=64)
-    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--train_batch_size", type=int, default=16)
+    parser.add_argument("--eval_batch_size", type=int, default=16)
     parser.add_argument("--num_epochs", type=int, default=100)
     parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
-    parser.add_argument("--lr", type=float, default=1e-4)
-    parser.add_argument("--warmup_steps", type=int, default=500)
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--lr_scheduler", type=str, default="cosine")
+    parser.add_argument("--lr_warmup_steps", type=int, default=500)
+    parser.add_argument("--adam_beta1", type=float, default=0.95)
+    parser.add_argument("--adam_beta2", type=float, default=0.999)
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-6)
+    parser.add_argument("--adam_epsilon", type=float, default=1e-3)
     parser.add_argument("--ema_inv_gamma", type=float, default=1.0)
     parser.add_argument("--ema_power", type=float, default=3 / 4)
-    parser.add_argument("--ema_max_decay", type=float, default=0.999)
+    parser.add_argument("--ema_max_decay", type=float, default=0.9999)
     parser.add_argument("--push_to_hub", action="store_true")
     parser.add_argument("--hub_token", type=str, default=None)
     parser.add_argument("--hub_model_id", type=str, default=None)
     parser.add_argument("--hub_private_repo", action="store_true")
+    parser.add_argument("--logging_dir", type=str, default="logs")
     parser.add_argument(
         "--mixed_precision",
         type=str,
 
@@ -87,6 +87,8 @@
     "regex!=2019.12.17",
     "requests",
     "torch>=1.4",
+    "tensorboard",
+    "modelcards=0.1.4"
 ]
 
 # this is a lookup table with items like:
@@ -172,6 +174,8 @@ def run(self):
     deps["requests"],
     deps["torch"],
     deps["Pillow"],
+    deps["tensorboard"],
+    deps["modelcards"],
 ]
 
 setup(
 
@@ -13,4 +13,5 @@
     "regex": "regex!=2019.12.17",
     "requests": "requests",
     "torch": "torch>=1.4",
+    "tensorboard": "tensorboard",
 }
@@ -19,20 +19,17 @@
 from pathlib import Path
 from typing import Optional
 
-import yaml
 from diffusers import DiffusionPipeline
 from huggingface_hub import HfFolder, Repository, whoami
+from modelcards import CardData, ModelCard
 
 from .utils import logging
 
 
 logger = logging.get_logger(__name__)
 
 
-AUTOGENERATED_TRAINER_COMMENT = """
-<!-- This model card has been generated automatically according to the information the Trainer had access to. You
-should probably proofread and complete it, then remove this comment. -->
-"""
+MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "utils" / "model_card_template.md"
 
 
 def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
@@ -152,17 +149,36 @@ def create_model_card(args, model_name):
     if args.local_rank not in [-1, 0]:
         return
 
-    # TODO: replace this placeholder model card generation
-    model_card = ""
-
-    metadata = {"license": "apache-2.0", "tags": ["pytorch", "diffusers"]}
-    metadata = yaml.dump(metadata, sort_keys=False)
-    if len(metadata) > 0:
-        model_card = f"---\n{metadata}---\n"
-
-    model_card += AUTOGENERATED_TRAINER_COMMENT
-
-    model_card += f"\n# {model_name}\n\n"
-
-    with open(os.path.join(args.output_dir, "README.md"), "w") as f:
-        f.write(model_card)
+    repo_name = get_full_repo_name(model_name, token=args.hub_token)
+
+    model_card = ModelCard.from_template(
+        card_data=CardData(  # Card metadata object that will be converted to YAML block
+            language="en",
+            license="apache-2.0",
+            library_name="diffusers",
+            tags=[],
+            datasets=args.dataset,
+            metrics=[],
+        ),
+        template_path=MODEL_CARD_TEMPLATE_PATH,
+        model_name=model_name,
+        repo_name=repo_name,
+        dataset_name=args.dataset,
+        learning_rate=args.learning_rate,
+        train_batch_size=args.train_batch_size,
+        eval_batch_size=args.eval_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        adam_beta1=args.adam_beta1,
+        adam_beta2=args.adam_beta2,
+        adam_weight_decay=args.adam_weight_decay,
+        adam_epsilon=args.adam_epsilon,
+        lr_scheduler=args.lr_scheduler,
+        lr_warmup_steps=args.lr_warmup_steps,
+        ema_inv_gamma=args.ema_inv_gamma,
+        ema_power=args.ema_power,
+        ema_max_decay=args.ema_max_decay,
+        mixed_precision=args.mixed_precision,
+    )
+
+    card_path = os.path.join(args.output_dir, "README.md")
+    model_card.save(card_path)
@@ -1 +1 @@
-from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline
+from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline, VQModel
@@ -145,5 +145,14 @@ def step(self, residual, sample, t, num_inference_steps, eta, use_clipped_residu
 
         return pred_prev_sample
 
+    def add_noise(self, original_samples, noise, timesteps):
+        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
+        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = self.match_shape(sqrt_one_minus_alpha_prod, original_samples)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
     def __len__(self):
         return self.config.timesteps
@@ -17,7 +17,6 @@
 import math
 
 import numpy as np
-import torch
 
 from ..configuration_utils import ConfigMixin
 from .scheduling_utils import SchedulerMixin
@@ -142,7 +141,7 @@ def step(self, residual, sample, t, predict_epsilon=True):
 
         return pred_prev_sample
 
-    def training_step(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor):
+    def add_noise(self, original_samples, noise, timesteps):
         sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
         sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
         sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
Original file line number	Diff line number	Diff line change
`@@ -13,4 +13,5 @@`
`13`	`13`	`"regex": "regex!=2019.12.17",`
`14`	`14`	`"requests": "requests",`
`15`	`15`	`"torch": "torch>=1.4",`
	`16`	`+ "tensorboard": "tensorboard",`
`16`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline`
	`1`	`+from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline, VQModel`