






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

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder







@dataclass
class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    42







class 	_UpperCAmelCase  (  _A				,							_A		):




      """simple docstring"""



      A          =    True
      @register_to_config
      def __init__(   self    ,	_lowerCAmelCase = 3    ,	_lowerCAmelCase = 3    ,	_lowerCAmelCase = ("DownEncoderBlock2D",)    ,	_lowerCAmelCase = ("UpDecoderBlock2D",)    ,	_lowerCAmelCase = (64,)    ,	_lowerCAmelCase = 1    ,	_lowerCAmelCase = "silu"    ,	_lowerCAmelCase = 4    ,	_lowerCAmelCase = 32    ,	_lowerCAmelCase = 32    ,	_lowerCAmelCase = 0.1_8215    ,	):



          '''simple docstring'''
          super().__init__()

          # pass init params to Encoder
          lowerCAmelCase__   :str      							=				Encoder(
              in_channels=UpperCamelCase_    ,	out_channels=UpperCamelCase_    ,	down_block_types=UpperCamelCase_    ,	block_out_channels=UpperCamelCase_    ,	layers_per_block=UpperCamelCase_    ,	act_fn=UpperCamelCase_    ,	norm_num_groups=UpperCamelCase_    ,	double_z=UpperCamelCase_    ,	)

          # pass init params to Decoder
          lowerCAmelCase__   :Optional[int]      							=				Decoder(
              in_channels=UpperCamelCase_    ,	out_channels=UpperCamelCase_    ,	up_block_types=UpperCamelCase_    ,	block_out_channels=UpperCamelCase_    ,	layers_per_block=UpperCamelCase_    ,	norm_num_groups=UpperCamelCase_    ,	act_fn=UpperCamelCase_    ,	)

          lowerCAmelCase__   :Optional[int]      							=				nn.Convad(2 * latent_channels    ,	2 * latent_channels    ,	1		)
          lowerCAmelCase__   :Any      							=				nn.Convad(UpperCamelCase_    ,	UpperCamelCase_    ,	1		)

          lowerCAmelCase__   :int      							=				False
          lowerCAmelCase__   :Union[str, Any]      							=				False

          # only relevant if vae tiling is enabled
          lowerCAmelCase__   :Dict      							=				self.config.sample_size
          lowerCAmelCase__   :Any      							=				(
              self.config.sample_size[0]
              if isinstance(self.config.sample_size    ,	(list, tuple)		)
              else self.config.sample_size
          )
          lowerCAmelCase__   :str      							=				int(sample_size / (2 ** (len(self.config.block_out_channels		) - 1))		)
          lowerCAmelCase__   :List[Any]      							=				0.25
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=False		):



          '''simple docstring'''





          if isinstance(UpperCamelCase_    ,	(Encoder, Decoder)		):
              lowerCAmelCase__   :str      							=				value
      def 			snake_case_      (   self    ,	_lowerCAmelCase = True		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				use_tiling
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.enable_tiling(UpperCamelCase_		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				True
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				False
      @property
      # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				{}

          def fn_recursive_add_processors(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):
              if hasattr(UpperCamelCase_    ,	"set_processor"		):
                  lowerCAmelCase__   :List[str]      							=				module.processor

              for sub_name, child in module.named_children():
                  fn_recursive_add_processors(f'''{name}.{sub_name}'''    ,	UpperCamelCase_    ,	UpperCamelCase_		)

              return processors

          for name, module in self.named_children():
              fn_recursive_add_processors(UpperCamelCase_    ,	UpperCamelCase_    ,	UpperCamelCase_		)

          return processors
      def 			snake_case_      (   self    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				len(self.attn_processors.keys()		)

          if isinstance(UpperCamelCase_    ,	UpperCamelCase_		) and len(UpperCamelCase_		) != count:
              raise ValueError(
                  f'''A dict of processors was passed, but the number of processors {len(UpperCamelCase_		)} does not match the'''
                  f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''		)

          def fn_recursive_attn_processor(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):
              if hasattr(UpperCamelCase_    ,	"set_processor"		):
                  if not isinstance(UpperCamelCase_    ,	UpperCamelCase_		):
                      module.set_processor(UpperCamelCase_		)
                  else:
                      module.set_processor(processor.pop(f'''{name}.processor'''		)		)

              for sub_name, child in module.named_children():
                  fn_recursive_attn_processor(f'''{name}.{sub_name}'''    ,	UpperCamelCase_    ,	UpperCamelCase_		)

          for name, module in self.named_children():
              fn_recursive_attn_processor(UpperCamelCase_    ,	UpperCamelCase_    ,	UpperCamelCase_		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.set_attn_processor(AttnProcessor()		)
      @apply_forward_hook
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = True		):



          '''simple docstring'''





          if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
              return self.tiled_encode(UpperCamelCase_    ,	return_dict=UpperCamelCase_		)

          if self.use_slicing and x.shape[0] > 1:
              lowerCAmelCase__   :Optional[Any]      							=				[self.encoder(UpperCamelCase_		) for x_slice in x.split(1		)]
              lowerCAmelCase__   :Optional[Any]      							=				torch.cat(UpperCamelCase_		)
          else:
              lowerCAmelCase__   :Any      							=				self.encoder(UpperCamelCase_		)

          lowerCAmelCase__   :Tuple      							=				self.quant_conv(UpperCamelCase_		)
          lowerCAmelCase__   :Any      							=				DiagonalGaussianDistribution(UpperCamelCase_		)

          if not return_dict:
              return (posterior,)

          return AutoencoderKLOutput(latent_dist=UpperCamelCase_		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = True		):



          '''simple docstring'''





          if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
              return self.tiled_decode(UpperCamelCase_    ,	return_dict=UpperCamelCase_		)

          lowerCAmelCase__   :Tuple      							=				self.post_quant_conv(UpperCamelCase_		)
          lowerCAmelCase__   :int      							=				self.decoder(UpperCamelCase_		)

          if not return_dict:
              return (dec,)

          return DecoderOutput(sample=UpperCamelCase_		)
      @apply_forward_hook
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = True		):



          '''simple docstring'''





          if self.use_slicing and z.shape[0] > 1:
              lowerCAmelCase__   :int      							=				[self._decode(UpperCamelCase_		).sample for z_slice in z.split(1		)]
              lowerCAmelCase__   :Dict      							=				torch.cat(UpperCamelCase_		)
          else:
              lowerCAmelCase__   :List[str]      							=				self._decode(UpperCamelCase_		).sample

          if not return_dict:
              return (decoded,)

          return DecoderOutput(sample=UpperCamelCase_		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				min(a.shape[2]    ,	b.shape[2]    ,	UpperCamelCase_		)
          for y in range(UpperCamelCase_		):
              lowerCAmelCase__   :Optional[int]      							=				a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
          return b
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				min(a.shape[3]    ,	b.shape[3]    ,	UpperCamelCase_		)
          for x in range(UpperCamelCase_		):
              lowerCAmelCase__   :Optional[int]      							=				a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
          return b
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = True		):



          '''simple docstring'''
          lowerCAmelCase__   :List[str]      							=				int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)		)
          lowerCAmelCase__   :Dict      							=				int(self.tile_latent_min_size * self.tile_overlap_factor		)
          lowerCAmelCase__   :List[Any]      							=				self.tile_latent_min_size - blend_extent

          # Split the image into 512x512 tiles and encode them separately.
          lowerCAmelCase__   :int      							=				[]
          for i in range(0    ,	x.shape[2]    ,	UpperCamelCase_		):
              lowerCAmelCase__   :List[str]      							=				[]
              for j in range(0    ,	x.shape[3]    ,	UpperCamelCase_		):
                  lowerCAmelCase__   :Dict      							=				x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                  lowerCAmelCase__   :Dict      							=				self.encoder(UpperCamelCase_		)
                  lowerCAmelCase__   :Optional[int]      							=				self.quant_conv(UpperCamelCase_		)
                  row.append(UpperCamelCase_		)
              rows.append(UpperCamelCase_		)
          lowerCAmelCase__   :Any      							=				[]
          for i, row in enumerate(UpperCamelCase_		):
              lowerCAmelCase__   :List[Any]      							=				[]
              for j, tile in enumerate(UpperCamelCase_		):
                  # blend the above tile and the left tile
                  # to the current tile and add the current tile to the result row
                  if i > 0:
                      lowerCAmelCase__   :List[str]      							=				self.blend_v(rows[i - 1][j]    ,	UpperCamelCase_    ,	UpperCamelCase_		)
                  if j > 0:
                      lowerCAmelCase__   :Any      							=				self.blend_h(row[j - 1]    ,	UpperCamelCase_    ,	UpperCamelCase_		)
                  result_row.append(tile[:, :, :row_limit, :row_limit]		)
              result_rows.append(torch.cat(UpperCamelCase_    ,	dim=3		)		)

          lowerCAmelCase__   :Tuple      							=				torch.cat(UpperCamelCase_    ,	dim=2		)
          lowerCAmelCase__   :int      							=				DiagonalGaussianDistribution(UpperCamelCase_		)

          if not return_dict:
              return (posterior,)

          return AutoencoderKLOutput(latent_dist=UpperCamelCase_		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = True		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)		)
          lowerCAmelCase__   :List[Any]      							=				int(self.tile_sample_min_size * self.tile_overlap_factor		)
          lowerCAmelCase__   :Dict      							=				self.tile_sample_min_size - blend_extent

          # Split z into overlapping 64x64 tiles and decode them separately.
          # The tiles have an overlap to avoid seams between tiles.
          lowerCAmelCase__   :str      							=				[]
          for i in range(0    ,	z.shape[2]    ,	UpperCamelCase_		):
              lowerCAmelCase__   :Dict      							=				[]
              for j in range(0    ,	z.shape[3]    ,	UpperCamelCase_		):
                  lowerCAmelCase__   :Optional[int]      							=				z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                  lowerCAmelCase__   :List[str]      							=				self.post_quant_conv(UpperCamelCase_		)
                  lowerCAmelCase__   :Union[str, Any]      							=				self.decoder(UpperCamelCase_		)
                  row.append(UpperCamelCase_		)
              rows.append(UpperCamelCase_		)
          lowerCAmelCase__   :int      							=				[]
          for i, row in enumerate(UpperCamelCase_		):
              lowerCAmelCase__   :Optional[Any]      							=				[]
              for j, tile in enumerate(UpperCamelCase_		):
                  # blend the above tile and the left tile
                  # to the current tile and add the current tile to the result row
                  if i > 0:
                      lowerCAmelCase__   :Tuple      							=				self.blend_v(rows[i - 1][j]    ,	UpperCamelCase_    ,	UpperCamelCase_		)
                  if j > 0:
                      lowerCAmelCase__   :Optional[int]      							=				self.blend_h(row[j - 1]    ,	UpperCamelCase_    ,	UpperCamelCase_		)
                  result_row.append(tile[:, :, :row_limit, :row_limit]		)
              result_rows.append(torch.cat(UpperCamelCase_    ,	dim=3		)		)

          lowerCAmelCase__   :Optional[int]      							=				torch.cat(UpperCamelCase_    ,	dim=2		)
          if not return_dict:
              return (dec,)

          return DecoderOutput(sample=UpperCamelCase_		)



      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = False    ,	_lowerCAmelCase = True    ,	_lowerCAmelCase = None    ,	):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				sample
          lowerCAmelCase__   :List[Any]      							=				self.encode(UpperCamelCase_		).latent_dist
          if sample_posterior:
              lowerCAmelCase__   :Optional[Any]      							=				posterior.sample(generator=UpperCamelCase_		)
          else:
              lowerCAmelCase__   :str      							=				posterior.mode()
          lowerCAmelCase__   :Union[str, Any]      							=				self.decode(UpperCamelCase_		).sample

          if not return_dict:
              return (dec,)

          return DecoderOutput(sample=UpperCamelCase_		)











from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_a    :   Tuple        =			logging.get_logger(__name__)

_a    :   Optional[int]        =			{
    """facebook/levit-128S""": """https://huggingface.co/facebook/levit-128S/resolve/main/config.json""",
    # See all LeViT models at https://huggingface.co/models?filter=levit
}







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    '''levit'''
      def __init__(   self    ,	_lowerCAmelCase=224    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=1    ,	_lowerCAmelCase=16    ,	_lowerCAmelCase=[128, 256, 384]    ,	_lowerCAmelCase=[4, 8, 12]    ,	_lowerCAmelCase=[4, 4, 4]    ,	_lowerCAmelCase=[16, 16, 16]    ,	_lowerCAmelCase=0    ,	_lowerCAmelCase=[2, 2, 2]    ,	_lowerCAmelCase=[2, 2, 2]    ,	_lowerCAmelCase=0.02    ,	**_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__(**_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				image_size
          lowerCAmelCase__   :Dict      							=				num_channels
          lowerCAmelCase__   :List[Any]      							=				kernel_size
          lowerCAmelCase__   :List[Any]      							=				stride
          lowerCAmelCase__   :Tuple      							=				padding
          lowerCAmelCase__   :str      							=				hidden_sizes
          lowerCAmelCase__   :Union[str, Any]      							=				num_attention_heads
          lowerCAmelCase__   :Optional[int]      							=				depths
          lowerCAmelCase__   :Union[str, Any]      							=				key_dim
          lowerCAmelCase__   :str      							=				drop_path_rate
          lowerCAmelCase__   :int      							=				patch_size
          lowerCAmelCase__   :Dict      							=				attention_ratio
          lowerCAmelCase__   :Dict      							=				mlp_ratio
          lowerCAmelCase__   :Optional[int]      							=				initializer_range
          lowerCAmelCase__   :Tuple      							=				[
              ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
              ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
          ]







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    version.parse('''1.11'''		)
      @property
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return OrderedDict(
              [
                  ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
              ]		)



      @property
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return 1e-4



