


'''simple docstring'''




from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block





@dataclass
class _snake_case   ( a__   ):
							"""simple docstring"""





							__a												=     42





class _snake_case   ( a__    ,   a__   ):
							"""simple docstring"""





							@register_to_config
							def __init__(   self			,				__A = None			,				__A = 3			,				__A = 3			,				__A = False			,				__A = "positional"			,				__A = 0			,				__A = True			,				__A = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")			,				__A = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")			,				__A = (224, 448, 672, 896)			,				__A = 2			,				__A = 1			,				__A = 1			,				__A = "conv"			,				__A = "conv"			,				__A = "silu"			,				__A = 8			,				__A = 32			,				__A = 1E-5			,				__A = "default"			,				__A = True			,				__A = None			,				__A = None			,				)     ->Optional[int]:



											"""simple docstring"""







											super().__init__()

											__lowerCamelCase		:      Union[str, Any]											=					sample_size
											__lowerCamelCase		:      Any											=					block_out_channels[0] * 4

											# Check inputs
											if len(__A	) != len(__A	):
															raise ValueError(
															    f'Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}.'	)

											if len(__A	) != len(__A	):
															raise ValueError(
															    f'Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.'	)

											# input
											__lowerCamelCase		:      Optional[Any]											=					nn.Convad(__A			,				block_out_channels[0]			,				kernel_size=3			,				padding=(1, 1)	)

											# time
											if time_embedding_type == "fourier":
															__lowerCamelCase		:      str											=					GaussianFourierProjection(embedding_size=block_out_channels[0]			,				scale=16	)
															__lowerCamelCase		:      Optional[Any]											=					2 * block_out_channels[0]
											elif time_embedding_type == "positional":
															__lowerCamelCase		:      Dict											=					Timesteps(block_out_channels[0]			,				__A			,				__A	)
															__lowerCamelCase		:      str											=					block_out_channels[0]

											__lowerCamelCase		:      Optional[Any]											=					TimestepEmbedding(__A			,				__A	)

											# class embedding
											if class_embed_type is None and num_class_embeds is not None:
															__lowerCamelCase		:      Tuple											=					nn.Embedding(__A			,				__A	)
											elif class_embed_type == "timestep":
															__lowerCamelCase		:      List[Any]											=					TimestepEmbedding(__A			,				__A	)
											elif class_embed_type == "identity":
															__lowerCamelCase		:      Any											=					nn.Identity(__A			,				__A	)
											else:
															__lowerCamelCase		:      List[str]											=					None

											__lowerCamelCase		:      Dict											=					nn.ModuleList([]	)
											__lowerCamelCase		:      Union[str, Any]											=					None
											__lowerCamelCase		:      str											=					nn.ModuleList([]	)

											# down
											__lowerCamelCase		:      Optional[int]											=					block_out_channels[0]
											for i, down_block_type in enumerate(__A	):
															__lowerCamelCase		:      List[Any]											=					output_channel
															__lowerCamelCase		:      Any											=					block_out_channels[i]
															__lowerCamelCase		:      Optional[int]											=					i == len(__A	) - 1

															__lowerCamelCase		:      Tuple											=					get_down_block(
															    __A			,				num_layers=__A			,				in_channels=__A			,				out_channels=__A			,				temb_channels=__A			,				add_downsample=not is_final_block			,				resnet_eps=__A			,				resnet_act_fn=__A			,				resnet_groups=__A			,				attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,				downsample_padding=__A			,				resnet_time_scale_shift=__A			,				downsample_type=__A			,				)
															self.down_blocks.append(__A	)

											# mid
											__lowerCamelCase		:      Dict											=					UNetMidBlockaD(
											    in_channels=block_out_channels[-1]			,				temb_channels=__A			,				resnet_eps=__A			,				resnet_act_fn=__A			,				output_scale_factor=__A			,				resnet_time_scale_shift=__A			,				attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]			,				resnet_groups=__A			,				add_attention=__A			,				)

											# up
											__lowerCamelCase		:      List[str]											=					list(reversed(__A	)	)
											__lowerCamelCase		:      Optional[int]											=					reversed_block_out_channels[0]
											for i, up_block_type in enumerate(__A	):
															__lowerCamelCase		:      List[Any]											=					output_channel
															__lowerCamelCase		:      Tuple											=					reversed_block_out_channels[i]
															__lowerCamelCase		:      Union[str, Any]											=					reversed_block_out_channels[min(i + 1			,				len(__A	) - 1	)]

															__lowerCamelCase		:      List[Any]											=					i == len(__A	) - 1

															__lowerCamelCase		:      Dict											=					get_up_block(
															    __A			,				num_layers=layers_per_block + 1			,				in_channels=__A			,				out_channels=__A			,				prev_output_channel=__A			,				temb_channels=__A			,				add_upsample=not is_final_block			,				resnet_eps=__A			,				resnet_act_fn=__A			,				resnet_groups=__A			,				attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,				resnet_time_scale_shift=__A			,				upsample_type=__A			,				)
															self.up_blocks.append(__A	)
															__lowerCamelCase		:      Optional[Any]											=					output_channel

											# out
											__lowerCamelCase		:      Tuple											=					norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,				32	)
											__lowerCamelCase		:      Union[str, Any]											=					nn.GroupNorm(num_channels=block_out_channels[0]			,				num_groups=__A			,				eps=__A	)
											__lowerCamelCase		:      Dict											=					nn.SiLU()
											__lowerCamelCase		:      str											=					nn.Convad(block_out_channels[0]			,				__A			,				kernel_size=3			,				padding=1	)







							def 		__SCREAMING_SNAKE_CASE(   self			,				__A			,				__A			,				__A = None			,				__A = True			,				)     ->Union[UNetaDOutput, Tuple]:



											"""simple docstring"""







											if self.config.center_input_sample:
															__lowerCamelCase		:      List[str]											=					2 * sample - 1.0

											# 1. time
											__lowerCamelCase		:      Any											=					timestep
											if not torch.is_tensor(__A	):
															__lowerCamelCase		:      List[str]											=					torch.tensor([timesteps]			,				dtype=torch.long			,				device=sample.device	)
											elif torch.is_tensor(__A	) and len(timesteps.shape	) == 0:
															__lowerCamelCase		:      List[str]											=					timesteps[None].to(sample.device	)

											# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
											__lowerCamelCase		:      Optional[int]											=					timesteps * torch.ones(sample.shape[0]			,				dtype=timesteps.dtype			,				device=timesteps.device	)

											__lowerCamelCase		:      Tuple											=					self.time_proj(__A	)

											# timesteps does not contain any weights and will always return f32 tensors
											# but time_embedding might actually be running in fp16. so we need to cast here.
											# there might be better ways to encapsulate this.
											__lowerCamelCase		:      List[Any]											=					t_emb.to(dtype=self.dtype	)
											__lowerCamelCase		:      List[str]											=					self.time_embedding(__A	)

											if self.class_embedding is not None:
															if class_labels is None:
																			raise ValueError("class_labels should be provided when doing class conditioning"	)

															if self.config.class_embed_type == "timestep":
																			__lowerCamelCase		:      Tuple											=					self.time_proj(__A	)

															__lowerCamelCase		:      int											=					self.class_embedding(__A	).to(dtype=self.dtype	)
															__lowerCamelCase		:      Optional[Any]											=					emb + class_emb

											# 2. pre-process
											__lowerCamelCase		:      str											=					sample
											__lowerCamelCase		:      int											=					self.conv_in(__A	)

											# 3. down
											__lowerCamelCase		:      Any											=					(sample,)
											for downsample_block in self.down_blocks:
															if hasattr(__A			,				"skip_conv"	):
																			__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      Optional[int]											=					downsample_block(
																			    hidden_states=__A			,				temb=__A			,				skip_sample=__A	)
															else:
																			__lowerCamelCase						,					__lowerCamelCase		:      int											=					downsample_block(hidden_states=__A			,				temb=__A	)

															down_block_res_samples += res_samples

											# 4. mid
											__lowerCamelCase		:      List[str]											=					self.mid_block(__A			,				__A	)

											# 5. up
											__lowerCamelCase		:      str											=					None
											for upsample_block in self.up_blocks:
															__lowerCamelCase		:      List[Any]											=					down_block_res_samples[-len(upsample_block.resnets	) :]
															__lowerCamelCase		:      Dict											=					down_block_res_samples[: -len(upsample_block.resnets	)]

															if hasattr(__A			,				"skip_conv"	):
																			__lowerCamelCase						,					__lowerCamelCase		:      Optional[Any]											=					upsample_block(__A			,				__A			,				__A			,				__A	)
															else:
																			__lowerCamelCase		:      Tuple											=					upsample_block(__A			,				__A			,				__A	)

        # 6. post-process
											__lowerCamelCase		:      List[Any]											=					self.conv_norm_out(__A	)
											__lowerCamelCase		:      Any											=					self.conv_act(__A	)
											__lowerCamelCase		:      Any											=					self.conv_out(__A	)

											if skip_sample is not None:
															sample += skip_sample

											if self.config.time_embedding_type == "fourier":
															__lowerCamelCase		:      Optional[Any]											=					timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]	)))	)
															__lowerCamelCase		:      List[str]											=					sample / timesteps

											if not return_dict:
															return (sample,)

											return UNetaDOutput(sample=__A	)









'''simple docstring'''




import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin





def 							a_       (						lowercase=2							,						lowercase=True							,						lowercase=False    )			->Optional[Any]:






				'''simple docstring'''

				__lowerCamelCase		:      List[str]											=					Accelerator(step_scheduler_with_optimizer=lowercase							,						split_batches=lowercase    )
				__lowerCamelCase		:      Union[str, Any]											=					torch.nn.Linear(2							,						4    )
				__lowerCamelCase		:      str											=					torch.optim.AdamW(model.parameters()							,						lr=1.0    )
				__lowerCamelCase		:      List[str]											=					torch.optim.lr_scheduler.OneCycleLR(lowercase							,						max_lr=0.0_1							,						steps_per_epoch=2							,						epochs=1    )
				__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      List[Any]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

				# Optimizer has stepped
				scheduler.step()
				if step_scheduler_with_optimizer or (num_processes == 1):
								assert (
								    scheduler.scheduler.last_epoch == num_processes
								), F'Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'
				else:
								assert (
								    scheduler.scheduler.last_epoch != num_processes
								), F'Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'





def 							a_       (						lowercase=2							,						lowercase=True							,						lowercase=False    )			->Dict:






				'''simple docstring'''

				__lowerCamelCase		:      Dict											=					Accelerator(step_scheduler_with_optimizer=lowercase							,						split_batches=lowercase    )
				__lowerCamelCase		:      Tuple											=					torch.nn.Linear(2							,						4    )
				__lowerCamelCase		:      Optional[Any]											=					torch.optim.AdamW(model.parameters()							,						lr=1.0    )
				__lowerCamelCase		:      Optional[Any]											=					torch.optim.lr_scheduler.LambdaLR(lowercase							,						lr_lambda=lambda lowercase    : 1 - n / 10    )
				__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      Optional[int]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

				# Optimizer has stepped
				__lowerCamelCase		:      List[Any]											=					False
				scheduler.step()
				__lowerCamelCase		:      List[Any]											=					1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
				assert (
				    scheduler.get_last_lr()[0] == expected_lr
				), F'Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'

				# Optimizer has not stepped
				__lowerCamelCase		:      Any											=					True
				scheduler.step()
				if not step_scheduler_with_optimizer:
								__lowerCamelCase		:      int											=					1 - 2 / 10
				assert (
				    scheduler.get_last_lr()[0] == expected_lr
				), F'Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'





def 							a_       (						lowercase = 2    )			->Dict:






				'''simple docstring'''

				from transformers import get_linear_schedule_with_warmup

				__lowerCamelCase		:      str											=					[1, 2, 4]
				for num_steps in steps:
								__lowerCamelCase		:      List[Any]											=					GradientAccumulationPlugin(num_steps=lowercase							,						adjust_scheduler=num_steps > 1    )
								__lowerCamelCase		:      Dict											=					Accelerator(gradient_accumulation_plugin=lowercase    )
								__lowerCamelCase		:      Any											=					torch.nn.Linear(2							,						4    )
								__lowerCamelCase		:      List[Any]											=					torch.optim.AdamW(model.parameters()							,						lr=1_0.0    )
								__lowerCamelCase		:      Optional[Any]											=					get_linear_schedule_with_warmup(optimizer=lowercase							,						num_warmup_steps=0							,						num_training_steps=20    )

								__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      Optional[int]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

								for i in range(10 * num_steps    ):
												with accelerator.accumulate(lowercase    ):
																optimizer.step()
																scheduler.step()

												if i == (10 * num_steps - 2):
																assert (
																    scheduler.get_last_lr()[0] != 0
																), F'Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'
								assert (
								    scheduler.get_last_lr()[0] == 0
								), F'Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'
								GradientState._reset_state()





@require_cpu
class _snake_case   ( unittest.TestCase   ):
							"""simple docstring"""





							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Tuple:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1	)			,				num_processes=1	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Tuple:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1	)			,				num_processes=1	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Optional[Any]:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1			,				step_scheduler_with_optimizer=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Dict:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1			,				step_scheduler_with_optimizer=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Optional[int]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(__A	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Union[str, Any]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(__A	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Any:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				step_scheduler_with_optimizer=__A	)	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->List[Any]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				step_scheduler_with_optimizer=__A	)	)







							@require_huggingface_suite
							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Dict:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				num_processes=1	)	)
											debug_launcher(__A	)





