





"""simple docstring"""







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class __magic_name__   (       A__		):
       UpperCamelCase    :						int      =	42





class __magic_name__   (       nn.Module		):
       def __init__(     self   ,  __magic_name__=3   ,  __magic_name__=3   ,  __magic_name__=("DownEncoderBlock2D",)   ,  __magic_name__=(6_4,)   ,  __magic_name__=2   ,  __magic_name__=3_2   ,  __magic_name__="silu"   ,  __magic_name__=True   ,  ):





         """simple docstring"""


         super().__init__()
         _lowerCAmelCase	      =					layers_per_block

         _lowerCAmelCase	      =					torch.nn.Convad(
             __magic_name__   ,  block_out_channels[0]   ,  kernel_size=3   ,  stride=1   ,  padding=1   ,  )

         _lowerCAmelCase	      =					None
         _lowerCAmelCase	      =					nn.ModuleList([]	)

         # down
         _lowerCAmelCase	      =					block_out_channels[0]
         for i, down_block_type in enumerate(__magic_name__	):
           _lowerCAmelCase	      =					output_channel
           _lowerCAmelCase	      =					block_out_channels[i]
           _lowerCAmelCase	      =					i == len(__magic_name__	) - 1

           _lowerCAmelCase	      =					get_down_block(
               __magic_name__   ,  num_layers=self.layers_per_block   ,  in_channels=__magic_name__   ,  out_channels=__magic_name__   ,  add_downsample=not is_final_block   ,  resnet_eps=1e-6   ,  downsample_padding=0   ,  resnet_act_fn=__magic_name__   ,  resnet_groups=__magic_name__   ,  attention_head_dim=__magic_name__   ,  temb_channels=__magic_name__   ,  )
           self.down_blocks.append(__magic_name__	)

         # mid
         _lowerCAmelCase	      =					UNetMidBlockaD(
             in_channels=block_out_channels[-1]   ,  resnet_eps=1e-6   ,  resnet_act_fn=__magic_name__   ,  output_scale_factor=1   ,  resnet_time_scale_shift='default'   ,  attention_head_dim=block_out_channels[-1]   ,  resnet_groups=__magic_name__   ,  temb_channels=__magic_name__   ,  )

         # out
         _lowerCAmelCase	      =					nn.GroupNorm(num_channels=block_out_channels[-1]   ,  num_groups=__magic_name__   ,  eps=1e-6	)
         _lowerCAmelCase	      =					nn.SiLU()

         _lowerCAmelCase	      =					2 * out_channels if double_z else out_channels
         _lowerCAmelCase	      =					nn.Convad(block_out_channels[-1]   ,  __magic_name__   ,  3   ,  padding=1	)

         _lowerCAmelCase	      =					False
       def 		_lowerCamelCase   (     self   ,  __magic_name__	):





         """simple docstring"""


         _lowerCAmelCase	      =					x
         _lowerCAmelCase	      =					self.conv_in(__magic_name__	)

         if self.training and self.gradient_checkpointing:

           def create_custom_forward(__magic_name__	):
             def custom_forward(*__magic_name__	):
               return module(*__magic_name__	)

             return custom_forward

           # down
           if is_torch_version('>='   ,  '1.11.0'	):
             for down_block in self.down_blocks:
               _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(
                   create_custom_forward(__magic_name__	)   ,  __magic_name__   ,  use_reentrant=__magic_name__	)
             # middle
             _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)   ,  __magic_name__   ,  use_reentrant=__magic_name__	)
           else:
             for down_block in self.down_blocks:
               _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(create_custom_forward(__magic_name__	)   ,  __magic_name__	)
             # middle
             _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block	)   ,  __magic_name__	)

         else:
           # down
           for down_block in self.down_blocks:
             _lowerCAmelCase	      =					down_block(__magic_name__	)

           # middle
           _lowerCAmelCase	      =					self.mid_block(__magic_name__	)

         # post-process
         _lowerCAmelCase	      =					self.conv_norm_out(__magic_name__	)
         _lowerCAmelCase	      =					self.conv_act(__magic_name__	)
         _lowerCAmelCase	      =					self.conv_out(__magic_name__	)

         return sample





class __magic_name__   (       nn.Module		):
       def __init__(     self   ,  __magic_name__=3   ,  __magic_name__=3   ,  __magic_name__=("UpDecoderBlock2D",)   ,  __magic_name__=(6_4,)   ,  __magic_name__=2   ,  __magic_name__=3_2   ,  __magic_name__="silu"   ,  __magic_name__="group"   ,  ):





         """simple docstring"""


         super().__init__()
         _lowerCAmelCase	      =					layers_per_block

         _lowerCAmelCase	      =					nn.Convad(
             __magic_name__   ,  block_out_channels[-1]   ,  kernel_size=3   ,  stride=1   ,  padding=1   ,  )

         _lowerCAmelCase	      =					None
         _lowerCAmelCase	      =					nn.ModuleList([]	)

         _lowerCAmelCase	      =					in_channels if norm_type == 'spatial' else None

         # mid
         _lowerCAmelCase	      =					UNetMidBlockaD(
             in_channels=block_out_channels[-1]   ,  resnet_eps=1e-6   ,  resnet_act_fn=__magic_name__   ,  output_scale_factor=1   ,  resnet_time_scale_shift='default' if norm_type == 'group' else norm_type   ,  attention_head_dim=block_out_channels[-1]   ,  resnet_groups=__magic_name__   ,  temb_channels=__magic_name__   ,  )

         # up
         _lowerCAmelCase	      =					list(reversed(__magic_name__	)	)
         _lowerCAmelCase	      =					reversed_block_out_channels[0]
         for i, up_block_type in enumerate(__magic_name__	):
           _lowerCAmelCase	      =					output_channel
           _lowerCAmelCase	      =					reversed_block_out_channels[i]

           _lowerCAmelCase	      =					i == len(__magic_name__	) - 1

           _lowerCAmelCase	      =					get_up_block(
               __magic_name__   ,  num_layers=self.layers_per_block + 1   ,  in_channels=__magic_name__   ,  out_channels=__magic_name__   ,  prev_output_channel=__magic_name__   ,  add_upsample=not is_final_block   ,  resnet_eps=1e-6   ,  resnet_act_fn=__magic_name__   ,  resnet_groups=__magic_name__   ,  attention_head_dim=__magic_name__   ,  temb_channels=__magic_name__   ,  resnet_time_scale_shift=__magic_name__   ,  )
           self.up_blocks.append(__magic_name__	)
           _lowerCAmelCase	      =					output_channel

         # out
         if norm_type == "spatial":
           _lowerCAmelCase	      =					SpatialNorm(block_out_channels[0]   ,  __magic_name__	)
         else:
           _lowerCAmelCase	      =					nn.GroupNorm(num_channels=block_out_channels[0]   ,  num_groups=__magic_name__   ,  eps=1e-6	)
         _lowerCAmelCase	      =					nn.SiLU()
         _lowerCAmelCase	      =					nn.Convad(block_out_channels[0]   ,  __magic_name__   ,  3   ,  padding=1	)

         _lowerCAmelCase	      =					False
       def 		_lowerCamelCase   (     self   ,  __magic_name__   ,  __magic_name__=None	):





         """simple docstring"""


         _lowerCAmelCase	      =					z
         _lowerCAmelCase	      =					self.conv_in(__magic_name__	)

         _lowerCAmelCase	      =					next(iter(self.up_blocks.parameters()	)	).dtype
         if self.training and self.gradient_checkpointing:

           def create_custom_forward(__magic_name__	):
             def custom_forward(*__magic_name__	):
               return module(*__magic_name__	)

             return custom_forward

           if is_torch_version('>='   ,  '1.11.0'	):
             # middle
             _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)   ,  __magic_name__   ,  __magic_name__   ,  use_reentrant=__magic_name__	)
             _lowerCAmelCase	      =					sample.to(__magic_name__	)

             # up
             for up_block in self.up_blocks:
               _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(
                   create_custom_forward(__magic_name__	)   ,  __magic_name__   ,  __magic_name__   ,  use_reentrant=__magic_name__	)
           else:
             # middle
             _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)   ,  __magic_name__   ,  __magic_name__	)
             _lowerCAmelCase	      =					sample.to(__magic_name__	)

             # up
             for up_block in self.up_blocks:
               _lowerCAmelCase	      =					torch.utils.checkpoint.checkpoint(create_custom_forward(__magic_name__	)   ,  __magic_name__   ,  __magic_name__	)
         else:
           # middle
           _lowerCAmelCase	      =					self.mid_block(__magic_name__   ,  __magic_name__	)
           _lowerCAmelCase	      =					sample.to(__magic_name__	)

           # up
           for up_block in self.up_blocks:
             _lowerCAmelCase	      =					up_block(__magic_name__   ,  __magic_name__	)

        # post-process
         if latent_embeds is None:
           _lowerCAmelCase	      =					self.conv_norm_out(__magic_name__	)
         else:
           _lowerCAmelCase	      =					self.conv_norm_out(__magic_name__   ,  __magic_name__	)
         _lowerCAmelCase	      =					self.conv_act(__magic_name__	)
         _lowerCAmelCase	      =					self.conv_out(__magic_name__	)

         return sample





class __magic_name__   (       nn.Module		):
       def __init__(     self   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__=None   ,  __magic_name__="random"   ,  __magic_name__=False   ,  __magic_name__=True	):





         """simple docstring"""


         super().__init__()
         _lowerCAmelCase	      =					n_e
         _lowerCAmelCase	      =					vq_embed_dim
         _lowerCAmelCase	      =					beta
         _lowerCAmelCase	      =					legacy

         _lowerCAmelCase	      =					nn.Embedding(self.n_e   ,  self.vq_embed_dim	)
         self.embedding.weight.data.uniform_(-1.0 / self.n_e   ,  1.0 / self.n_e	)

         _lowerCAmelCase	      =					remap
         if self.remap is not None:
           self.register_buffer('used'   ,  torch.tensor(np.load(self.remap	)	)	)
           _lowerCAmelCase	      =					self.used.shape[0]
           _lowerCAmelCase	      =					unknown_index  # "random" or "extra" or integer
           if self.unknown_index == "extra":
             _lowerCAmelCase	      =					self.re_embed
             _lowerCAmelCase	      =					self.re_embed + 1
           print(
               F'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
               F'''Using {self.unknown_index} for unknown indices.'''	)
         else:
           _lowerCAmelCase	      =					n_e

         _lowerCAmelCase	      =					sane_index_shape
       def 		_lowerCamelCase   (     self   ,  __magic_name__	):





         """simple docstring"""


         _lowerCAmelCase	      =					inds.shape
         assert len(__magic_name__	) > 1
         _lowerCAmelCase	      =					inds.reshape(ishape[0]   ,  -1	)
         _lowerCAmelCase	      =					self.used.to(__magic_name__	)
         _lowerCAmelCase	      =					(inds[:, :, None] == used[None, None, ...]).long()
         _lowerCAmelCase	      =					match.argmax(-1	)
         _lowerCAmelCase	      =					match.sum(2	) < 1
         if self.unknown_index == "random":
           _lowerCAmelCase	      =					torch.randint(0   ,  self.re_embed   ,  size=new[unknown].shape	).to(device=new.device	)
         else:
           _lowerCAmelCase	      =					self.unknown_index
         return new.reshape(__magic_name__	)
       def 		_lowerCamelCase   (     self   ,  __magic_name__	):





         """simple docstring"""


         _lowerCAmelCase	      =					inds.shape
         assert len(__magic_name__	) > 1
         _lowerCAmelCase	      =					inds.reshape(ishape[0]   ,  -1	)
         _lowerCAmelCase	      =					self.used.to(__magic_name__	)
         if self.re_embed > self.used.shape[0]:  # extra token
           _lowerCAmelCase	      =					0  # simply set to zero
         _lowerCAmelCase	      =					torch.gather(used[None, :][inds.shape[0] * [0], :]   ,  1   ,  __magic_name__	)
         return back.reshape(__magic_name__	)
       def 		_lowerCamelCase   (     self   ,  __magic_name__	):





         """simple docstring"""


         _lowerCAmelCase	      =					z.permute(0   ,  2   ,  3   ,  1	).contiguous()
         _lowerCAmelCase	      =					z.view(-1   ,  self.vq_embed_dim	)

         # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
         _lowerCAmelCase	      =					torch.argmin(torch.cdist(__magic_name__   ,  self.embedding.weight	)   ,  dim=1	)

         _lowerCAmelCase	      =					self.embedding(__magic_name__	).view(z.shape	)
         _lowerCAmelCase	      =					None
         _lowerCAmelCase	      =					None

         # compute loss for embedding
         if not self.legacy:
           _lowerCAmelCase	      =					self.beta * torch.mean((z_q.detach() - z) ** 2	) + torch.mean((z_q - z.detach()) ** 2	)
         else:
           _lowerCAmelCase	      =					torch.mean((z_q.detach() - z) ** 2	) + self.beta * torch.mean((z_q - z.detach()) ** 2	)

         # preserve gradients
         _lowerCAmelCase	      =					z + (z_q - z).detach()

         # reshape back to match original input shape
         _lowerCAmelCase	      =					z_q.permute(0   ,  3   ,  1   ,  2	).contiguous()

         if self.remap is not None:
           _lowerCAmelCase	      =					min_encoding_indices.reshape(z.shape[0]   ,  -1	)  # add batch axis
           _lowerCAmelCase	      =					self.remap_to_used(__magic_name__	)
           _lowerCAmelCase	      =					min_encoding_indices.reshape(-1   ,  1	)  # flatten

         if self.sane_index_shape:
           _lowerCAmelCase	      =					min_encoding_indices.reshape(z_q.shape[0]   ,  z_q.shape[2]   ,  z_q.shape[3]	)

         return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
       def 		_lowerCamelCase   (     self   ,  __magic_name__   ,  __magic_name__	):





         """simple docstring"""


         if self.remap is not None:
           _lowerCAmelCase	      =					indices.reshape(shape[0]   ,  -1	)  # add batch axis
           _lowerCAmelCase	      =					self.unmap_to_all(__magic_name__	)
           _lowerCAmelCase	      =					indices.reshape(-1	)  # flatten again

         # get quantized latent vectors
         _lowerCAmelCase	      =					self.embedding(__magic_name__	)

         if shape is not None:
           _lowerCAmelCase	      =					z_q.view(__magic_name__	)
           # reshape back to match original input shape
           _lowerCAmelCase	      =					z_q.permute(0   ,  3   ,  1   ,  2	).contiguous()

         return z_q





class __magic_name__   (       A__		):
       def __init__(     self   ,  __magic_name__   ,  __magic_name__=False	):





         """simple docstring"""


         _lowerCAmelCase	      =					parameters
         _lowerCAmelCase      ,   _lowerCAmelCase	      =					torch.chunk(__magic_name__   ,  2   ,  dim=1	)
         _lowerCAmelCase	      =					torch.clamp(self.logvar   ,  -30.0   ,  20.0	)
         _lowerCAmelCase	      =					deterministic
         _lowerCAmelCase	      =					torch.exp(0.5 * self.logvar	)
         _lowerCAmelCase	      =					torch.exp(self.logvar	)
         if self.deterministic:
           _lowerCAmelCase	      =					_lowerCAmelCase	      =					torch.zeros_like(
               self.mean   ,  device=self.parameters.device   ,  dtype=self.parameters.dtype	)
       def 		_lowerCamelCase   (     self   ,  __magic_name__ = None	):





         """simple docstring"""


         _lowerCAmelCase	      =					randn_tensor(
             self.mean.shape   ,  generator=__magic_name__   ,  device=self.parameters.device   ,  dtype=self.parameters.dtype	)
         _lowerCAmelCase	      =					self.mean + self.std * sample
         return x
       def 		_lowerCamelCase   (     self   ,  __magic_name__=None	):





         """simple docstring"""


         if self.deterministic:
           return torch.Tensor([0.0]	)
         else:
           if other is None:
             return 0.5 * torch.sum(torch.pow(self.mean   ,  2	) + self.var - 1.0 - self.logvar   ,  dim=[1, 2, 3]	)
           else:
             return 0.5 * torch.sum(
                 torch.pow(self.mean - other.mean   ,  2	) / other.var
                 + self.var / other.var
                 - 1.0
                 - self.logvar
                 + other.logvar   ,  dim=[1, 2, 3]   ,  )
       def 		_lowerCamelCase   (     self   ,  __magic_name__   ,  __magic_name__=[1, 2, 3]	):





         """simple docstring"""


         if self.deterministic:
           return torch.Tensor([0.0]	)
         _lowerCAmelCase	      =					np.log(2.0 * np.pi	)
         return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean   ,  2	) / self.var   ,  dim=__magic_name__	)
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


         return self.mean











# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

from accelerate.test_utils import execute_subprocess_async






def     a     (		lowerCamelCase_=None							):






				'''simple docstring'''
				if subparsers is not None:
								lowercase__  					=     subparsers.add_parser('''test'''							)
				else:
								lowercase__  					=     argparse.ArgumentParser('''Accelerate test command'''							)

				parser.add_argument(
				    '''--config_file'''						,  default=lowerCamelCase_						,  help=(
				        '''The path to use to store the config file. Will default to a file named default_config.yaml in the cache '''
				        '''location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '''
				        '''such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '''
				        '''with \'huggingface\'.'''
				    )						,  )

				if subparsers is not None:
								parser.set_defaults(func=lowerCamelCase_							)
				return parser






def     a     (		lowerCamelCase_							):






				'''simple docstring'''
				lowercase__  					=     os.path.sep.join(__file__.split(os.path.sep							)[:-2] + ['''test_utils''', '''scripts''', '''test_script.py''']							)

				if args.config_file is None:
								lowercase__  					=     script_name
				else:
								lowercase__  					=     F"""--config_file={args.config_file} {script_name}"""

				lowercase__  					=     ['''accelerate-launch'''] + test_args.split()
				lowercase__  					=     execute_subprocess_async(lowerCamelCase_						,  env=os.environ.copy()							)
				if result.returncode == 0:
								print('''Test is a success! You are ready for your distributed training!'''							)






def     a     (		):






				'''simple docstring'''
				lowercase__  					=     test_command_parser()
				lowercase__  					=     parser.parse_args()
				test_command(lowerCamelCase_							)


if __name__ == "__main__":
							main()




