



'''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     :





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

					'''simple docstring'''
					__UpperCAmelCase      :					List[Any]   =   parent
					__UpperCAmelCase      :					Any   =   batch_size
					__UpperCAmelCase      :					Dict   =   image_size
					__UpperCAmelCase      :					Optional[Any]   =   num_channels
					__UpperCAmelCase      :					Dict   =   num_stages
					__UpperCAmelCase      :					str   =   hidden_sizes
					__UpperCAmelCase      :					Optional[Any]   =   depths
					__UpperCAmelCase      :					int   =   is_training
					__UpperCAmelCase      :					str   =   use_labels
					__UpperCAmelCase      :					str   =   intermediate_size
					__UpperCAmelCase      :					str   =   hidden_act
					__UpperCAmelCase      :					Dict   =   type_sequence_label_size
					__UpperCAmelCase      :					Optional[int]   =   initializer_range
					__UpperCAmelCase      :					Optional[int]   =   out_features
					__UpperCAmelCase      :					Tuple   =   num_labels
					__UpperCAmelCase      :					List[str]   =   scope
					__UpperCAmelCase      :					Optional[int]   =   num_stages





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

					'''simple docstring'''
					__UpperCAmelCase      :					str   =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

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

					return config, pixel_values, labels





		def 							__A				(   self )				->   str:

					'''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 							__A				(   self )				->   Tuple:

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





		def 							__A				(   self       ,				__UpperCAmelCase       ,				__UpperCAmelCase       ,				__UpperCAmelCase )				->   Dict:

					'''simple docstring'''
					__UpperCAmelCase      :					Any   =   UperNetForSemanticSegmentation(config=__UpperCAmelCase )
					model.to(__UpperCAmelCase )
					model.eval()
					__UpperCAmelCase      :					Optional[int]   =   model(__UpperCAmelCase )
					self.parent.assertEqual(
					    result.logits.shape       ,				(self.batch_size, self.num_labels, self.image_size, self.image_size) )





		def 							__A				(   self )				->   str:

					'''simple docstring'''
					__UpperCAmelCase      :					Dict   =   self.prepare_config_and_inputs()
					(
					    (
					    __UpperCAmelCase     
					) ,     (
					    __UpperCAmelCase     
					) ,     (
					    __UpperCAmelCase     
					) ,          
					)      :					List[str]   =   config_and_inputs
					__UpperCAmelCase      :					Optional[int]   =   {"""pixel_values""": pixel_values}
					return config, inputs_dict



@require_torch
class _A     (			__A		, __A		, unittest.TestCase			):
		_SCREAMING_SNAKE_CASE    :			int  					=	(UperNetForSemanticSegmentation,) if is_torch_available() else ()
		_SCREAMING_SNAKE_CASE    :			int  					=	{'image-segmentation': UperNetForSemanticSegmentation} if is_torch_available() else {}
		_SCREAMING_SNAKE_CASE    :			Optional[int]  					=	False
		_SCREAMING_SNAKE_CASE    :			List[Any]  					=	False
		_SCREAMING_SNAKE_CASE    :			Tuple  					=	False
		_SCREAMING_SNAKE_CASE    :			Optional[Any]  					=	False
		_SCREAMING_SNAKE_CASE    :			Union[str, Any]  					=	False
		_SCREAMING_SNAKE_CASE    :			Dict  					=	False





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

					'''simple docstring'''
					__UpperCAmelCase      :					str   =   UperNetModelTester(self )
					__UpperCAmelCase      :					List[Any]   =   ConfigTester(self       ,				config_class=__UpperCAmelCase       ,				has_text_modality=__UpperCAmelCase       ,				hidden_size=37 )





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

					'''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 							__A				(   self )				->   int:

					'''simple docstring'''
					return





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

					'''simple docstring'''
					__UpperCAmelCase ,     __UpperCAmelCase      :					Optional[Any]   =   self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
								__UpperCAmelCase      :					Optional[int]   =   model_class(__UpperCAmelCase )
								__UpperCAmelCase      :					Union[str, Any]   =   inspect.signature(model.forward )
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								__UpperCAmelCase      :					Any   =   [*signature.parameters.keys()]

								__UpperCAmelCase      :					List[Any]   =   ["""pixel_values"""]
								self.assertListEqual(arg_names[:1]       ,				__UpperCAmelCase )





		def 							__A				(   self )				->   Tuple:

					'''simple docstring'''
					__UpperCAmelCase      :					int   =   self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCAmelCase )





		@unittest.skip(reason="""UperNet does not use inputs_embeds""" )
		def 							__A				(   self )				->   Optional[int]:

					'''simple docstring'''
					pass





		@unittest.skip(reason="""UperNet does not support input and output embeddings""" )
		def 							__A				(   self )				->   Tuple:

					'''simple docstring'''
					pass





		@unittest.skip(reason="""UperNet does not have a base model""" )
		def 							__A				(   self )				->   List[Any]:

					'''simple docstring'''
					pass





		@unittest.skip(reason="""UperNet does not have a base model""" )
		def 							__A				(   self )				->   Dict:

					'''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 							__A				(   self )				->   Any:

					'''simple docstring'''
					pass





		@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""" )
		def 							__A				(   self )				->   Dict:

					'''simple docstring'''
					pass





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

					'''simple docstring'''
					def check_hidden_states_output(__UpperCAmelCase       ,				__UpperCAmelCase       ,				__UpperCAmelCase ):
								__UpperCAmelCase      :					Optional[int]   =   model_class(__UpperCAmelCase )
								model.to(__UpperCAmelCase )
								model.eval()

								with torch.no_grad():
											__UpperCAmelCase      :					str   =   model(**self._prepare_for_class(__UpperCAmelCase       ,				__UpperCAmelCase ) )

								__UpperCAmelCase      :					Any   =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

								__UpperCAmelCase      :					int   =   self.model_tester.num_stages
								self.assertEqual(len(__UpperCAmelCase )       ,				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]       ,				)

					__UpperCAmelCase ,     __UpperCAmelCase      :					Optional[Any]   =   self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
								__UpperCAmelCase      :					int   =   True
								check_hidden_states_output(__UpperCAmelCase       ,				__UpperCAmelCase       ,				__UpperCAmelCase )

								# check that output_hidden_states also work using config
								del inputs_dict["output_hidden_states"]
								__UpperCAmelCase      :					Optional[int]   =   True

								check_hidden_states_output(__UpperCAmelCase       ,				__UpperCAmelCase       ,				__UpperCAmelCase )





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

					'''simple docstring'''
					__UpperCAmelCase ,     __UpperCAmelCase      :					int   =   self.model_tester.prepare_config_and_inputs_for_common()

					__UpperCAmelCase      :					Tuple   =   _config_zero_init(__UpperCAmelCase )
					__UpperCAmelCase      :					Dict   =   _config_zero_init(configs_no_init.backbone_config )
					for model_class in self.all_model_classes:
								__UpperCAmelCase      :					Tuple   =   model_class(config=__UpperCAmelCase )
								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 							__A				(   self )				->   str:

					'''simple docstring'''
					pass





		@slow
		def 							__A				(   self )				->   List[str]:

					'''simple docstring'''
					for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__UpperCAmelCase      :					Dict   =   UperNetForSemanticSegmentation.from_pretrained(__UpperCAmelCase )
								self.assertIsNotNone(__UpperCAmelCase )


def      lowercase_		(      ):

			"""simple docstring"""





			__UpperCAmelCase      :					Tuple   =   hf_hub_download(
			    repo_id="""hf-internal-testing/fixtures_ade20k"""	,							repo_type="""dataset"""	,							filename="""ADE_val_00000001.jpg"""       )
			__UpperCAmelCase      :					Tuple   =   Image.open(UpperCamelCase__       ).convert("""RGB"""       )
			return image






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





		def 							__A				(   self )				->   Dict:

					'''simple docstring'''
					__UpperCAmelCase      :					Union[str, Any]   =   AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny""" )
					__UpperCAmelCase      :					List[str]   =   UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny""" ).to(__UpperCAmelCase )

					__UpperCAmelCase      :					int   =   prepare_img()
					__UpperCAmelCase      :					Tuple   =   processor(images=__UpperCAmelCase       ,				return_tensors="""pt""" ).to(__UpperCAmelCase )

					with torch.no_grad():
								__UpperCAmelCase      :					Any   =   model(**__UpperCAmelCase )

					__UpperCAmelCase      :					Optional[Any]   =   torch.Size((1, model.config.num_labels, 512, 512) )
					self.assertEqual(outputs.logits.shape       ,				__UpperCAmelCase )

					__UpperCAmelCase      :					Optional[int]   =   torch.tensor(
					    [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]] ).to(__UpperCAmelCase )
					self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,				__UpperCAmelCase       ,				atol=1E-4 ) )





		def 							__A				(   self )				->   str:

					'''simple docstring'''
					__UpperCAmelCase      :					Optional[int]   =   AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny""" )
					__UpperCAmelCase      :					List[str]   =   UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny""" ).to(__UpperCAmelCase )

					__UpperCAmelCase      :					List[str]   =   prepare_img()
					__UpperCAmelCase      :					Any   =   processor(images=__UpperCAmelCase       ,				return_tensors="""pt""" ).to(__UpperCAmelCase )

					with torch.no_grad():
								__UpperCAmelCase      :					Optional[int]   =   model(**__UpperCAmelCase )

					__UpperCAmelCase      :					int   =   torch.Size((1, model.config.num_labels, 512, 512) )
					self.assertEqual(outputs.logits.shape       ,				__UpperCAmelCase )

					__UpperCAmelCase      :					Optional[Any]   =   torch.tensor(
					    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]] ).to(__UpperCAmelCase )
					self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,				__UpperCAmelCase       ,				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       ):
						lowerCAmelCase__     :  torch.FloatTensor






class 			UpperCamelCase__(	nn.Module       ):





						def __init__(				self						,__UpperCAmelCase=3						,__UpperCAmelCase=3						,__UpperCAmelCase=("DownEncoderBlock2D",)						,__UpperCAmelCase=(64,)						,__UpperCAmelCase=2						,__UpperCAmelCase=32						,__UpperCAmelCase="silu"						,__UpperCAmelCase=True						,)    -> Union[str, Any]:
													super().__init__()
													A__		   =							layers_per_block

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

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

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

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

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

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

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

													A__		   =							False





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> Optional[int]:
													A__		   =							x
													A__		   =							self.conv_in(__UpperCAmelCase   )

													if self.training and self.gradient_checkpointing:

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

																											return custom_forward

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

													else:
																				# down
																				for down_block in self.down_blocks:
																											A__		   =							down_block(__UpperCAmelCase   )

																				# middle
																				A__		   =							self.mid_block(__UpperCAmelCase   )

													# post-process
													A__		   =							self.conv_norm_out(__UpperCAmelCase   )
													A__		   =							self.conv_act(__UpperCAmelCase   )
													A__		   =							self.conv_out(__UpperCAmelCase   )

													return sample






class 			UpperCamelCase__(	nn.Module       ):





						def __init__(				self						,__UpperCAmelCase=3						,__UpperCAmelCase=3						,__UpperCAmelCase=("UpDecoderBlock2D",)						,__UpperCAmelCase=(64,)						,__UpperCAmelCase=2						,__UpperCAmelCase=32						,__UpperCAmelCase="silu"						,__UpperCAmelCase="group"						,)    -> Any:
													super().__init__()
													A__		   =							layers_per_block

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

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

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

													# mid
													A__		   =							UNetMidBlockaD(
													    in_channels=block_out_channels[-1]						,resnet_eps=1e-6						,resnet_act_fn=__UpperCAmelCase						,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=__UpperCAmelCase						,temb_channels=__UpperCAmelCase						,)

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

																				A__		   =							i == len(__UpperCAmelCase   ) - 1

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

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

													A__		   =							False





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase=None   )    -> Dict:
													A__		   =							z
													A__		   =							self.conv_in(__UpperCAmelCase   )

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

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

																											return custom_forward

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

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

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

																				# up
																				for up_block in self.up_blocks:
																											A__		   =							up_block(__UpperCAmelCase						,__UpperCAmelCase   )

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

													return sample






class 			UpperCamelCase__(	nn.Module       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase=None						,__UpperCAmelCase="random"						,__UpperCAmelCase=False						,__UpperCAmelCase=True   )    -> Tuple:
													super().__init__()
													A__		   =							n_e
													A__		   =							vq_embed_dim
													A__		   =							beta
													A__		   =							legacy

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

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

													A__		   =							sane_index_shape





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





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





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> Dict:
													# reshape z -> (batch, height, width, channel) and flatten
													A__		   =							z.permute(0						,2						,3						,1   ).contiguous()
													A__		   =							z.view(-1						,self.vq_embed_dim   )

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

													A__		   =							self.embedding(__UpperCAmelCase   ).view(z.shape   )
													A__		   =							None
													A__		   =							None

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

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

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

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

													if self.sane_index_shape:
																				A__		   =							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  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> str:
													# shape specifying (batch, height, width, channel)
													if self.remap is not None:
																				A__		   =							indices.reshape(shape[0]						,-1   )  # add batch axis
																				A__		   =							self.unmap_to_all(__UpperCAmelCase   )
																				A__		   =							indices.reshape(-1   )  # flatten again

													# get quantized latent vectors
													A__		   =							self.embedding(__UpperCAmelCase   )

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

													return z_q






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase=False   )    -> Any:
													A__		   =							parameters
													A__   ,     A__		   =							torch.chunk(__UpperCAmelCase						,2						,dim=1   )
													A__		   =							torch.clamp(self.logvar						,-3_0.0						,2_0.0   )
													A__		   =							deterministic
													A__		   =							torch.exp(0.5 * self.logvar   )
													A__		   =							torch.exp(self.logvar   )
													if self.deterministic:
																				A__		   =							A__		   =							torch.zeros_like(
																				    self.mean						,device=self.parameters.device						,dtype=self.parameters.dtype   )





						def  snake_case__		(				self						,__UpperCAmelCase = None   )    -> torch.FloatTensor:
													# make sure sample is on the same device as the parameters and has same dtype
													A__		   =							randn_tensor(
													    self.mean.shape						,generator=__UpperCAmelCase						,device=self.parameters.device						,dtype=self.parameters.dtype   )
													A__		   =							self.mean + self.std * sample
													return x





						def  snake_case__		(				self						,__UpperCAmelCase=None   )    -> int:
													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  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase=[1, 2, 3]   )    -> List[Any]:
													if self.deterministic:
																				return torch.Tensor([0.0]   )
													A__		   =							np.log(2.0 * np.pi   )
													return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean						,2   ) / self.var						,dim=__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Optional[Any]:
													return self.mean





