


import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

							if is_vision_available():
														from transformers import MaskaFormerImageProcessor

if is_vision_available():
							from PIL import Image

class     A__				:




						def __init__(							self       :   List[Any]					,     _a       :   int					,     _a       :   Optional[Any]=2					,     _a       :   Any=True					,     _a       :   Tuple=False					,     _a       :   Tuple=10					,     _a       :   Optional[int]=3					,     _a       :   Union[str, Any]=32 * 8					,     _a       :   Optional[int]=32 * 8					,     _a       :   Union[str, Any]=4					,     _a       :   Dict=64					,     )				->  Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=parent
										_SCREAMING_SNAKE_CASE      	=batch_size
										_SCREAMING_SNAKE_CASE      	=is_training
										_SCREAMING_SNAKE_CASE      	=use_auxiliary_loss
										_SCREAMING_SNAKE_CASE      	=num_queries
										_SCREAMING_SNAKE_CASE      	=num_channels
										_SCREAMING_SNAKE_CASE      	=min_size
										_SCREAMING_SNAKE_CASE      	=max_size
										_SCREAMING_SNAKE_CASE      	=num_labels
										_SCREAMING_SNAKE_CASE      	=hidden_dim
										_SCREAMING_SNAKE_CASE      	=hidden_dim




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]	).to(
										    snake_case__	)

										_SCREAMING_SNAKE_CASE      	=torch.ones([self.batch_size, self.min_size, self.max_size]					,     device=snake_case__	)

										_SCREAMING_SNAKE_CASE      	=(
										    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]					,     device=snake_case__	) > 0.5
										).float()
										_SCREAMING_SNAKE_CASE      	=(torch.rand((self.batch_size, self.num_labels)					,     device=snake_case__	) > 0.5).long()

										_SCREAMING_SNAKE_CASE      	=self.get_config()
										return config, pixel_values, pixel_mask, mask_labels, class_labels




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerConfig(
										    hidden_size=self.hidden_dim					,     )
										_SCREAMING_SNAKE_CASE      	=self.num_queries
										_SCREAMING_SNAKE_CASE      	=self.num_labels

										_SCREAMING_SNAKE_CASE      	=[1, 1, 1, 1]
										_SCREAMING_SNAKE_CASE      	=self.num_channels

										_SCREAMING_SNAKE_CASE      	=64
										_SCREAMING_SNAKE_CASE      	=128
										_SCREAMING_SNAKE_CASE      	=self.hidden_dim
										_SCREAMING_SNAKE_CASE      	=self.hidden_dim
										_SCREAMING_SNAKE_CASE      	=self.hidden_dim
										return config




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.prepare_config_and_inputs()
										_SCREAMING_SNAKE_CASE      	={"pixel_values": pixel_values, "pixel_mask": pixel_mask}
										return config, inputs_dict




						def 				__UpperCamelCase					(							self       :   int					,     _a       :   Union[str, Any]					,     _a       :   Dict	)				->  str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=output.encoder_hidden_states
										_SCREAMING_SNAKE_CASE      	=output.pixel_decoder_hidden_states
										_SCREAMING_SNAKE_CASE      	=output.transformer_decoder_hidden_states

										self.parent.assertTrue(len(snake_case__	)					,     len(config.backbone_config.depths	)	)
										self.parent.assertTrue(len(snake_case__	)					,     len(config.backbone_config.depths	)	)
										self.parent.assertTrue(len(snake_case__	)					,     config.decoder_layers	)




						def 				__UpperCamelCase					(							self       :   Optional[int]					,     _a       :   List[str]					,     _a       :   List[str]					,     _a       :   str					,     _a       :   List[Any]=False	)				->  List[str]:


										"""simple docstring"""





										with torch.no_grad():
														_SCREAMING_SNAKE_CASE      	=MaskaFormerModel(config=snake_case__	)
														model.to(snake_case__	)
														model.eval()

														_SCREAMING_SNAKE_CASE      	=model(pixel_values=snake_case__					,     pixel_mask=snake_case__	)
														_SCREAMING_SNAKE_CASE      	=model(snake_case__					,     output_hidden_states=snake_case__	)

										self.parent.assertEqual(
										    output.transformer_decoder_last_hidden_state.shape					,     (self.batch_size, self.num_queries, self.hidden_dim)					,     )
										# let's ensure the other two hidden state exists
										self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None	)
										self.parent.assertTrue(output.encoder_last_hidden_state is not None	)

										if output_hidden_states:
														self.check_output_hidden_state(snake_case__					,     snake_case__	)




						def 				__UpperCamelCase					(							self       :   Tuple					,     _a       :   List[Any]					,     _a       :   Any					,     _a       :   Union[str, Any]					,     _a       :   Any					,     _a       :   Optional[int]	)				->  Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerForUniversalSegmentation(config=snake_case__	)
										model.to(snake_case__	)
										model.eval()

										def comm_check_on_output(_a       :   Union[str, Any]	):
														# let's still check that all the required stuff is there
														self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None	)
														self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None	)
														self.parent.assertTrue(result.encoder_last_hidden_state is not None	)
														# okay, now we need to check the logits shape
														# due to the encoder compression, masks have a //4 spatial size
														self.parent.assertEqual(
														    result.masks_queries_logits.shape					,     (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)					,     )
														# + 1 for null class
														self.parent.assertEqual(
														    result.class_queries_logits.shape					,     (self.batch_size, self.num_queries, self.num_labels + 1)	)

										with torch.no_grad():
														_SCREAMING_SNAKE_CASE      	=model(pixel_values=snake_case__					,     pixel_mask=snake_case__	)
														_SCREAMING_SNAKE_CASE      	=model(snake_case__	)

														comm_check_on_output(snake_case__	)

														_SCREAMING_SNAKE_CASE      	=model(
														    pixel_values=snake_case__					,     pixel_mask=snake_case__					,     mask_labels=snake_case__					,     class_labels=snake_case__	)

										comm_check_on_output(snake_case__	)

										self.parent.assertTrue(result.loss is not None	)
										self.parent.assertEqual(result.loss.shape					,     torch.Size([1]	)	)

@require_torch
class     A__				(		UpperCamelCase__ ,  UpperCamelCase__ ,  unittest.TestCase	):
						UpperCAmelCase   							=    (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
						UpperCAmelCase   							=    {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

						UpperCAmelCase   							=    False
						UpperCAmelCase   							=    False
						UpperCAmelCase   							=    False
						UpperCAmelCase   							=    False




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerModelTester(self	)
										_SCREAMING_SNAKE_CASE      	=ConfigTester(self					,     config_class=snake_case__					,     has_text_modality=snake_case__	)




						def 				__UpperCamelCase					(							self       :   Dict	)				->  Optional[int]:


										"""simple docstring"""





										self.config_tester.run_common_tests()




						def 				__UpperCamelCase					(							self       :   List[str]	)				->  List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs_for_common()
										self.model_tester.create_and_check_maskaformer_model(snake_case__					,     **snake_case__					,     output_hidden_states=snake_case__	)




						def 				__UpperCamelCase					(							self       :   Dict	)				->  List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*snake_case__	)




						@unittest.skip(reason='''Mask2Former does not use inputs_embeds'''	)
						def 				__UpperCamelCase					(							self       :   Dict	)				->  Tuple:


										"""simple docstring"""





										pass




						@unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''	)
						def 				__UpperCamelCase					(							self       :   Optional[int]	)				->  str:


										"""simple docstring"""





										pass




						@unittest.skip(reason='''Mask2Former is not a generative model'''	)
						def 				__UpperCamelCase					(							self       :   str	)				->  str:


										"""simple docstring"""





										pass




						@unittest.skip(reason='''Mask2Former does not use token embeddings'''	)
						def 				__UpperCamelCase					(							self       :   Optional[int]	)				->  Tuple:


										"""simple docstring"""





										pass




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


										"""simple docstring"""





										pass




						@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''	)
						def 				__UpperCamelCase					(							self       :   List[str]	)				->  Any:


										"""simple docstring"""





										pass




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs_for_common()

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

														_SCREAMING_SNAKE_CASE      	=["pixel_values"]
														self.assertListEqual(arg_names[:1]					,     snake_case__	)




						@slow
						def 				__UpperCamelCase					(							self       :   Dict	)				->  int:


										"""simple docstring"""





										for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
														_SCREAMING_SNAKE_CASE      	=MaskaFormerModel.from_pretrained(snake_case__	)
														self.assertIsNotNone(snake_case__	)




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=(self.model_tester.min_size,) * 2
										_SCREAMING_SNAKE_CASE      	={
										    "pixel_values": torch.randn((2, 3, *size)					,     device=snake_case__	),
										    "mask_labels": torch.randn((2, 10, *size)					,     device=snake_case__	),
										    "class_labels": torch.zeros(2					,     10					,     device=snake_case__	).long(),
										}
										_SCREAMING_SNAKE_CASE      	=self.model_tester.get_config()

										_SCREAMING_SNAKE_CASE      	=MaskaFormerForUniversalSegmentation(snake_case__	).to(snake_case__	)
										_SCREAMING_SNAKE_CASE      	=model(**snake_case__	)
										self.assertTrue(outputs.loss is not None	)




						def 				__UpperCamelCase					(							self       :   Optional[int]	)				->  Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs_for_common()
										self.model_tester.create_and_check_maskaformer_model(snake_case__					,     **snake_case__					,     output_hidden_states=snake_case__	)




						def 				__UpperCamelCase					(							self       :   Dict	)				->  List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_SCREAMING_SNAKE_CASE      	=model_class(snake_case__	).to(snake_case__	)
														_SCREAMING_SNAKE_CASE      	=model(**snake_case__					,     output_attentions=snake_case__	)
														self.assertTrue(outputs.attentions is not None	)




						def 				__UpperCamelCase					(							self       :   List[str]	)				->  Dict:


										"""simple docstring"""





										if not self.model_tester.is_training:
														return

										_SCREAMING_SNAKE_CASE      	=self.all_model_classes[1]
										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs()

										_SCREAMING_SNAKE_CASE      	=model_class(snake_case__	)
										model.to(snake_case__	)
										model.train()

										_SCREAMING_SNAKE_CASE      	=model(snake_case__					,     mask_labels=snake_case__					,     class_labels=snake_case__	).loss
										loss.backward()




						def 				__UpperCamelCase					(							self       :   List[str]	)				->  int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.all_model_classes[1]
										_SCREAMING_SNAKE_CASE      	=self.model_tester.prepare_config_and_inputs()
										_SCREAMING_SNAKE_CASE      	=True
										_SCREAMING_SNAKE_CASE      	=True

										_SCREAMING_SNAKE_CASE      	=model_class(snake_case__	).to(snake_case__	)
										model.train()

										_SCREAMING_SNAKE_CASE      	=model(snake_case__					,     mask_labels=snake_case__					,     class_labels=snake_case__	)

										_SCREAMING_SNAKE_CASE      	=outputs.encoder_hidden_states[0]
										encoder_hidden_states.retain_grad()

										_SCREAMING_SNAKE_CASE      	=outputs.pixel_decoder_hidden_states[0]
										pixel_decoder_hidden_states.retain_grad()

										_SCREAMING_SNAKE_CASE      	=outputs.transformer_decoder_hidden_states[0]
										transformer_decoder_hidden_states.retain_grad()

										_SCREAMING_SNAKE_CASE      	=outputs.attentions[0]
										attentions.retain_grad()

										outputs.loss.backward(retain_graph=snake_case__	)

										self.assertIsNotNone(encoder_hidden_states.grad	)
										self.assertIsNotNone(pixel_decoder_hidden_states.grad	)
										self.assertIsNotNone(transformer_decoder_hidden_states.grad	)
										self.assertIsNotNone(attentions.grad	)
snake_case_    :      str				    =				1e-4







def     lowerCamelCase( ):
				_SCREAMING_SNAKE_CASE      	=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
				return image


@require_vision
@slow
class     A__				(		unittest.TestCase	):




						@cached_property
						def 				__UpperCamelCase					(							self       :   Tuple	)				->  Optional[int]:


										"""simple docstring"""





										return "facebook/mask2former-swin-small-coco-instance"




						@cached_property
						def 				__UpperCamelCase					(							self       :   List[Any]	)				->  Optional[Any]:


										"""simple docstring"""





										return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints	) if is_vision_available() else None




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerModel.from_pretrained(self.model_checkpoints	).to(snake_case__	)
										_SCREAMING_SNAKE_CASE      	=self.default_image_processor
										_SCREAMING_SNAKE_CASE      	=prepare_img()
										_SCREAMING_SNAKE_CASE      	=image_processor(snake_case__					,     return_tensors='''pt'''	).to(snake_case__	)
										_SCREAMING_SNAKE_CASE      	=inputs["pixel_values"].shape
										# check size is divisible by 32
										self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
										# check size
										self.assertEqual(snake_case__					,     (1, 3, 384, 384)	)

										with torch.no_grad():
														_SCREAMING_SNAKE_CASE      	=model(**snake_case__	)

										_SCREAMING_SNAKE_CASE      	=torch.tensor(
										    [[-0.27_90, -1.07_17, -1.16_68], [-0.51_28, -0.31_28, -0.49_87], [-0.58_32, 0.19_71, -0.01_97]]	).to(snake_case__	)
										self.assertTrue(
										    torch.allclose(
										        outputs.encoder_last_hidden_state[0, 0, :3, :3]					,     snake_case__					,     atol=snake_case__	)	)

										_SCREAMING_SNAKE_CASE      	=torch.tensor(
										    [[0.89_73, 1.18_47, 1.17_76], [1.19_34, 1.50_40, 1.51_28], [1.11_53, 1.44_86, 1.49_51]]	).to(snake_case__	)
										self.assertTrue(
										    torch.allclose(
										        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]					,     snake_case__					,     atol=snake_case__	)	)

										_SCREAMING_SNAKE_CASE      	=torch.tensor(
										    [[2.11_52, 1.70_00, -0.86_03], [1.58_08, 1.80_04, -0.93_53], [1.60_43, 1.74_95, -0.59_99]]	).to(snake_case__	)
										self.assertTrue(
										    torch.allclose(
										        outputs.transformer_decoder_last_hidden_state[0, :3, :3]					,     snake_case__					,     atol=snake_case__	)	)




						def 				__UpperCamelCase					(							self       :   Tuple	)				->  List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(snake_case__	).eval()
										_SCREAMING_SNAKE_CASE      	=self.default_image_processor
										_SCREAMING_SNAKE_CASE      	=prepare_img()
										_SCREAMING_SNAKE_CASE      	=image_processor(snake_case__					,     return_tensors='''pt'''	).to(snake_case__	)
										_SCREAMING_SNAKE_CASE      	=inputs["pixel_values"].shape
										# check size is divisible by 32
										self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
										# check size
										self.assertEqual(snake_case__					,     (1, 3, 384, 384)	)

										with torch.no_grad():
														_SCREAMING_SNAKE_CASE      	=model(**snake_case__	)
										# masks_queries_logits
										_SCREAMING_SNAKE_CASE      	=outputs.masks_queries_logits
										self.assertEqual(
										    masks_queries_logits.shape					,     (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)	)
										_SCREAMING_SNAKE_CASE      	=[
										    [-8.78_39, -9.00_56, -8.81_21],
										    [-7.41_04, -7.03_13, -6.54_01],
										    [-6.61_05, -6.34_27, -6.46_75],
										]
										_SCREAMING_SNAKE_CASE      	=torch.tensor(snake_case__	).to(snake_case__	)
										self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]					,     snake_case__					,     atol=snake_case__	)	)
										# class_queries_logits
										_SCREAMING_SNAKE_CASE      	=outputs.class_queries_logits
										self.assertEqual(class_queries_logits.shape					,     (1, model.config.num_queries, model.config.num_labels + 1)	)
										_SCREAMING_SNAKE_CASE      	=torch.tensor(
										    [
										        [1.83_24, -8.08_35, -4.19_22],
										        [0.84_50, -9.00_50, -3.60_53],
										        [0.30_45, -7.72_93, -3.02_75],
										    ]	).to(snake_case__	)
										self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]					,     snake_case__					,     atol=snake_case__	)	)




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(snake_case__	).eval()
										_SCREAMING_SNAKE_CASE      	=self.default_image_processor

										_SCREAMING_SNAKE_CASE      	=image_processor(
										    [np.zeros((3, 800, 1333)	), np.zeros((3, 800, 1333)	)]					,     segmentation_maps=[np.zeros((384, 384)	).astype(np.floataa	), np.zeros((384, 384)	).astype(np.floataa	)]					,     return_tensors='''pt'''					,     )

										_SCREAMING_SNAKE_CASE      	=inputs["pixel_values"].to(snake_case__	)
										_SCREAMING_SNAKE_CASE      	=[el.to(snake_case__	) for el in inputs["mask_labels"]]
										_SCREAMING_SNAKE_CASE      	=[el.to(snake_case__	) for el in inputs["class_labels"]]

										with torch.no_grad():
														_SCREAMING_SNAKE_CASE      	=model(**snake_case__	)

										self.assertTrue(outputs.loss is not None	)

'''simple docstring'''





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

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer






@dataclass
class 							UpperCAmelCase_							(    A				):

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






class 							UpperCAmelCase_							(    A					,				A				):

 '''simple docstring'''







 @register_to_config
 def __init__(					self		:  List[str]			,				snake_case__		:  int = 3			,				snake_case__		:  int = 3			,				snake_case__		:  Tuple[str] = ("DownEncoderBlock2D",)			,				snake_case__		:  Tuple[str] = ("UpDecoderBlock2D",)			,				snake_case__		:  Tuple[int] = (64,)			,				snake_case__		:  int = 1			,				snake_case__		:  str = "silu"			,				snake_case__		:  int = 3			,				snake_case__		:  int = 32			,				snake_case__		:  int = 2_56			,				snake_case__		:  int = 32			,				snake_case__		:  Optional[int] = None			,				snake_case__		:  float = 0.18215			,				snake_case__		:  str = "group"			,				):



   '''simple docstring'''




   super().__init__()

   # pass init params to Encoder
   UpperCAmelCase__      :					Tuple				       =							Encoder(
       in_channels=snake_case__			,				out_channels=snake_case__			,				down_block_types=snake_case__			,				block_out_channels=snake_case__			,				layers_per_block=snake_case__			,				act_fn=snake_case__			,				norm_num_groups=snake_case__			,				double_z=snake_case__			,				)

   UpperCAmelCase__      :					str				       =							vq_embed_dim if vq_embed_dim is not None else latent_channels

   UpperCAmelCase__      :					Tuple				       =							nn.Convad(snake_case__			,				snake_case__			,				1 )
   UpperCAmelCase__      :					int				       =							VectorQuantizer(snake_case__			,				snake_case__			,				beta=0.25			,				remap=snake_case__			,				sane_index_shape=snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							nn.Convad(snake_case__			,				snake_case__			,				1 )

   # pass init params to Decoder
   UpperCAmelCase__      :					Union[str, Any]				       =							Decoder(
       in_channels=snake_case__			,				out_channels=snake_case__			,				up_block_types=snake_case__			,				block_out_channels=snake_case__			,				layers_per_block=snake_case__			,				act_fn=snake_case__			,				norm_num_groups=snake_case__			,				norm_type=snake_case__			,				)







 @apply_forward_hook
 def 							UpperCamelCase  (					self		:  Union[str, Any]			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							self.encoder(snake_case__ )
   UpperCAmelCase__      :					List[str]				       =							self.quant_conv(snake_case__ )

   if not return_dict:
     return (h,)

   return VQEncoderOutput(latents=snake_case__ )







 @apply_forward_hook
 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = False			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   if not force_not_quantize:
     UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Any				       =							self.quantize(snake_case__ )
   else:
     UpperCAmelCase__      :					Tuple				       =							h
   UpperCAmelCase__      :					int				       =							self.post_quant_conv(snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							self.decoder(snake_case__			,				quant if self.config.norm_type == "spatial" else None )

   if not return_dict:
     return (dec,)

   return DecoderOutput(sample=snake_case__ )







 def 							UpperCamelCase  (					self		:  Any			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							sample
   UpperCAmelCase__      :					int				       =							self.encode(snake_case__ ).latents
   UpperCAmelCase__      :					Any				       =							self.decode(snake_case__ ).sample

   if not return_dict:
     return (dec,)

   return DecoderOutput(sample=snake_case__ )



