We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Example unconditional training script does not work
### Describe the bug

When running the [example unconditional training script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py)
on GPU using `--mixed_precision=fp16`, the autocasting of accelerate will break:

```
Traceback (most recent call last):
  File "<path>\train_unconditional.py", line 251, in <module>
    main(args)
  File "<path>\train_unconditional.py", line 145, in main
    loss = F.mse_loss(noise_pred, noise)
  File "<lib-path>\lib\site-packages\torch\nn\functional.py", line 3269, in mse_loss
    if not (target.size() == input.size()):
AttributeError: 'dict' object has no attribute 'size'
```

This happens due to the fact that in `accelerate.utils.operations.recursively_apply` the "Mapping" class is handled using a dict initialization (=> https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L87). This breaks the UNet2DOutput class, because it will wrap the output in a dict (because UNet2DOutput does inherit from OrderedDict).
i.e. `UNet2DOutput(sample=sample)` will become equivalent to `UNet2DOutput(sample={'sample': sample})`.

now this could be changed in the train_unconditional.py. unfortunately this just moves the problem, as it will fail in a later stage in library code that isn't accessible to the user.

### Reproduction

Run the example script with mixed-precision set to fp16:
```
train_unconditional.py --dataset_name="huggan/pokemon" --output_dir="ddpm-test-64" --mixed_precision=fp16
```

### Logs

_No response_

### System Info

- `diffusers` version: 0.3.0
- Platform: Windows-10-10.0.22000-SP0
- Python version: 3.9.13
- PyTorch version (GPU?): 1.12.1+cu116 (True)
- Huggingface_hub version: 0.9.1
- Transformers version: 4.21.3
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

some comments:
*  i'm running on windows 11, wrongly reported as 10 using `diffusers-cli env`
* `accelerate` version: 0.12.0

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/unconditional_image_generation/train_unconditional.py`
Content:
```
1 import argparse
2 import math
3 import os
4 
5 import torch
6 import torch.nn.functional as F
7 
8 from accelerate import Accelerator
9 from accelerate.logging import get_logger
10 from datasets import load_dataset
11 from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
12 from diffusers.hub_utils import init_git_repo, push_to_hub
13 from diffusers.optimization import get_scheduler
14 from diffusers.training_utils import EMAModel
15 from torchvision.transforms import (
16     CenterCrop,
17     Compose,
18     InterpolationMode,
19     Normalize,
20     RandomHorizontalFlip,
21     Resize,
22     ToTensor,
23 )
24 from tqdm.auto import tqdm
25 
26 
27 logger = get_logger(__name__)
28 
29 
30 def main(args):
31     logging_dir = os.path.join(args.output_dir, args.logging_dir)
32     accelerator = Accelerator(
33         gradient_accumulation_steps=args.gradient_accumulation_steps,
34         mixed_precision=args.mixed_precision,
35         log_with="tensorboard",
36         logging_dir=logging_dir,
37     )
38 
39     model = UNet2DModel(
40         sample_size=args.resolution,
41         in_channels=3,
42         out_channels=3,
43         layers_per_block=2,
44         block_out_channels=(128, 128, 256, 256, 512, 512),
45         down_block_types=(
46             "DownBlock2D",
47             "DownBlock2D",
48             "DownBlock2D",
49             "DownBlock2D",
50             "AttnDownBlock2D",
51             "DownBlock2D",
52         ),
53         up_block_types=(
54             "UpBlock2D",
55             "AttnUpBlock2D",
56             "UpBlock2D",
57             "UpBlock2D",
58             "UpBlock2D",
59             "UpBlock2D",
60         ),
61     )
62     noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
63     optimizer = torch.optim.AdamW(
64         model.parameters(),
65         lr=args.learning_rate,
66         betas=(args.adam_beta1, args.adam_beta2),
67         weight_decay=args.adam_weight_decay,
68         eps=args.adam_epsilon,
69     )
70 
71     augmentations = Compose(
72         [
73             Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
74             CenterCrop(args.resolution),
75             RandomHorizontalFlip(),
76             ToTensor(),
77             Normalize([0.5], [0.5]),
78         ]
79     )
80 
81     if args.dataset_name is not None:
82         dataset = load_dataset(
83             args.dataset_name,
84             args.dataset_config_name,
85             cache_dir=args.cache_dir,
86             use_auth_token=True if args.use_auth_token else None,
87             split="train",
88         )
89     else:
90         dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")
91 
92     def transforms(examples):
93         images = [augmentations(image.convert("RGB")) for image in examples["image"]]
94         return {"input": images}
95 
96     dataset.set_transform(transforms)
97     train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.train_batch_size, shuffle=True)
98 
99     lr_scheduler = get_scheduler(
100         args.lr_scheduler,
101         optimizer=optimizer,
102         num_warmup_steps=args.lr_warmup_steps,
103         num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
104     )
105 
106     model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
107         model, optimizer, train_dataloader, lr_scheduler
108     )
109 
110     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
111 
112     ema_model = EMAModel(model, inv_gamma=args.ema_inv_gamma, power=args.ema_power, max_value=args.ema_max_decay)
113 
114     if args.push_to_hub:
115         repo = init_git_repo(args, at_init=True)
116 
117     if accelerator.is_main_process:
118         run = os.path.split(__file__)[-1].split(".")[0]
119         accelerator.init_trackers(run)
120 
121     global_step = 0
122     for epoch in range(args.num_epochs):
123         model.train()
124         progress_bar = tqdm(total=num_update_steps_per_epoch, disable=not accelerator.is_local_main_process)
125         progress_bar.set_description(f"Epoch {epoch}")
126         for step, batch in enumerate(train_dataloader):
127             clean_images = batch["input"]
128             # Sample noise that we'll add to the images
129             noise = torch.randn(clean_images.shape).to(clean_images.device)
130             bsz = clean_images.shape[0]
131             # Sample a random timestep for each image
132             timesteps = torch.randint(
133                 0, noise_scheduler.config.num_train_timesteps, (bsz,), device=clean_images.device
134             ).long()
135 
136             # Add noise to the clean images according to the noise magnitude at each timestep
137             # (this is the forward diffusion process)
138             noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
139 
140             with accelerator.accumulate(model):
141                 # Predict the noise residual
142                 noise_pred = model(noisy_images, timesteps).sample
143                 loss = F.mse_loss(noise_pred, noise)
144                 accelerator.backward(loss)
145 
146                 accelerator.clip_grad_norm_(model.parameters(), 1.0)
147                 optimizer.step()
148                 lr_scheduler.step()
149                 if args.use_ema:
150                     ema_model.step(model)
151                 optimizer.zero_grad()
152 
153             # Checks if the accelerator has performed an optimization step behind the scenes
154             if accelerator.sync_gradients:
155                 progress_bar.update(1)
156                 global_step += 1
157 
158             logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
159             if args.use_ema:
160                 logs["ema_decay"] = ema_model.decay
161             progress_bar.set_postfix(**logs)
162             accelerator.log(logs, step=global_step)
163         progress_bar.close()
164 
165         accelerator.wait_for_everyone()
166 
167         # Generate sample images for visual inspection
168         if accelerator.is_main_process:
169             if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
170                 pipeline = DDPMPipeline(
171                     unet=accelerator.unwrap_model(ema_model.averaged_model if args.use_ema else model),
172                     scheduler=noise_scheduler,
173                 )
174 
175                 generator = torch.manual_seed(0)
176                 # run pipeline in inference (sample random noise and denoise)
177                 images = pipeline(generator=generator, batch_size=args.eval_batch_size, output_type="numpy").images
178 
179                 # denormalize the images and save to tensorboard
180                 images_processed = (images * 255).round().astype("uint8")
181                 accelerator.trackers[0].writer.add_images(
182                     "test_samples", images_processed.transpose(0, 3, 1, 2), epoch
183                 )
184 
185             if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
186                 # save the model
187                 if args.push_to_hub:
188                     push_to_hub(args, pipeline, repo, commit_message=f"Epoch {epoch}", blocking=False)
189                 else:
190                     pipeline.save_pretrained(args.output_dir)
191         accelerator.wait_for_everyone()
192 
193     accelerator.end_training()
194 
195 
196 if __name__ == "__main__":
197     parser = argparse.ArgumentParser(description="Simple example of a training script.")
198     parser.add_argument("--local_rank", type=int, default=-1)
199     parser.add_argument("--dataset_name", type=str, default=None)
200     parser.add_argument("--dataset_config_name", type=str, default=None)
201     parser.add_argument("--train_data_dir", type=str, default=None, help="A folder containing the training data.")
202     parser.add_argument("--output_dir", type=str, default="ddpm-model-64")
203     parser.add_argument("--overwrite_output_dir", action="store_true")
204     parser.add_argument("--cache_dir", type=str, default=None)
205     parser.add_argument("--resolution", type=int, default=64)
206     parser.add_argument("--train_batch_size", type=int, default=16)
207     parser.add_argument("--eval_batch_size", type=int, default=16)
208     parser.add_argument("--num_epochs", type=int, default=100)
209     parser.add_argument("--save_images_epochs", type=int, default=10)
210     parser.add_argument("--save_model_epochs", type=int, default=10)
211     parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
212     parser.add_argument("--learning_rate", type=float, default=1e-4)
213     parser.add_argument("--lr_scheduler", type=str, default="cosine")
214     parser.add_argument("--lr_warmup_steps", type=int, default=500)
215     parser.add_argument("--adam_beta1", type=float, default=0.95)
216     parser.add_argument("--adam_beta2", type=float, default=0.999)
217     parser.add_argument("--adam_weight_decay", type=float, default=1e-6)
218     parser.add_argument("--adam_epsilon", type=float, default=1e-08)
219     parser.add_argument("--use_ema", action="store_true", default=True)
220     parser.add_argument("--ema_inv_gamma", type=float, default=1.0)
221     parser.add_argument("--ema_power", type=float, default=3 / 4)
222     parser.add_argument("--ema_max_decay", type=float, default=0.9999)
223     parser.add_argument("--push_to_hub", action="store_true")
224     parser.add_argument("--use_auth_token", action="store_true")
225     parser.add_argument("--hub_token", type=str, default=None)
226     parser.add_argument("--hub_model_id", type=str, default=None)
227     parser.add_argument("--hub_private_repo", action="store_true")
228     parser.add_argument("--logging_dir", type=str, default="logs")
229     parser.add_argument(
230         "--mixed_precision",
231         type=str,
232         default="no",
233         choices=["no", "fp16", "bf16"],
234         help=(
235             "Whether to use mixed precision. Choose"
236             "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
237             "and an Nvidia Ampere GPU."
238         ),
239     )
240 
241     args = parser.parse_args()
242     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
243     if env_local_rank != -1 and env_local_rank != args.local_rank:
244         args.local_rank = env_local_rank
245 
246     if args.dataset_name is None and args.train_data_dir is None:
247         raise ValueError("You must specify either a dataset name from the hub or a train data directory.")
248 
249     main(args)
250
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py
--- a/examples/unconditional_image_generation/train_unconditional.py
+++ b/examples/unconditional_image_generation/train_unconditional.py
@@ -143,7 +143,8 @@
                 loss = F.mse_loss(noise_pred, noise)
                 accelerator.backward(loss)
 
-                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                 optimizer.step()
                 lr_scheduler.step()
                 if args.use_ema:

{"golden_diff": "diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py\n--- a/examples/unconditional_image_generation/train_unconditional.py\n+++ b/examples/unconditional_image_generation/train_unconditional.py\n@@ -143,7 +143,8 @@\n                 loss = F.mse_loss(noise_pred, noise)\n                 accelerator.backward(loss)\n \n-                accelerator.clip_grad_norm_(model.parameters(), 1.0)\n+                if accelerator.sync_gradients:\n+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)\n                 optimizer.step()\n                 lr_scheduler.step()\n                 if args.use_ema:\n", "issue": "Example unconditional training script does not work\n### Describe the bug\n\nWhen running the [example unconditional training script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py)\r\non GPU using `--mixed_precision=fp16`, the autocasting of accelerate will break:\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"<path>\\train_unconditional.py\", line 251, in <module>\r\n    main(args)\r\n  File \"<path>\\train_unconditional.py\", line 145, in main\r\n    loss = F.mse_loss(noise_pred, noise)\r\n  File \"<lib-path>\\lib\\site-packages\\torch\\nn\\functional.py\", line 3269, in mse_loss\r\n    if not (target.size() == input.size()):\r\nAttributeError: 'dict' object has no attribute 'size'\r\n```\r\n\r\nThis happens due to the fact that in `accelerate.utils.operations.recursively_apply` the \"Mapping\" class is handled using a dict initialization (=> https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L87). This breaks the UNet2DOutput class, because it will wrap the output in a dict (because UNet2DOutput does inherit from OrderedDict).\r\ni.e. `UNet2DOutput(sample=sample)` will become equivalent to `UNet2DOutput(sample={'sample': sample})`.\r\n\r\nnow this could be changed in the train_unconditional.py. unfortunately this just moves the problem, as it will fail in a later stage in library code that isn't accessible to the user.\n\n### Reproduction\n\nRun the example script with mixed-precision set to fp16:\r\n```\r\ntrain_unconditional.py --dataset_name=\"huggan/pokemon\" --output_dir=\"ddpm-test-64\" --mixed_precision=fp16\r\n```\n\n### Logs\n\n_No response_\n\n### System Info\n\n- `diffusers` version: 0.3.0\r\n- Platform: Windows-10-10.0.22000-SP0\r\n- Python version: 3.9.13\r\n- PyTorch version (GPU?): 1.12.1+cu116 (True)\r\n- Huggingface_hub version: 0.9.1\r\n- Transformers version: 4.21.3\r\n- Using GPU in script?: yes\r\n- Using distributed or parallel set-up in script?: no\r\n\r\nsome comments:\r\n*  i'm running on windows 11, wrongly reported as 10 using `diffusers-cli env`\r\n* `accelerate` version: 0.12.0\n", "before_files": [{"content": "import argparse\nimport math\nimport os\n\nimport torch\nimport torch.nn.functional as F\n\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom datasets import load_dataset\nfrom diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel\nfrom diffusers.hub_utils import init_git_repo, push_to_hub\nfrom diffusers.optimization import get_scheduler\nfrom diffusers.training_utils import EMAModel\nfrom torchvision.transforms import (\n    CenterCrop,\n    Compose,\n    InterpolationMode,\n    Normalize,\n    RandomHorizontalFlip,\n    Resize,\n    ToTensor,\n)\nfrom tqdm.auto import tqdm\n\n\nlogger = get_logger(__name__)\n\n\ndef main(args):\n    logging_dir = os.path.join(args.output_dir, args.logging_dir)\n    accelerator = Accelerator(\n        gradient_accumulation_steps=args.gradient_accumulation_steps,\n        mixed_precision=args.mixed_precision,\n        log_with=\"tensorboard\",\n        logging_dir=logging_dir,\n    )\n\n    model = UNet2DModel(\n        sample_size=args.resolution,\n        in_channels=3,\n        out_channels=3,\n        layers_per_block=2,\n        block_out_channels=(128, 128, 256, 256, 512, 512),\n        down_block_types=(\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"AttnDownBlock2D\",\n            \"DownBlock2D\",\n        ),\n        up_block_types=(\n            \"UpBlock2D\",\n            \"AttnUpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n        ),\n    )\n    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=args.learning_rate,\n        betas=(args.adam_beta1, args.adam_beta2),\n        weight_decay=args.adam_weight_decay,\n        eps=args.adam_epsilon,\n    )\n\n    augmentations = Compose(\n        [\n            Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),\n            CenterCrop(args.resolution),\n            RandomHorizontalFlip(),\n            ToTensor(),\n            Normalize([0.5], [0.5]),\n        ]\n    )\n\n    if args.dataset_name is not None:\n        dataset = load_dataset(\n            args.dataset_name,\n            args.dataset_config_name,\n            cache_dir=args.cache_dir,\n            use_auth_token=True if args.use_auth_token else None,\n            split=\"train\",\n        )\n    else:\n        dataset = load_dataset(\"imagefolder\", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split=\"train\")\n\n    def transforms(examples):\n        images = [augmentations(image.convert(\"RGB\")) for image in examples[\"image\"]]\n        return {\"input\": images}\n\n    dataset.set_transform(transforms)\n    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.train_batch_size, shuffle=True)\n\n    lr_scheduler = get_scheduler(\n        args.lr_scheduler,\n        optimizer=optimizer,\n        num_warmup_steps=args.lr_warmup_steps,\n        num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,\n    )\n\n    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n        model, optimizer, train_dataloader, lr_scheduler\n    )\n\n    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n\n    ema_model = EMAModel(model, inv_gamma=args.ema_inv_gamma, power=args.ema_power, max_value=args.ema_max_decay)\n\n    if args.push_to_hub:\n        repo = init_git_repo(args, at_init=True)\n\n    if accelerator.is_main_process:\n        run = os.path.split(__file__)[-1].split(\".\")[0]\n        accelerator.init_trackers(run)\n\n    global_step = 0\n    for epoch in range(args.num_epochs):\n        model.train()\n        progress_bar = tqdm(total=num_update_steps_per_epoch, disable=not accelerator.is_local_main_process)\n        progress_bar.set_description(f\"Epoch {epoch}\")\n        for step, batch in enumerate(train_dataloader):\n            clean_images = batch[\"input\"]\n            # Sample noise that we'll add to the images\n            noise = torch.randn(clean_images.shape).to(clean_images.device)\n            bsz = clean_images.shape[0]\n            # Sample a random timestep for each image\n            timesteps = torch.randint(\n                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=clean_images.device\n            ).long()\n\n            # Add noise to the clean images according to the noise magnitude at each timestep\n            # (this is the forward diffusion process)\n            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)\n\n            with accelerator.accumulate(model):\n                # Predict the noise residual\n                noise_pred = model(noisy_images, timesteps).sample\n                loss = F.mse_loss(noise_pred, noise)\n                accelerator.backward(loss)\n\n                accelerator.clip_grad_norm_(model.parameters(), 1.0)\n                optimizer.step()\n                lr_scheduler.step()\n                if args.use_ema:\n                    ema_model.step(model)\n                optimizer.zero_grad()\n\n            # Checks if the accelerator has performed an optimization step behind the scenes\n            if accelerator.sync_gradients:\n                progress_bar.update(1)\n                global_step += 1\n\n            logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0], \"step\": global_step}\n            if args.use_ema:\n                logs[\"ema_decay\"] = ema_model.decay\n            progress_bar.set_postfix(**logs)\n            accelerator.log(logs, step=global_step)\n        progress_bar.close()\n\n        accelerator.wait_for_everyone()\n\n        # Generate sample images for visual inspection\n        if accelerator.is_main_process:\n            if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:\n                pipeline = DDPMPipeline(\n                    unet=accelerator.unwrap_model(ema_model.averaged_model if args.use_ema else model),\n                    scheduler=noise_scheduler,\n                )\n\n                generator = torch.manual_seed(0)\n                # run pipeline in inference (sample random noise and denoise)\n                images = pipeline(generator=generator, batch_size=args.eval_batch_size, output_type=\"numpy\").images\n\n                # denormalize the images and save to tensorboard\n                images_processed = (images * 255).round().astype(\"uint8\")\n                accelerator.trackers[0].writer.add_images(\n                    \"test_samples\", images_processed.transpose(0, 3, 1, 2), epoch\n                )\n\n            if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:\n                # save the model\n                if args.push_to_hub:\n                    push_to_hub(args, pipeline, repo, commit_message=f\"Epoch {epoch}\", blocking=False)\n                else:\n                    pipeline.save_pretrained(args.output_dir)\n        accelerator.wait_for_everyone()\n\n    accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description=\"Simple example of a training script.\")\n    parser.add_argument(\"--local_rank\", type=int, default=-1)\n    parser.add_argument(\"--dataset_name\", type=str, default=None)\n    parser.add_argument(\"--dataset_config_name\", type=str, default=None)\n    parser.add_argument(\"--train_data_dir\", type=str, default=None, help=\"A folder containing the training data.\")\n    parser.add_argument(\"--output_dir\", type=str, default=\"ddpm-model-64\")\n    parser.add_argument(\"--overwrite_output_dir\", action=\"store_true\")\n    parser.add_argument(\"--cache_dir\", type=str, default=None)\n    parser.add_argument(\"--resolution\", type=int, default=64)\n    parser.add_argument(\"--train_batch_size\", type=int, default=16)\n    parser.add_argument(\"--eval_batch_size\", type=int, default=16)\n    parser.add_argument(\"--num_epochs\", type=int, default=100)\n    parser.add_argument(\"--save_images_epochs\", type=int, default=10)\n    parser.add_argument(\"--save_model_epochs\", type=int, default=10)\n    parser.add_argument(\"--gradient_accumulation_steps\", type=int, default=1)\n    parser.add_argument(\"--learning_rate\", type=float, default=1e-4)\n    parser.add_argument(\"--lr_scheduler\", type=str, default=\"cosine\")\n    parser.add_argument(\"--lr_warmup_steps\", type=int, default=500)\n    parser.add_argument(\"--adam_beta1\", type=float, default=0.95)\n    parser.add_argument(\"--adam_beta2\", type=float, default=0.999)\n    parser.add_argument(\"--adam_weight_decay\", type=float, default=1e-6)\n    parser.add_argument(\"--adam_epsilon\", type=float, default=1e-08)\n    parser.add_argument(\"--use_ema\", action=\"store_true\", default=True)\n    parser.add_argument(\"--ema_inv_gamma\", type=float, default=1.0)\n    parser.add_argument(\"--ema_power\", type=float, default=3 / 4)\n    parser.add_argument(\"--ema_max_decay\", type=float, default=0.9999)\n    parser.add_argument(\"--push_to_hub\", action=\"store_true\")\n    parser.add_argument(\"--use_auth_token\", action=\"store_true\")\n    parser.add_argument(\"--hub_token\", type=str, default=None)\n    parser.add_argument(\"--hub_model_id\", type=str, default=None)\n    parser.add_argument(\"--hub_private_repo\", action=\"store_true\")\n    parser.add_argument(\"--logging_dir\", type=str, default=\"logs\")\n    parser.add_argument(\n        \"--mixed_precision\",\n        type=str,\n        default=\"no\",\n        choices=[\"no\", \"fp16\", \"bf16\"],\n        help=(\n            \"Whether to use mixed precision. Choose\"\n            \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n            \"and an Nvidia Ampere GPU.\"\n        ),\n    )\n\n    args = parser.parse_args()\n    env_local_rank = int(os.environ.get(\"LOCAL_RANK\", -1))\n    if env_local_rank != -1 and env_local_rank != args.local_rank:\n        args.local_rank = env_local_rank\n\n    if args.dataset_name is None and args.train_data_dir is None:\n        raise ValueError(\"You must specify either a dataset name from the hub or a train data directory.\")\n\n    main(args)\n", "path": "examples/unconditional_image_generation/train_unconditional.py"}], "after_files": [{"content": "import argparse\nimport math\nimport os\n\nimport torch\nimport torch.nn.functional as F\n\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom datasets import load_dataset\nfrom diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel\nfrom diffusers.hub_utils import init_git_repo, push_to_hub\nfrom diffusers.optimization import get_scheduler\nfrom diffusers.training_utils import EMAModel\nfrom torchvision.transforms import (\n    CenterCrop,\n    Compose,\n    InterpolationMode,\n    Normalize,\n    RandomHorizontalFlip,\n    Resize,\n    ToTensor,\n)\nfrom tqdm.auto import tqdm\n\n\nlogger = get_logger(__name__)\n\n\ndef main(args):\n    logging_dir = os.path.join(args.output_dir, args.logging_dir)\n    accelerator = Accelerator(\n        gradient_accumulation_steps=args.gradient_accumulation_steps,\n        mixed_precision=args.mixed_precision,\n        log_with=\"tensorboard\",\n        logging_dir=logging_dir,\n    )\n\n    model = UNet2DModel(\n        sample_size=args.resolution,\n        in_channels=3,\n        out_channels=3,\n        layers_per_block=2,\n        block_out_channels=(128, 128, 256, 256, 512, 512),\n        down_block_types=(\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"DownBlock2D\",\n            \"AttnDownBlock2D\",\n            \"DownBlock2D\",\n        ),\n        up_block_types=(\n            \"UpBlock2D\",\n            \"AttnUpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n            \"UpBlock2D\",\n        ),\n    )\n    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=args.learning_rate,\n        betas=(args.adam_beta1, args.adam_beta2),\n        weight_decay=args.adam_weight_decay,\n        eps=args.adam_epsilon,\n    )\n\n    augmentations = Compose(\n        [\n            Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),\n            CenterCrop(args.resolution),\n            RandomHorizontalFlip(),\n            ToTensor(),\n            Normalize([0.5], [0.5]),\n        ]\n    )\n\n    if args.dataset_name is not None:\n        dataset = load_dataset(\n            args.dataset_name,\n            args.dataset_config_name,\n            cache_dir=args.cache_dir,\n            use_auth_token=True if args.use_auth_token else None,\n            split=\"train\",\n        )\n    else:\n        dataset = load_dataset(\"imagefolder\", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split=\"train\")\n\n    def transforms(examples):\n        images = [augmentations(image.convert(\"RGB\")) for image in examples[\"image\"]]\n        return {\"input\": images}\n\n    dataset.set_transform(transforms)\n    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.train_batch_size, shuffle=True)\n\n    lr_scheduler = get_scheduler(\n        args.lr_scheduler,\n        optimizer=optimizer,\n        num_warmup_steps=args.lr_warmup_steps,\n        num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,\n    )\n\n    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n        model, optimizer, train_dataloader, lr_scheduler\n    )\n\n    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n\n    ema_model = EMAModel(model, inv_gamma=args.ema_inv_gamma, power=args.ema_power, max_value=args.ema_max_decay)\n\n    if args.push_to_hub:\n        repo = init_git_repo(args, at_init=True)\n\n    if accelerator.is_main_process:\n        run = os.path.split(__file__)[-1].split(\".\")[0]\n        accelerator.init_trackers(run)\n\n    global_step = 0\n    for epoch in range(args.num_epochs):\n        model.train()\n        progress_bar = tqdm(total=num_update_steps_per_epoch, disable=not accelerator.is_local_main_process)\n        progress_bar.set_description(f\"Epoch {epoch}\")\n        for step, batch in enumerate(train_dataloader):\n            clean_images = batch[\"input\"]\n            # Sample noise that we'll add to the images\n            noise = torch.randn(clean_images.shape).to(clean_images.device)\n            bsz = clean_images.shape[0]\n            # Sample a random timestep for each image\n            timesteps = torch.randint(\n                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=clean_images.device\n            ).long()\n\n            # Add noise to the clean images according to the noise magnitude at each timestep\n            # (this is the forward diffusion process)\n            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)\n\n            with accelerator.accumulate(model):\n                # Predict the noise residual\n                noise_pred = model(noisy_images, timesteps).sample\n                loss = F.mse_loss(noise_pred, noise)\n                accelerator.backward(loss)\n\n                if accelerator.sync_gradients:\n                    accelerator.clip_grad_norm_(model.parameters(), 1.0)\n                optimizer.step()\n                lr_scheduler.step()\n                if args.use_ema:\n                    ema_model.step(model)\n                optimizer.zero_grad()\n\n            # Checks if the accelerator has performed an optimization step behind the scenes\n            if accelerator.sync_gradients:\n                progress_bar.update(1)\n                global_step += 1\n\n            logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0], \"step\": global_step}\n            if args.use_ema:\n                logs[\"ema_decay\"] = ema_model.decay\n            progress_bar.set_postfix(**logs)\n            accelerator.log(logs, step=global_step)\n        progress_bar.close()\n\n        accelerator.wait_for_everyone()\n\n        # Generate sample images for visual inspection\n        if accelerator.is_main_process:\n            if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:\n                pipeline = DDPMPipeline(\n                    unet=accelerator.unwrap_model(ema_model.averaged_model if args.use_ema else model),\n                    scheduler=noise_scheduler,\n                )\n\n                generator = torch.manual_seed(0)\n                # run pipeline in inference (sample random noise and denoise)\n                images = pipeline(generator=generator, batch_size=args.eval_batch_size, output_type=\"numpy\").images\n\n                # denormalize the images and save to tensorboard\n                images_processed = (images * 255).round().astype(\"uint8\")\n                accelerator.trackers[0].writer.add_images(\n                    \"test_samples\", images_processed.transpose(0, 3, 1, 2), epoch\n                )\n\n            if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:\n                # save the model\n                if args.push_to_hub:\n                    push_to_hub(args, pipeline, repo, commit_message=f\"Epoch {epoch}\", blocking=False)\n                else:\n                    pipeline.save_pretrained(args.output_dir)\n        accelerator.wait_for_everyone()\n\n    accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description=\"Simple example of a training script.\")\n    parser.add_argument(\"--local_rank\", type=int, default=-1)\n    parser.add_argument(\"--dataset_name\", type=str, default=None)\n    parser.add_argument(\"--dataset_config_name\", type=str, default=None)\n    parser.add_argument(\"--train_data_dir\", type=str, default=None, help=\"A folder containing the training data.\")\n    parser.add_argument(\"--output_dir\", type=str, default=\"ddpm-model-64\")\n    parser.add_argument(\"--overwrite_output_dir\", action=\"store_true\")\n    parser.add_argument(\"--cache_dir\", type=str, default=None)\n    parser.add_argument(\"--resolution\", type=int, default=64)\n    parser.add_argument(\"--train_batch_size\", type=int, default=16)\n    parser.add_argument(\"--eval_batch_size\", type=int, default=16)\n    parser.add_argument(\"--num_epochs\", type=int, default=100)\n    parser.add_argument(\"--save_images_epochs\", type=int, default=10)\n    parser.add_argument(\"--save_model_epochs\", type=int, default=10)\n    parser.add_argument(\"--gradient_accumulation_steps\", type=int, default=1)\n    parser.add_argument(\"--learning_rate\", type=float, default=1e-4)\n    parser.add_argument(\"--lr_scheduler\", type=str, default=\"cosine\")\n    parser.add_argument(\"--lr_warmup_steps\", type=int, default=500)\n    parser.add_argument(\"--adam_beta1\", type=float, default=0.95)\n    parser.add_argument(\"--adam_beta2\", type=float, default=0.999)\n    parser.add_argument(\"--adam_weight_decay\", type=float, default=1e-6)\n    parser.add_argument(\"--adam_epsilon\", type=float, default=1e-08)\n    parser.add_argument(\"--use_ema\", action=\"store_true\", default=True)\n    parser.add_argument(\"--ema_inv_gamma\", type=float, default=1.0)\n    parser.add_argument(\"--ema_power\", type=float, default=3 / 4)\n    parser.add_argument(\"--ema_max_decay\", type=float, default=0.9999)\n    parser.add_argument(\"--push_to_hub\", action=\"store_true\")\n    parser.add_argument(\"--use_auth_token\", action=\"store_true\")\n    parser.add_argument(\"--hub_token\", type=str, default=None)\n    parser.add_argument(\"--hub_model_id\", type=str, default=None)\n    parser.add_argument(\"--hub_private_repo\", action=\"store_true\")\n    parser.add_argument(\"--logging_dir\", type=str, default=\"logs\")\n    parser.add_argument(\n        \"--mixed_precision\",\n        type=str,\n        default=\"no\",\n        choices=[\"no\", \"fp16\", \"bf16\"],\n        help=(\n            \"Whether to use mixed precision. Choose\"\n            \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n            \"and an Nvidia Ampere GPU.\"\n        ),\n    )\n\n    args = parser.parse_args()\n    env_local_rank = int(os.environ.get(\"LOCAL_RANK\", -1))\n    if env_local_rank != -1 and env_local_rank != args.local_rank:\n        args.local_rank = env_local_rank\n\n    if args.dataset_name is None and args.train_data_dir is None:\n        raise ValueError(\"You must specify either a dataset name from the hub or a train data directory.\")\n\n    main(args)\n", "path": "examples/unconditional_image_generation/train_unconditional.py"}]}