




'''simple docstring'''



import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import UperNetForSemanticSegmentation
      from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor






class 							_a				:





       '''simple docstring'''






       def __init__(      self,							A,							A=13,							A=32,							A=3,							A=4,							A=[10, 20, 30, 40],							A=[2, 2, 3, 2],							A=True,							A=True,							A=37,							A="gelu",							A=10,							A=0.02,							A=["stage2", "stage3", "stage4"],							A=3,							A=None,							):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[str]    					=						parent
             SCREAMING_SNAKE_CASE   :      str    					=						batch_size
             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						image_size
             SCREAMING_SNAKE_CASE   :      Tuple    					=						num_channels
             SCREAMING_SNAKE_CASE   :      int    					=						num_stages
             SCREAMING_SNAKE_CASE   :      List[str]    					=						hidden_sizes
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						depths
             SCREAMING_SNAKE_CASE   :      Dict    					=						is_training
             SCREAMING_SNAKE_CASE   :      Dict    					=						use_labels
             SCREAMING_SNAKE_CASE   :      int    					=						intermediate_size
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						hidden_act
             SCREAMING_SNAKE_CASE   :      Dict    					=						type_sequence_label_size
             SCREAMING_SNAKE_CASE   :      Any    					=						initializer_range
             SCREAMING_SNAKE_CASE   :      List[Any]    					=						out_features
             SCREAMING_SNAKE_CASE   :      Any    					=						num_labels
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						scope
             SCREAMING_SNAKE_CASE   :      Any    					=						num_stages







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[Any]    					=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

             SCREAMING_SNAKE_CASE   :      int    					=						None
             if self.use_labels:
                   SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						ids_tensor([self.batch_size],							self.type_sequence_label_size  )

             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						self.get_config()

             return config, pixel_values, labels







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             return ConvNextConfig(
                 num_channels=self.num_channels,							num_stages=self.num_stages,							hidden_sizes=self.hidden_sizes,							depths=self.depths,							is_training=self.is_training,							intermediate_size=self.intermediate_size,							hidden_act=self.hidden_act,							out_features=self.out_features,							)







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             return UperNetConfig(
                 backbone_config=self.get_backbone_config(),							hidden_size=512,							pool_scales=[1, 2, 3, 6],							use_auxiliary_head=A,							auxiliary_loss_weight=0.4,							auxiliary_in_channels=40,							auxiliary_channels=256,							auxiliary_num_convs=1,							auxiliary_concat_input=A,							loss_ignore_index=255,							num_labels=self.num_labels,							)







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      int    					=						UperNetForSemanticSegmentation(config=A  )
             model.to(A  )
             model.eval()
             SCREAMING_SNAKE_CASE   :      Any    					=						model(A  )
             self.parent.assertEqual(
                 result.logits.shape,							(self.batch_size, self.num_labels, self.image_size, self.image_size)  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[Any]    					=						self.prepare_config_and_inputs()
             (
                 (
                 SCREAMING_SNAKE_CASE
             )					,    (
                 SCREAMING_SNAKE_CASE
             )					,    (
                 SCREAMING_SNAKE_CASE
             )					,    
             )   :      Dict    					=						config_and_inputs
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						{'pixel_values': pixel_values}
             return config, inputs_dict






@require_torch
class 							_a				(			SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE ,      unittest.TestCase ):





       '''simple docstring'''






       A						:				str        =						(UperNetForSemanticSegmentation,) if is_torch_available() else ()
       A						:				Dict        =						{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
       A						:				List[Any]        =						False
       A						:				Tuple        =						False
       A						:				Optional[Any]        =						False
       A						:				str        =						False
       A						:				List[str]        =						False
       A						:				int        =						False







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						UperNetModelTester(self  )
             SCREAMING_SNAKE_CASE   :      Any    					=						ConfigTester(self,							config_class=A,							has_text_modality=A,							hidden_size=37  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             return







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      str    					=						self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   SCREAMING_SNAKE_CASE   :      List[Any]    					=						model_class(A  )
                   SCREAMING_SNAKE_CASE   :      Any    					=						inspect.signature(model.forward  )
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   SCREAMING_SNAKE_CASE   :      Any    					=						[*signature.parameters.keys()]

                   SCREAMING_SNAKE_CASE   :      Any    					=						['pixel_values']
                   self.assertListEqual(arg_names[:1],							A  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      str    					=						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*A  )







       @unittest.skip(reason='UperNet does not use inputs_embeds'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @unittest.skip(reason='UperNet does not support input and output embeddings'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @unittest.skip(reason='UperNet does not have a base model'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @unittest.skip(reason='UperNet does not have a base model'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @require_torch_multi_gpu
       @unittest.skip(reason='UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             def check_hidden_states_output(A,							A,							A  ):
                   SCREAMING_SNAKE_CASE   :      str    					=						model_class(A  )
                   model.to(A  )
                   model.eval()

                   with torch.no_grad():
                         SCREAMING_SNAKE_CASE   :      Optional[int]    					=						model(**self._prepare_for_class(A,							A  )  )

                   SCREAMING_SNAKE_CASE   :      int    					=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                   SCREAMING_SNAKE_CASE   :      Optional[int]    					=						self.model_tester.num_stages
                   self.assertEqual(len(A  ),							expected_num_stages + 1  )

                   # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
                   self.assertListEqual(
                       list(hidden_states[0].shape[-2:]  ),							[self.model_tester.image_size // 4, self.model_tester.image_size // 4],							)

             SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      Dict    					=						self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   SCREAMING_SNAKE_CASE   :      List[str]    					=						True
                   check_hidden_states_output(A,							A,							A  )

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   SCREAMING_SNAKE_CASE   :      Any    					=						True

                   check_hidden_states_output(A,							A,							A  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      List[str]    					=						self.model_tester.prepare_config_and_inputs_for_common()

             SCREAMING_SNAKE_CASE   :      Any    					=						_config_zero_init(A  )
             SCREAMING_SNAKE_CASE   :      Tuple    					=						_config_zero_init(configs_no_init.backbone_config  )
             for model_class in self.all_model_classes:
                   SCREAMING_SNAKE_CASE   :      str    					=						model_class(config=A  )
                   for name, param in model.named_parameters():
                         if param.requires_grad:
                               self.assertIn(
                                   ((param.data.mean() * 1E9).round() / 1E9).item(),							[0.0, 1.0],							msg=F"Parameter {name} of model {model_class} seems not properly initialized",							)







       @unittest.skip(reason='UperNet does not have tied weights'  )
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







       @slow
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   SCREAMING_SNAKE_CASE   :      Dict    					=						UperNetForSemanticSegmentation.from_pretrained(A  )
                   self.assertIsNotNone(A  )




def   lowercase__(			):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      int    					=						hf_hub_download(
          repo_id='hf-internal-testing/fixtures_ade20k'      ,repo_type='dataset'      ,filename='ADE_val_00000001.jpg'    )
      SCREAMING_SNAKE_CASE   :      Optional[int]    					=						Image.open(__UpperCamelCase    ).convert('RGB'    )
      return image






@require_torch
@require_vision
@slow
class 							_a				(			unittest.TestCase ):





       '''simple docstring'''






       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      int    					=						AutoImageProcessor.from_pretrained('openmmlab/upernet-swin-tiny'  )
             SCREAMING_SNAKE_CASE   :      Any    					=						UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-swin-tiny'  ).to(A  )

             SCREAMING_SNAKE_CASE   :      str    					=						prepare_img()
             SCREAMING_SNAKE_CASE   :      str    					=						processor(images=A,							return_tensors='pt'  ).to(A  )

             with torch.no_grad():
                   SCREAMING_SNAKE_CASE   :      Any    					=						model(**A  )

             SCREAMING_SNAKE_CASE   :      str    					=						torch.Size((1, model.config.num_labels, 512, 512)  )
             self.assertEqual(outputs.logits.shape,							A  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.tensor(
                 [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]  ).to(A  )
             self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],							A,							atol=1E-4  )  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						AutoImageProcessor.from_pretrained('openmmlab/upernet-convnext-tiny'  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-convnext-tiny'  ).to(A  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						prepare_img()
             SCREAMING_SNAKE_CASE   :      int    					=						processor(images=A,							return_tensors='pt'  ).to(A  )

             with torch.no_grad():
                   SCREAMING_SNAKE_CASE   :      str    					=						model(**A  )

             SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size((1, model.config.num_labels, 512, 512)  )
             self.assertEqual(outputs.logits.shape,							A  )

             SCREAMING_SNAKE_CASE   :      Any    					=						torch.tensor(
                 [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]  ).to(A  )
             self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],							A,							atol=1E-4  )  )




'''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 							UpperCAmelCase_							(    A				):

 '''simple docstring'''
 lowercase_					:    torch.FloatTensor






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  Union[str, Any]			,				snake_case__		:  Tuple=3			,				snake_case__		:  Dict=3			,				snake_case__		:  Dict=("DownEncoderBlock2D",)			,				snake_case__		:  Optional[Any]=(64,)			,				snake_case__		:  List[Any]=2			,				snake_case__		:  Any=32			,				snake_case__		:  Tuple="silu"			,				snake_case__		:  Tuple=True			,				):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Tuple				       =							layers_per_block

   UpperCAmelCase__      :					Optional[int]				       =							torch.nn.Convad(
       snake_case__			,				block_out_channels[0]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

   UpperCAmelCase__      :					Tuple				       =							None
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.ModuleList([] )

   # down
   UpperCAmelCase__      :					Any				       =							block_out_channels[0]
   for i, down_block_type in enumerate(snake_case__ ):
     UpperCAmelCase__      :					List[str]				       =							output_channel
     UpperCAmelCase__      :					Dict				       =							block_out_channels[i]
     UpperCAmelCase__      :					Tuple				       =							i == len(snake_case__ ) - 1

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

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				output_scale_factor=1			,				resnet_time_scale_shift="default"			,				attention_head_dim=block_out_channels[-1]			,				resnet_groups=snake_case__			,				temb_channels=snake_case__			,				)

   # out
   UpperCAmelCase__      :					Tuple				       =							nn.GroupNorm(num_channels=block_out_channels[-1]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.SiLU()

   UpperCAmelCase__      :					Dict				       =							2 * out_channels if double_z else out_channels
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[-1]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							x
   UpperCAmelCase__      :					Dict				       =							self.conv_in(snake_case__ )

   if self.training and self.gradient_checkpointing:

     def create_custom_forward(snake_case__		:  Dict ):
       def custom_forward(*snake_case__		:  List[str] ):
         return module(*snake_case__ )

       return custom_forward

     # down
     if is_torch_version(">="			,				"1.11.0" ):
       for down_block in self.down_blocks:
         UpperCAmelCase__      :					List[str]				       =							torch.utils.checkpoint.checkpoint(
             create_custom_forward(snake_case__ )			,				snake_case__			,				use_reentrant=snake_case__ )
       # middle
       UpperCAmelCase__      :					int				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				use_reentrant=snake_case__ )
     else:
       for down_block in self.down_blocks:
         UpperCAmelCase__      :					Optional[Any]				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(snake_case__ )			,				snake_case__ )
       # middle
       UpperCAmelCase__      :					int				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block )			,				snake_case__ )

   else:
     # down
     for down_block in self.down_blocks:
       UpperCAmelCase__      :					Dict				       =							down_block(snake_case__ )

     # middle
     UpperCAmelCase__      :					Optional[Any]				       =							self.mid_block(snake_case__ )

   # post-process
   UpperCAmelCase__      :					Tuple				       =							self.conv_norm_out(snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							self.conv_act(snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							self.conv_out(snake_case__ )

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  str			,				snake_case__		:  int=3			,				snake_case__		:  str=3			,				snake_case__		:  Union[str, Any]=("UpDecoderBlock2D",)			,				snake_case__		:  Dict=(64,)			,				snake_case__		:  Optional[Any]=2			,				snake_case__		:  Dict=32			,				snake_case__		:  str="silu"			,				snake_case__		:  Any="group"			,				):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Any				       =							layers_per_block

   UpperCAmelCase__      :					Any				       =							nn.Convad(
       snake_case__			,				block_out_channels[-1]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

   UpperCAmelCase__      :					str				       =							None
   UpperCAmelCase__      :					Optional[int]				       =							nn.ModuleList([] )

   UpperCAmelCase__      :					str				       =							in_channels if norm_type == "spatial" else None

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				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=snake_case__			,				temb_channels=snake_case__			,				)

   # up
   UpperCAmelCase__      :					Tuple				       =							list(reversed(snake_case__ ) )
   UpperCAmelCase__      :					Optional[Any]				       =							reversed_block_out_channels[0]
   for i, up_block_type in enumerate(snake_case__ ):
     UpperCAmelCase__      :					Dict				       =							output_channel
     UpperCAmelCase__      :					List[Any]				       =							reversed_block_out_channels[i]

     UpperCAmelCase__      :					List[Any]				       =							i == len(snake_case__ ) - 1

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

   # out
   if norm_type == "spatial":
     UpperCAmelCase__      :					Optional[Any]				       =							SpatialNorm(block_out_channels[0]			,				snake_case__ )
   else:
     UpperCAmelCase__      :					Optional[int]				       =							nn.GroupNorm(num_channels=block_out_channels[0]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Dict				       =							nn.SiLU()
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[0]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  List[str]			,				snake_case__		:  Tuple=None ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							z
   UpperCAmelCase__      :					List[str]				       =							self.conv_in(snake_case__ )

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

     def create_custom_forward(snake_case__		:  Dict ):
       def custom_forward(*snake_case__		:  List[Any] ):
         return module(*snake_case__ )

       return custom_forward

     if is_torch_version(">="			,				"1.11.0" ):
       # middle
       UpperCAmelCase__      :					str				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				snake_case__			,				use_reentrant=snake_case__ )
       UpperCAmelCase__      :					List[Any]				       =							sample.to(snake_case__ )

       # up
       for up_block in self.up_blocks:
         UpperCAmelCase__      :					Tuple				       =							torch.utils.checkpoint.checkpoint(
             create_custom_forward(snake_case__ )			,				snake_case__			,				snake_case__			,				use_reentrant=snake_case__ )
     else:
       # middle
       UpperCAmelCase__      :					List[Any]				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				snake_case__ )
       UpperCAmelCase__      :					int				       =							sample.to(snake_case__ )

       # up
       for up_block in self.up_blocks:
         UpperCAmelCase__      :					str				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(snake_case__ )			,				snake_case__			,				snake_case__ )
   else:
     # middle
     UpperCAmelCase__      :					Union[str, Any]				       =							self.mid_block(snake_case__			,				snake_case__ )
     UpperCAmelCase__      :					Optional[Any]				       =							sample.to(snake_case__ )

     # up
     for up_block in self.up_blocks:
       UpperCAmelCase__      :					int				       =							up_block(snake_case__			,				snake_case__ )

        # post-process
   if latent_embeds is None:
     UpperCAmelCase__      :					List[Any]				       =							self.conv_norm_out(snake_case__ )
   else:
     UpperCAmelCase__      :					Any				       =							self.conv_norm_out(snake_case__			,				snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							self.conv_act(snake_case__ )
   UpperCAmelCase__      :					Optional[Any]				       =							self.conv_out(snake_case__ )

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  str			,				snake_case__		:  str			,				snake_case__		:  int			,				snake_case__		:  Tuple			,				snake_case__		:  Union[str, Any]=None			,				snake_case__		:  Optional[Any]="random"			,				snake_case__		:  Any=False			,				snake_case__		:  Any=True ):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Any				       =							n_e
   UpperCAmelCase__      :					str				       =							vq_embed_dim
   UpperCAmelCase__      :					List[Any]				       =							beta
   UpperCAmelCase__      :					List[Any]				       =							legacy

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

   UpperCAmelCase__      :					Optional[Any]				       =							remap
   if self.remap is not None:
     self.register_buffer("used"			,				torch.tensor(np.load(self.remap ) ) )
     UpperCAmelCase__      :					Optional[Any]				       =							self.used.shape[0]
     UpperCAmelCase__      :					Any				       =							unknown_index  # "random" or "extra" or integer
     if self.unknown_index == "extra":
       UpperCAmelCase__      :					Union[str, Any]				       =							self.re_embed
       UpperCAmelCase__      :					Optional[Any]				       =							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:
     UpperCAmelCase__      :					int				       =							n_e

   UpperCAmelCase__      :					Dict				       =							sane_index_shape







 def 							UpperCamelCase  (					self		:  Tuple			,				snake_case__		:  int ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							inds.shape
   assert len(snake_case__ ) > 1
   UpperCAmelCase__      :					Tuple				       =							inds.reshape(ishape[0]			,				-1 )
   UpperCAmelCase__      :					Optional[Any]				       =							self.used.to(snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							(inds[:, :, None] == used[None, None, ...]).long()
   UpperCAmelCase__      :					Optional[int]				       =							match.argmax(-1 )
   UpperCAmelCase__      :					str				       =							match.sum(2 ) < 1
   if self.unknown_index == "random":
     UpperCAmelCase__      :					Union[str, Any]				       =							torch.randint(0			,				self.re_embed			,				size=new[unknown].shape ).to(device=new.device )
   else:
     UpperCAmelCase__      :					Tuple				       =							self.unknown_index
   return new.reshape(snake_case__ )







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							inds.shape
   assert len(snake_case__ ) > 1
   UpperCAmelCase__      :					List[Any]				       =							inds.reshape(ishape[0]			,				-1 )
   UpperCAmelCase__      :					int				       =							self.used.to(snake_case__ )
   if self.re_embed > self.used.shape[0]:  # extra token
     UpperCAmelCase__      :					Any				       =							0  # simply set to zero
   UpperCAmelCase__      :					List[Any]				       =							torch.gather(used[None, :][inds.shape[0] * [0], :]			,				1			,				snake_case__ )
   return back.reshape(snake_case__ )







 def 							UpperCamelCase  (					self		:  Any			,				snake_case__		:  Any ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							z.permute(0			,				2			,				3			,				1 ).contiguous()
   UpperCAmelCase__      :					str				       =							z.view(-1			,				self.vq_embed_dim )

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

   UpperCAmelCase__      :					Any				       =							self.embedding(snake_case__ ).view(z.shape )
   UpperCAmelCase__      :					Union[str, Any]				       =							None
   UpperCAmelCase__      :					Any				       =							None

   # compute loss for embedding
   if not self.legacy:
     UpperCAmelCase__      :					Optional[int]				       =							self.beta * torch.mean((z_q.detach() - z) ** 2 ) + torch.mean((z_q - z.detach()) ** 2 )
   else:
     UpperCAmelCase__      :					Tuple				       =							torch.mean((z_q.detach() - z) ** 2 ) + self.beta * torch.mean((z_q - z.detach()) ** 2 )

   # preserve gradients
   UpperCAmelCase__      :					Any				       =							z + (z_q - z).detach()

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

   if self.remap is not None:
     UpperCAmelCase__      :					Dict				       =							min_encoding_indices.reshape(z.shape[0]			,				-1 )  # add batch axis
     UpperCAmelCase__      :					Optional[int]				       =							self.remap_to_used(snake_case__ )
     UpperCAmelCase__      :					Optional[Any]				       =							min_encoding_indices.reshape(-1			,				1 )  # flatten

   if self.sane_index_shape:
     UpperCAmelCase__      :					Union[str, Any]				       =							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 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str			,				snake_case__		:  int ):



   '''simple docstring'''




   if self.remap is not None:
     UpperCAmelCase__      :					List[Any]				       =							indices.reshape(shape[0]			,				-1 )  # add batch axis
     UpperCAmelCase__      :					Dict				       =							self.unmap_to_all(snake_case__ )
     UpperCAmelCase__      :					Dict				       =							indices.reshape(-1 )  # flatten again

   # get quantized latent vectors
   UpperCAmelCase__      :					Optional[Any]				       =							self.embedding(snake_case__ )

   if shape is not None:
     UpperCAmelCase__      :					List[str]				       =							z_q.view(snake_case__ )
     # reshape back to match original input shape
     UpperCAmelCase__      :					List[Any]				       =							z_q.permute(0			,				3			,				1			,				2 ).contiguous()

   return z_q






class 							UpperCAmelCase_							(    A				):

 '''simple docstring'''







 def __init__(					self		:  int			,				snake_case__		:  Optional[Any]			,				snake_case__		:  List[str]=False ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							parameters
   UpperCAmelCase__       ,			UpperCAmelCase__      :					Optional[int]				       =							torch.chunk(snake_case__			,				2			,				dim=1 )
   UpperCAmelCase__      :					int				       =							torch.clamp(self.logvar			,				-30.0			,				20.0 )
   UpperCAmelCase__      :					Optional[Any]				       =							deterministic
   UpperCAmelCase__      :					Dict				       =							torch.exp(0.5 * self.logvar )
   UpperCAmelCase__      :					List[str]				       =							torch.exp(self.logvar )
   if self.deterministic:
     UpperCAmelCase__      :					Dict				       =							torch.zeros_like(
         self.mean			,				device=self.parameters.device			,				dtype=self.parameters.dtype )







 def 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Optional[torch.Generator] = None ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							randn_tensor(
       self.mean.shape			,				generator=snake_case__			,				device=self.parameters.device			,				dtype=self.parameters.dtype )
   UpperCAmelCase__      :					List[str]				       =							self.mean + self.std * sample
   return x







 def 							UpperCamelCase  (					self		:  Optional[Any]			,				snake_case__		:  Any=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 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Tuple			,				snake_case__		:  Tuple=[1, 2, 3] ):



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   return self.mean



