code: 
  
  
  
  
  
  
  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__      ):
     __lowerCAmelCase      =						42
  
  
  
  
  
  
  class 					snake_case_							(       a__      ,			a__      ):
  
     @register_to_config
     def __init__(    self:	str   ,       snake_case__:	Any = None   ,       snake_case__:	Any = 3   ,       snake_case__:	Any = 3   ,       snake_case__:	Any = False   ,       snake_case__:	int = "positional"   ,       snake_case__:	Tuple = 0   ,       snake_case__:	Dict = True   ,       snake_case__:	Dict = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")   ,       snake_case__:	List[Any] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")   ,       snake_case__:	List[Any] = (224, 448, 672, 896)   ,       snake_case__:	Optional[Any] = 2   ,       snake_case__:	Tuple = 1   ,       snake_case__:	Any = 1   ,       snake_case__:	Dict = "conv"   ,       snake_case__:	List[Any] = "conv"   ,       snake_case__:	Union[str, Any] = "silu"   ,       snake_case__:	Optional[int] = 8   ,       snake_case__:	Dict = 32   ,       snake_case__:	List[Any] = 1E-5   ,       snake_case__:	Any = "default"   ,       snake_case__:	List[Any] = True   ,       snake_case__:	List[str] = None   ,       snake_case__:	List[Any] = None   ,       )  ->		Optional[int]:
  
  
         """simple docstring"""
  
  
  
         super().__init__()
  
         _lowercase     =	sample_size
         _lowercase     =	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
         _lowercase     =	nn.Convad(__A   ,       block_out_channels[0]   ,       kernel_size=3   ,       padding=(1, 1)	)
  
         # time
         if time_embedding_type == "fourier":
             _lowercase     =	GaussianFourierProjection(embedding_size=block_out_channels[0]   ,       scale=16	)
             _lowercase     =	2 * block_out_channels[0]
         elif time_embedding_type == "positional":
             _lowercase     =	Timesteps(block_out_channels[0]   ,       __A   ,       __A	)
             _lowercase     =	block_out_channels[0]
  
         _lowercase     =	TimestepEmbedding(__A   ,       __A	)
  
         # class embedding
         if class_embed_type is None and num_class_embeds is not None:
             _lowercase     =	nn.Embedding(__A   ,       __A	)
         elif class_embed_type == "timestep":
             _lowercase     =	TimestepEmbedding(__A   ,       __A	)
         elif class_embed_type == "identity":
             _lowercase     =	nn.Identity(__A   ,       __A	)
         else:
             _lowercase     =	None
  
         _lowercase     =	nn.ModuleList([]	)
         _lowercase     =	None
         _lowercase     =	nn.ModuleList([]	)
  
         # down
         _lowercase     =	block_out_channels[0]
         for i, down_block_type in enumerate(__A	):
             _lowercase     =	output_channel
             _lowercase     =	block_out_channels[i]
             _lowercase     =	i == len(__A	) - 1
  
             _lowercase     =	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
         _lowercase     =	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
         _lowercase     =	list(reversed(__A	)	)
         _lowercase     =	reversed_block_out_channels[0]
         for i, up_block_type in enumerate(__A	):
             _lowercase     =	output_channel
             _lowercase     =	reversed_block_out_channels[i]
             _lowercase     =	reversed_block_out_channels[min(i + 1   ,       len(__A	) - 1	)]
  
             _lowercase     =	i == len(__A	) - 1
  
             _lowercase     =	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	)
             _lowercase     =	output_channel
  
         # out
         _lowercase     =	norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4   ,       32	)
         _lowercase     =	nn.GroupNorm(num_channels=block_out_channels[0]   ,       num_groups=__A   ,       eps=__A	)
         _lowercase     =	nn.SiLU()
         _lowercase     =	nn.Convad(block_out_channels[0]   ,       __A   ,       kernel_size=3   ,       padding=1	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	str   ,       snake_case__:	Optional[int]   ,       snake_case__:	Dict   ,       snake_case__:	Union[str, Any] = None   ,       snake_case__:	Any = True   ,       )  ->		Union[UNetaDOutput, Tuple]:
  
  
         """simple docstring"""
  
  
  
         if self.config.center_input_sample:
             _lowercase     =	2 * sample - 1.0
  
         # 1. time
         _lowercase     =	timestep
         if not torch.is_tensor(__A	):
             _lowercase     =	torch.tensor([timesteps]   ,       dtype=torch.long   ,       device=sample.device	)
         elif torch.is_tensor(__A	) and len(timesteps.shape	) == 0:
             _lowercase     =	timesteps[None].to(sample.device	)
  
         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         _lowercase     =	timesteps * torch.ones(sample.shape[0]   ,       dtype=timesteps.dtype   ,       device=timesteps.device	)
  
         _lowercase     =	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.
         _lowercase     =	t_emb.to(dtype=self.dtype	)
         _lowercase     =	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":
                 _lowercase     =	self.time_proj(__A	)
  
             _lowercase     =	self.class_embedding(__A	).to(dtype=self.dtype	)
             _lowercase     =	emb + class_emb
  
         # 2. pre-process
         _lowercase     =	sample
         _lowercase     =	self.conv_in(__A	)
  
         # 3. down
         _lowercase     =	(sample,)
         for downsample_block in self.down_blocks:
             if hasattr(__A   ,       'skip_conv'	):
                 _lowercase     =	downsample_block(
                     hidden_states=__A   ,       temb=__A   ,       skip_sample=__A	)
             else:
                 _lowercase     =	downsample_block(hidden_states=__A   ,       temb=__A	)
  
             down_block_res_samples += res_samples
  
         # 4. mid
         _lowercase     =	self.mid_block(__A   ,       __A	)
  
         # 5. up
         _lowercase     =	None
         for upsample_block in self.up_blocks:
             _lowercase     =	down_block_res_samples[-len(upsample_block.resnets	) :]
             _lowercase     =	down_block_res_samples[: -len(upsample_block.resnets	)]
  
             if hasattr(__A   ,       'skip_conv'	):
                 _lowercase     =	upsample_block(__A   ,       __A   ,       __A   ,       __A	)
             else:
                 _lowercase     =	upsample_block(__A   ,       __A   ,       __A	)
  
          # 6. post-process
         _lowercase     =	self.conv_norm_out(__A	)
         _lowercase     =	self.conv_act(__A	)
         _lowercase     =	self.conv_out(__A	)
  
         if skip_sample is not None:
             sample += skip_sample
  
         if self.config.time_embedding_type == "fourier":
             _lowercase     =	timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]	)))	)
             _lowercase     =	sample / timesteps
  
         if not return_dict:
             return (sample,)
  
         return UNetaDOutput(sample=__A	)
  
style_context: 
  
  
  
  
  
  
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  
  from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
  from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
  from ...image_utils import (
      IMAGENET_STANDARD_MEAN,
      IMAGENET_STANDARD_STD,
      ChannelDimension,
      ImageInput,
      PILImageResampling,
      make_list_of_images,
      to_numpy_array,
      valid_images,
  )
  from ...utils import TensorType, is_vision_available, logging
  
  
  if is_vision_available():
     import PIL
  
  
  lowerCamelCase										=  logging.get_logger(__name__)
  
  class 					snake_case_							(       lowercase_      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
     __lowerCAmelCase      =						["""pixel_values"""]
  
     def __init__(    self:	int   ,       snake_case__:	bool = True   ,       snake_case__:	Dict[str, int] = None   ,       snake_case__:	PILImageResampling = PIL.Image.NEAREST   ,       snake_case__:	bool = False   ,       snake_case__:	Dict[str, int] = None   ,       snake_case__:	Union[int, float] = 1 / 255   ,       snake_case__:	bool = False   ,       snake_case__:	bool = True   ,       snake_case__:	bool = True   ,       snake_case__:	Optional[Union[float, List[float]]] = None   ,       snake_case__:	Optional[Union[float, List[float]]] = None   ,       snake_case__:	bool = True   ,       **snake_case__:	Optional[Any]   ,       )  ->		None:
  
  
         """simple docstring"""
  
  
  
         super().__init__(**snake_case__	)
         _lowercase     =	size if size is not None else {'height': 346, 'width': 346}
         _lowercase     =	get_size_dict(snake_case__	)
         _lowercase     =	crop_size if crop_size is not None else {'height': 289, 'width': 289}
         _lowercase     =	get_size_dict(snake_case__   ,       param_name='crop_size'	)
  
         _lowercase     =	do_resize
         _lowercase     =	size
         _lowercase     =	resample
         _lowercase     =	do_center_crop
         _lowercase     =	crop_size
         _lowercase     =	do_rescale
         _lowercase     =	rescale_factor
         _lowercase     =	rescale_offset
         _lowercase     =	do_normalize
         _lowercase     =	image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
         _lowercase     =	image_std if image_std is not None else IMAGENET_STANDARD_STD
         _lowercase     =	include_top
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	str   ,       snake_case__:	np.ndarray   ,       snake_case__:	Dict[str, int]   ,       snake_case__:	PILImageResampling = PIL.Image.NEAREST   ,       snake_case__:	Optional[Union[str, ChannelDimension]] = None   ,       **snake_case__:	Optional[int]   ,       )  ->		np.ndarray:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	get_size_dict(snake_case__	)
         if "height" not in size or "width" not in size:
             raise ValueError(f'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'''	)
         return resize(
             snake_case__   ,       size=(size['height'], size['width'])   ,       resample=snake_case__   ,       data_format=snake_case__   ,       **snake_case__	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Optional[Any]   ,       snake_case__:	np.ndarray   ,       snake_case__:	Dict[str, int]   ,       snake_case__:	Optional[Union[str, ChannelDimension]] = None   ,       **snake_case__:	str   ,       )  ->		np.ndarray:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	get_size_dict(snake_case__	)
         if "height" not in size or "width" not in size:
             raise ValueError(f'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'''	)
         return center_crop(snake_case__   ,       size=(size['height'], size['width'])   ,       data_format=snake_case__   ,       **snake_case__	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Optional[Any]   ,       snake_case__:	np.ndarray   ,       snake_case__:	Union[int, float]   ,       snake_case__:	bool = True   ,       snake_case__:	Optional[Union[str, ChannelDimension]] = None   ,       **snake_case__:	List[str]   ,       )  ->		Optional[int]:
  
  
         """simple docstring"""
  
  
  
         if offset:
             _lowercase     =	(image - 1_27.5) * scale
             if data_format is not None:
                 _lowercase     =	to_channel_dimension_format(snake_case__   ,       snake_case__	)
             _lowercase     =	rescaled_image.astype(np.floataa	)
         else:
             _lowercase     =	rescale(snake_case__   ,       scale=snake_case__   ,       data_format=snake_case__   ,       **snake_case__	)
         return rescaled_image
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Dict   ,       snake_case__:	np.ndarray   ,       snake_case__:	Union[float, List[float]]   ,       snake_case__:	Union[float, List[float]]   ,       snake_case__:	Optional[Union[str, ChannelDimension]] = None   ,       **snake_case__:	Any   ,       )  ->		np.ndarray:
  
  
         """simple docstring"""
  
  
  
         return normalize(snake_case__   ,       mean=snake_case__   ,       std=snake_case__   ,       data_format=snake_case__   ,       **snake_case__	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	Optional[int]   ,       snake_case__:	ImageInput   ,       snake_case__:	bool = None   ,       snake_case__:	Dict[str, int] = None   ,       snake_case__:	Union[str, Any]=None   ,       snake_case__:	bool = None   ,       snake_case__:	Dict[str, int] = None   ,       snake_case__:	bool = None   ,       snake_case__:	float = None   ,       snake_case__:	bool = None   ,       snake_case__:	bool = None   ,       snake_case__:	Optional[Union[float, List[float]]] = None   ,       snake_case__:	Optional[Union[float, List[float]]] = None   ,       snake_case__:	bool = None   ,       snake_case__:	Optional[Union[str, TensorType]] = None   ,       snake_case__:	ChannelDimension = ChannelDimension.FIRST   ,       **snake_case__:	int   ,       )  ->		PIL.Image.Image:
  
  
         """simple docstring"""
  
  
  
         _lowercase     =	do_resize if do_resize is not None else self.do_resize
         _lowercase     =	resample if resample is not None else self.resample
         _lowercase     =	do_center_crop if do_center_crop is not None else self.do_center_crop
         _lowercase     =	do_rescale if do_rescale is not None else self.do_rescale
         _lowercase     =	rescale_factor if rescale_factor is not None else self.rescale_factor
         _lowercase     =	rescale_offset if rescale_offset is not None else self.rescale_offset
         _lowercase     =	do_normalize if do_normalize is not None else self.do_normalize
         _lowercase     =	image_mean if image_mean is not None else self.image_mean
         _lowercase     =	image_std if image_std is not None else self.image_std
         _lowercase     =	include_top if include_top is not None else self.include_top
  
         _lowercase     =	size if size is not None else self.size
         _lowercase     =	get_size_dict(snake_case__	)
         _lowercase     =	crop_size if crop_size is not None else self.crop_size
         _lowercase     =	get_size_dict(snake_case__   ,       param_name='crop_size'	)
  
         _lowercase     =	make_list_of_images(snake_case__	)
  
         if not valid_images(snake_case__	):
             raise ValueError(
                 'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
                 'torch.Tensor, tf.Tensor or jax.ndarray.'	)
  
         if do_resize and size is None or resample is None:
             raise ValueError('Size and resample must be specified if do_resize is True.'	)
  
         if do_center_crop and crop_size is None:
             raise ValueError('Crop size must be specified if do_center_crop is True.'	)
  
         if do_rescale and rescale_factor is None:
             raise ValueError('Rescale factor must be specified if do_rescale is True.'	)
  
         if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('Image mean and std must be specified if do_normalize is True.'	)
  
         # All transformations expect numpy arrays.
         _lowercase     =	[to_numpy_array(snake_case__	) for image in images]
  
         if do_resize:
             _lowercase     =	[self.resize(image=snake_case__   ,       size=snake_case__   ,       resample=snake_case__	) for image in images]
  
         if do_center_crop:
             _lowercase     =	[self.center_crop(image=snake_case__   ,       size=snake_case__	) for image in images]
  
         if do_rescale:
             _lowercase     =	[self.rescale(image=snake_case__   ,       scale=snake_case__   ,       offset=snake_case__	) for image in images]
  
         if do_normalize:
             _lowercase     =	[self.normalize(image=snake_case__   ,       mean=snake_case__   ,       std=snake_case__	) for image in images]
  
         if include_top:
             _lowercase     =	[self.normalize(image=snake_case__   ,       mean=[0, 0, 0]   ,       std=snake_case__	) for image in images]
  
         _lowercase     =	[to_channel_dimension_format(snake_case__   ,       snake_case__	) for image in images]
  
         _lowercase     =	{'pixel_values': images}
         return BatchFeature(data=snake_case__   ,       tensor_type=snake_case__	)
  