'''simple docstring'''




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


			'''simple docstring'''






			def __init__(  self			,       __lowerCAmelCase			,       __lowerCAmelCase=2			,       __lowerCAmelCase=True			,       __lowerCAmelCase=False			,       __lowerCAmelCase=10			,       __lowerCAmelCase=3			,       __lowerCAmelCase=32 * 8			,       __lowerCAmelCase=32 * 8			,       __lowerCAmelCase=4			,       __lowerCAmelCase=64			,       ):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]).to(
					    __lowerCAmelCase)

					lowerCAmelCase     =    torch.ones([self.batch_size, self.min_size, self.max_size]			,       device=__lowerCAmelCase)

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

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



			def 				a_			(  self):




					"""simple docstring"""


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

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

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



			def 				a_			(  self):




					"""simple docstring"""


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



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    output.encoder_hidden_states
					lowerCAmelCase     =    output.pixel_decoder_hidden_states
					lowerCAmelCase     =    output.transformer_decoder_hidden_states

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



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase=False):




					"""simple docstring"""


					with torch.no_grad():
							lowerCAmelCase     =    MaskaFormerModel(config=__lowerCAmelCase)
							model.to(__lowerCAmelCase)
							model.eval()

							lowerCAmelCase     =    model(pixel_values=__lowerCAmelCase			,       pixel_mask=__lowerCAmelCase)
							lowerCAmelCase     =    model(__lowerCAmelCase			,       output_hidden_states=__lowerCAmelCase)

					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(__lowerCAmelCase			,       __lowerCAmelCase)







			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    MaskaFormerForUniversalSegmentation(config=__lowerCAmelCase)
					model.to(__lowerCAmelCase)
					model.eval()

					def comm_check_on_output(__lowerCAmelCase):
							# 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():
							lowerCAmelCase     =    model(pixel_values=__lowerCAmelCase			,       pixel_mask=__lowerCAmelCase)
							lowerCAmelCase     =    model(__lowerCAmelCase)

							comm_check_on_output(__lowerCAmelCase)

							lowerCAmelCase     =    model(
							    pixel_values=__lowerCAmelCase			,       pixel_mask=__lowerCAmelCase			,       mask_labels=__lowerCAmelCase			,       class_labels=__lowerCAmelCase)

					comm_check_on_output(__lowerCAmelCase)

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







@require_torch
class 							a__(   lowerCAmelCase__					,		lowerCAmelCase__					,		unittest.TestCase      ):


			'''simple docstring'''






			UpperCAmelCase_     :	str						=			(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
			UpperCAmelCase_     :	Tuple						=			{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

			UpperCAmelCase_     :	Union[str, Any]						=			False
			UpperCAmelCase_     :	List[str]						=			False
			UpperCAmelCase_     :	Dict						=			False
			UpperCAmelCase_     :	Dict						=			False



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    MaskaFormerModelTester(self)
					lowerCAmelCase     =    ConfigTester(self			,       config_class=__lowerCAmelCase			,       has_text_modality=__lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					self.config_tester.run_common_tests()



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase   ,	lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs_for_common()
					self.model_tester.create_and_check_maskaformer_model(__lowerCAmelCase			,       **__lowerCAmelCase			,       output_hidden_states=__lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__lowerCAmelCase)



			@unittest.skip(reason="""Mask2Former does not use inputs_embeds""")
			def 				a_			(  self):




					"""simple docstring"""


					pass



			@unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method""")
			def 				a_			(  self):




					"""simple docstring"""


					pass



			@unittest.skip(reason="""Mask2Former is not a generative model""")
			def 				a_			(  self):




					"""simple docstring"""


					pass



			@unittest.skip(reason="""Mask2Former does not use token embeddings""")
			def 				a_			(  self):




					"""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 				a_			(  self):




					"""simple docstring"""


					pass



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




					"""simple docstring"""


					pass



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase   ,	lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs_for_common()

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

							lowerCAmelCase     =    ["""pixel_values"""]
							self.assertListEqual(arg_names[:1]			,       __lowerCAmelCase)



			@slow
			def 				a_			(  self):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


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

					lowerCAmelCase     =    MaskaFormerForUniversalSegmentation(__lowerCAmelCase).to(__lowerCAmelCase)
					lowerCAmelCase     =    model(**__lowerCAmelCase)
					self.assertTrue(outputs.loss is not None)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase   ,	lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs_for_common()
					self.model_tester.create_and_check_maskaformer_model(__lowerCAmelCase			,       **__lowerCAmelCase			,       output_hidden_states=__lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase   ,	lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs_for_common()

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



			def 				a_			(  self):




					"""simple docstring"""


					if not self.model_tester.is_training:
							return

					lowerCAmelCase     =    self.all_model_classes[1]
					lowerCAmelCase   ,	lowerCAmelCase   ,	lowerCAmelCase   ,	lowerCAmelCase   ,	lowerCAmelCase     =    self.model_tester.prepare_config_and_inputs()

					lowerCAmelCase     =    model_class(__lowerCAmelCase)
					model.to(__lowerCAmelCase)
					model.train()

					lowerCAmelCase     =    model(__lowerCAmelCase			,       mask_labels=__lowerCAmelCase			,       class_labels=__lowerCAmelCase).loss
					loss.backward()







			def 				a_			(  self):




					"""simple docstring"""


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

					lowerCAmelCase     =    model_class(__lowerCAmelCase).to(__lowerCAmelCase)
					model.train()

					lowerCAmelCase     =    model(__lowerCAmelCase			,       mask_labels=__lowerCAmelCase			,       class_labels=__lowerCAmelCase)

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

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

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

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

					outputs.loss.backward(retain_graph=__lowerCAmelCase)

					self.assertIsNotNone(encoder_hidden_states.grad)
					self.assertIsNotNone(pixel_decoder_hidden_states.grad)
					self.assertIsNotNone(transformer_decoder_hidden_states.grad)
					self.assertIsNotNone(attentions.grad)







__lowercase        =    1e-4






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

		'''simple docstring'''


		lowerCAmelCase     =    Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
		return image






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


			'''simple docstring'''






			@cached_property
			def 				a_			(  self):




					"""simple docstring"""


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



			@cached_property
			def 				a_			(  self):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    MaskaFormerModel.from_pretrained(self.model_checkpoints).to(__lowerCAmelCase)
					lowerCAmelCase     =    self.default_image_processor
					lowerCAmelCase     =    prepare_img()
					lowerCAmelCase     =    image_processor(__lowerCAmelCase			,       return_tensors="""pt""").to(__lowerCAmelCase)
					lowerCAmelCase     =    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(__lowerCAmelCase			,       (1, 3, 384, 384))

					with torch.no_grad():
							lowerCAmelCase     =    model(**__lowerCAmelCase)

					lowerCAmelCase     =    torch.tensor(
					    [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]).to(__lowerCAmelCase)
					self.assertTrue(
					    torch.allclose(
					        outputs.encoder_last_hidden_state[0, 0, :3, :3]			,       __lowerCAmelCase			,       atol=__lowerCAmelCase))

					lowerCAmelCase     =    torch.tensor(
					    [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]).to(__lowerCAmelCase)
					self.assertTrue(
					    torch.allclose(
					        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,       __lowerCAmelCase			,       atol=__lowerCAmelCase))

					lowerCAmelCase     =    torch.tensor(
					    [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]).to(__lowerCAmelCase)
					self.assertTrue(
					    torch.allclose(
					        outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,       __lowerCAmelCase			,       atol=__lowerCAmelCase))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints).to(__lowerCAmelCase).eval()
					lowerCAmelCase     =    self.default_image_processor
					lowerCAmelCase     =    prepare_img()
					lowerCAmelCase     =    image_processor(__lowerCAmelCase			,       return_tensors="""pt""").to(__lowerCAmelCase)
					lowerCAmelCase     =    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(__lowerCAmelCase			,       (1, 3, 384, 384))

					with torch.no_grad():
							lowerCAmelCase     =    model(**__lowerCAmelCase)
					# masks_queries_logits
					lowerCAmelCase     =    outputs.masks_queries_logits
					self.assertEqual(
					    masks_queries_logits.shape			,       (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4))
					lowerCAmelCase     =    [
					    [-8.7839, -9.0056, -8.8121],
					    [-7.4104, -7.0313, -6.5401],
					    [-6.6105, -6.3427, -6.4675],
					]
					lowerCAmelCase     =    torch.tensor(__lowerCAmelCase).to(__lowerCAmelCase)
					self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,       __lowerCAmelCase			,       atol=__lowerCAmelCase))
					# class_queries_logits
					lowerCAmelCase     =    outputs.class_queries_logits
					self.assertEqual(class_queries_logits.shape			,       (1, model.config.num_queries, model.config.num_labels + 1))
					lowerCAmelCase     =    torch.tensor(
					    [
					        [1.8324, -8.0835, -4.1922],
					        [0.8450, -9.0050, -3.6053],
					        [0.3045, -7.7293, -3.0275],
					    ]).to(__lowerCAmelCase)
					self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,       __lowerCAmelCase			,       atol=__lowerCAmelCase))







			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints).to(__lowerCAmelCase).eval()
					lowerCAmelCase     =    self.default_image_processor

					lowerCAmelCase     =    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"""			,       )

					lowerCAmelCase     =    inputs["""pixel_values"""].to(__lowerCAmelCase)
					lowerCAmelCase     =    [el.to(__lowerCAmelCase) for el in inputs["""mask_labels"""]]
					lowerCAmelCase     =    [el.to(__lowerCAmelCase) for el in inputs["""class_labels"""]]

					with torch.no_grad():
							lowerCAmelCase     =    model(**__lowerCAmelCase)

					self.assertTrue(outputs.loss is not None)





'''simple docstring'''




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

import torch
import torch.nn as nn

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






@dataclass
class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	"DiagonalGaussianDistribution"







class 							a__(   lowerCAmelCase__					,		lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	List[Any]						=			True



			@register_to_config
			def __init__(  self			,       __lowerCAmelCase = 3			,       __lowerCAmelCase = 3			,       __lowerCAmelCase = ("DownEncoderBlock2D",)			,       __lowerCAmelCase = ("UpDecoderBlock2D",)			,       __lowerCAmelCase = (64,)			,       __lowerCAmelCase = 1			,       __lowerCAmelCase = "silu"			,       __lowerCAmelCase = 4			,       __lowerCAmelCase = 32			,       __lowerCAmelCase = 32			,       __lowerCAmelCase = 0.18215			,       ):




					"""simple docstring"""


					super().__init__()

					# pass init params to Encoder
					lowerCAmelCase     =    Encoder(
					    in_channels=__lowerCAmelCase			,       out_channels=__lowerCAmelCase			,       down_block_types=__lowerCAmelCase			,       block_out_channels=__lowerCAmelCase			,       layers_per_block=__lowerCAmelCase			,       act_fn=__lowerCAmelCase			,       norm_num_groups=__lowerCAmelCase			,       double_z=__lowerCAmelCase			,       )

					# pass init params to Decoder
					lowerCAmelCase     =    Decoder(
					    in_channels=__lowerCAmelCase			,       out_channels=__lowerCAmelCase			,       up_block_types=__lowerCAmelCase			,       block_out_channels=__lowerCAmelCase			,       layers_per_block=__lowerCAmelCase			,       norm_num_groups=__lowerCAmelCase			,       act_fn=__lowerCAmelCase			,       )

					lowerCAmelCase     =    nn.Convad(2 * latent_channels			,       2 * latent_channels			,       1)
					lowerCAmelCase     =    nn.Convad(__lowerCAmelCase			,       __lowerCAmelCase			,       1)

					lowerCAmelCase     =    False
					lowerCAmelCase     =    False

					# only relevant if vae tiling is enabled
					lowerCAmelCase     =    self.config.sample_size
					lowerCAmelCase     =    (
					    self.config.sample_size[0]
					    if isinstance(self.config.sample_size			,       (list, tuple))
					    else self.config.sample_size
					)
					lowerCAmelCase     =    int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
					lowerCAmelCase     =    0.25



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase=False):




					"""simple docstring"""


					if isinstance(__lowerCAmelCase			,       (Encoder, Decoder)):
							lowerCAmelCase     =    value



			def 				a_			(  self			,       __lowerCAmelCase = True):




					"""simple docstring"""


					lowerCAmelCase     =    use_tiling



			def 				a_			(  self):




					"""simple docstring"""


					self.enable_tiling(__lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    True



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    False



			@property
			# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    {}

					def fn_recursive_add_processors(__lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase):
							if hasattr(__lowerCAmelCase			,       """set_processor"""):
									lowerCAmelCase     =    module.processor

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

							return processors

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

					return processors



			def 				a_			(  self			,       __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    len(self.attn_processors.keys())

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

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

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

					for name, module in self.named_children():
							fn_recursive_attn_processor(__lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					self.set_attn_processor(AttnProcessor())



			@apply_forward_hook
			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = True):




					"""simple docstring"""


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

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

					lowerCAmelCase     =    self.quant_conv(__lowerCAmelCase)
					lowerCAmelCase     =    DiagonalGaussianDistribution(__lowerCAmelCase)

					if not return_dict:
							return (posterior,)

					return AutoencoderKLOutput(latent_dist=__lowerCAmelCase)



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = True):




					"""simple docstring"""


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

					lowerCAmelCase     =    self.post_quant_conv(__lowerCAmelCase)
					lowerCAmelCase     =    self.decoder(__lowerCAmelCase)

					if not return_dict:
							return (dec,)

					return DecoderOutput(sample=__lowerCAmelCase)



			@apply_forward_hook
			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = True):




					"""simple docstring"""


					if self.use_slicing and z.shape[0] > 1:
							lowerCAmelCase     =    [self._decode(__lowerCAmelCase).sample for z_slice in z.split(1)]
							lowerCAmelCase     =    torch.cat(__lowerCAmelCase)
					else:
							lowerCAmelCase     =    self._decode(__lowerCAmelCase).sample

					if not return_dict:
							return (decoded,)

					return DecoderOutput(sample=__lowerCAmelCase)



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    min(a.shape[2]			,       b.shape[2]			,       __lowerCAmelCase)
					for y in range(__lowerCAmelCase):
							lowerCAmelCase     =    a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
					return b



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    min(a.shape[3]			,       b.shape[3]			,       __lowerCAmelCase)
					for x in range(__lowerCAmelCase):
							lowerCAmelCase     =    a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
					return b



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = True):




					"""simple docstring"""


					lowerCAmelCase     =    int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
					lowerCAmelCase     =    int(self.tile_latent_min_size * self.tile_overlap_factor)
					lowerCAmelCase     =    self.tile_latent_min_size - blend_extent

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

					lowerCAmelCase     =    torch.cat(__lowerCAmelCase			,       dim=2)
					lowerCAmelCase     =    DiagonalGaussianDistribution(__lowerCAmelCase)

					if not return_dict:
							return (posterior,)

					return AutoencoderKLOutput(latent_dist=__lowerCAmelCase)



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = True):




					"""simple docstring"""


					lowerCAmelCase     =    int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
					lowerCAmelCase     =    int(self.tile_sample_min_size * self.tile_overlap_factor)
					lowerCAmelCase     =    self.tile_sample_min_size - blend_extent

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

					lowerCAmelCase     =    torch.cat(__lowerCAmelCase			,       dim=2)
					if not return_dict:
							return (dec,)

					return DecoderOutput(sample=__lowerCAmelCase)







			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = False			,       __lowerCAmelCase = True			,       __lowerCAmelCase = None			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    sample
					lowerCAmelCase     =    self.encode(__lowerCAmelCase).latent_dist
					if sample_posterior:
							lowerCAmelCase     =    posterior.sample(generator=__lowerCAmelCase)
					else:
							lowerCAmelCase     =    posterior.mode()
					lowerCAmelCase     =    self.decode(__lowerCAmelCase).sample

					if not return_dict:
							return (dec,)

					return DecoderOutput(sample=__lowerCAmelCase)




