import copy
import unittest

from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
         MODEL_FOR_QUESTION_ANSWERING_MAPPING,
         MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
         MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
         LayoutLMvaConfig,
         LayoutLMvaForQuestionAnswering,
         LayoutLMvaForSequenceClassification,
         LayoutLMvaForTokenClassification,
         LayoutLMvaModel,
     )
     from transformers.models.layoutlmva.modeling_layoutlmva import LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import LayoutLMvaImageProcessor



class     _A   :






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=2   ,				SCREAMING_SNAKE_CASE_=3   ,				SCREAMING_SNAKE_CASE_=4   ,				SCREAMING_SNAKE_CASE_=2   ,				SCREAMING_SNAKE_CASE_=7   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=99   ,				SCREAMING_SNAKE_CASE_=36   ,				SCREAMING_SNAKE_CASE_=3   ,				SCREAMING_SNAKE_CASE_=4   ,				SCREAMING_SNAKE_CASE_=37   ,				SCREAMING_SNAKE_CASE_="gelu"   ,				SCREAMING_SNAKE_CASE_=0.1   ,				SCREAMING_SNAKE_CASE_=0.1   ,				SCREAMING_SNAKE_CASE_=512   ,				SCREAMING_SNAKE_CASE_=16   ,				SCREAMING_SNAKE_CASE_=2   ,				SCREAMING_SNAKE_CASE_=0.02   ,				SCREAMING_SNAKE_CASE_=6   ,				SCREAMING_SNAKE_CASE_=6   ,				SCREAMING_SNAKE_CASE_=3   ,				SCREAMING_SNAKE_CASE_=4   ,				SCREAMING_SNAKE_CASE_=None   ,				SCREAMING_SNAKE_CASE_=1000   ,				)		->      Optional[int]:





       '''simple docstring'''





       UpperCamelCase__										=      parent
       UpperCamelCase__										=      batch_size
       UpperCamelCase__										=      num_channels
       UpperCamelCase__										=      image_size
       UpperCamelCase__										=      patch_size
       UpperCamelCase__										=      text_seq_length
       UpperCamelCase__										=      is_training
       UpperCamelCase__										=      use_input_mask
       UpperCamelCase__										=      use_token_type_ids
       UpperCamelCase__										=      use_labels
       UpperCamelCase__										=      vocab_size
       UpperCamelCase__										=      hidden_size
       UpperCamelCase__										=      num_hidden_layers
       UpperCamelCase__										=      num_attention_heads
       UpperCamelCase__										=      intermediate_size
       UpperCamelCase__										=      hidden_act
       UpperCamelCase__										=      hidden_dropout_prob
       UpperCamelCase__										=      attention_probs_dropout_prob
       UpperCamelCase__										=      max_position_embeddings
       UpperCamelCase__										=      type_vocab_size
       UpperCamelCase__										=      type_sequence_label_size
       UpperCamelCase__										=      initializer_range
       UpperCamelCase__										=      coordinate_size
       UpperCamelCase__										=      shape_size
       UpperCamelCase__										=      num_labels
       UpperCamelCase__										=      num_choices
       UpperCamelCase__										=      scope
       UpperCamelCase__										=      range_bbox

       # LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
       UpperCamelCase__										=      text_seq_length
       UpperCamelCase__										=      (image_size // patch_size) ** 2 + 1
       UpperCamelCase__										=      self.text_seq_length + self.image_seq_length






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      ids_tensor([self.batch_size, self.text_seq_length]   ,				self.vocab_size	)

       UpperCamelCase__										=      ids_tensor([self.batch_size, self.text_seq_length, 4]   ,				self.range_bbox	)
       # Ensure that bbox is legal
       for i in range(bbox.shape[0]	):
            for j in range(bbox.shape[1]	):
                 if bbox[i, j, 3] < bbox[i, j, 1]:
                      UpperCamelCase__										=      bbox[i, j, 3]
                      UpperCamelCase__										=      bbox[i, j, 1]
                      UpperCamelCase__										=      t
                 if bbox[i, j, 2] < bbox[i, j, 0]:
                      UpperCamelCase__										=      bbox[i, j, 2]
                      UpperCamelCase__										=      bbox[i, j, 0]
                      UpperCamelCase__										=      t

       UpperCamelCase__										=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

       UpperCamelCase__										=      None
       if self.use_input_mask:
            UpperCamelCase__										=      random_attention_mask([self.batch_size, self.text_seq_length]	)

       UpperCamelCase__										=      None
       if self.use_token_type_ids:
            UpperCamelCase__										=      ids_tensor([self.batch_size, self.text_seq_length]   ,				self.type_vocab_size	)

       UpperCamelCase__										=      None
       UpperCamelCase__										=      None
       if self.use_labels:
            UpperCamelCase__										=      ids_tensor([self.batch_size]   ,				self.type_sequence_label_size	)
            UpperCamelCase__										=      ids_tensor([self.batch_size, self.text_seq_length]   ,				self.num_labels	)

       UpperCamelCase__										=      LayoutLMvaConfig(
           vocab_size=self.vocab_size   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				intermediate_size=self.intermediate_size   ,				hidden_act=self.hidden_act   ,				hidden_dropout_prob=self.hidden_dropout_prob   ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,				max_position_embeddings=self.max_position_embeddings   ,				type_vocab_size=self.type_vocab_size   ,				initializer_range=self.initializer_range   ,				coordinate_size=self.coordinate_size   ,				shape_size=self.shape_size   ,				input_size=self.image_size   ,				patch_size=self.patch_size   ,				)

       return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      LayoutLMvaModel(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()

       # text + image
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      model(
           SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_   ,				token_type_ids=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				token_type_ids=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_	)

       self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)	)

       # text only
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(
           result.last_hidden_state.shape   ,				(self.batch_size, self.text_seq_length, self.hidden_size)	)

       # image only
       UpperCamelCase__										=      model(pixel_values=SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(
           result.last_hidden_state.shape   ,				(self.batch_size, self.image_seq_length, self.hidden_size)	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      self.num_labels
       UpperCamelCase__										=      LayoutLMvaForSequenceClassification(SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       UpperCamelCase__										=      model(
           SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_   ,				token_type_ids=SCREAMING_SNAKE_CASE_   ,				labels=SCREAMING_SNAKE_CASE_   ,				)
       self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      self.num_labels
       UpperCamelCase__										=      LayoutLMvaForTokenClassification(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       UpperCamelCase__										=      model(
           SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_   ,				token_type_ids=SCREAMING_SNAKE_CASE_   ,				labels=SCREAMING_SNAKE_CASE_   ,				)
       self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.text_seq_length, self.num_labels)	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      LayoutLMvaForQuestionAnswering(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       UpperCamelCase__										=      model(
           SCREAMING_SNAKE_CASE_   ,				bbox=SCREAMING_SNAKE_CASE_   ,				pixel_values=SCREAMING_SNAKE_CASE_   ,				attention_mask=SCREAMING_SNAKE_CASE_   ,				token_type_ids=SCREAMING_SNAKE_CASE_   ,				start_positions=SCREAMING_SNAKE_CASE_   ,				end_positions=SCREAMING_SNAKE_CASE_   ,				)
       self.parent.assertEqual(result.start_logits.shape   ,				(self.batch_size, self.seq_length)	)
       self.parent.assertEqual(result.end_logits.shape   ,				(self.batch_size, self.seq_length)	)






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       (
           UpperCamelCase__
       )						,       
       )										=      config_and_inputs
       UpperCamelCase__										=      {
           '''input_ids''': input_ids,
           '''bbox''': bbox,
           '''pixel_values''': pixel_values,
           '''token_type_ids''': token_type_ids,
           '''attention_mask''': input_mask,
       }
       return config, inputs_dict



@require_torch
class     _A   (    __UpperCamelCase	, __UpperCamelCase	, unittest.TestCase							):
  SCREAMING_SNAKE_CASE_	:	List[Any]						=False
  SCREAMING_SNAKE_CASE_	:	str						=False
  SCREAMING_SNAKE_CASE_	:	List[str]						=False

  SCREAMING_SNAKE_CASE_	:	Optional[Any]						=(
      (
          LayoutLMvaModel,
          LayoutLMvaForSequenceClassification,
          LayoutLMvaForTokenClassification,
          LayoutLMvaForQuestionAnswering,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE_	:	Any						=(
      {"document-question-answering": LayoutLMvaForQuestionAnswering, "feature-extraction": LayoutLMvaModel}
      if is_torch_available()
      else {}
  )






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      Tuple:





       '''simple docstring'''





       return True






  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       UpperCamelCase__										=      LayoutLMvaModelTester(self	)
       UpperCamelCase__										=      ConfigTester(self   ,				config_class=SCREAMING_SNAKE_CASE_   ,				hidden_size=37	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=False	)		->      int:





       '''simple docstring'''





       UpperCamelCase__										=      copy.deepcopy(SCREAMING_SNAKE_CASE_	)
       if model_class in get_values(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      {
                k: v.unsqueeze(1	).expand(-1   ,				self.model_tester.num_choices   ,				-1	).contiguous()
                if isinstance(SCREAMING_SNAKE_CASE_   ,				torch.Tensor	) and v.ndim > 1
                else v
                for k, v in inputs_dict.items()
            }
       if return_labels:
            if model_class in get_values(SCREAMING_SNAKE_CASE_	):
                 UpperCamelCase__										=      torch.ones(self.model_tester.batch_size   ,				dtype=torch.long   ,				device=SCREAMING_SNAKE_CASE_	)
            elif model_class in get_values(SCREAMING_SNAKE_CASE_	):
                 UpperCamelCase__										=      torch.zeros(
                     self.model_tester.batch_size   ,				dtype=torch.long   ,				device=SCREAMING_SNAKE_CASE_	)
                 UpperCamelCase__										=      torch.zeros(
                     self.model_tester.batch_size   ,				dtype=torch.long   ,				device=SCREAMING_SNAKE_CASE_	)
            elif model_class in [
                *get_values(SCREAMING_SNAKE_CASE_	),
            ]:
                 UpperCamelCase__										=      torch.zeros(
                     self.model_tester.batch_size   ,				dtype=torch.long   ,				device=SCREAMING_SNAKE_CASE_	)
            elif model_class in [
                *get_values(SCREAMING_SNAKE_CASE_	),
            ]:
                 UpperCamelCase__										=      torch.zeros(
                     (self.model_tester.batch_size, self.model_tester.text_seq_length)   ,				dtype=torch.long   ,				device=SCREAMING_SNAKE_CASE_   ,				)

       return inputs_dict






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





       '''simple docstring'''





       self.config_tester.run_common_tests()






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





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            UpperCamelCase__										=      type
            self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_	)






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





       '''simple docstring'''





       for model_name in LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCamelCase__										=      LayoutLMvaModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
            self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)







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



@require_torch
class     _A   (    unittest.TestCase							):






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





       '''simple docstring'''





       return LayoutLMvaImageProcessor(apply_ocr=SCREAMING_SNAKE_CASE_	) if is_vision_available() else None






  @slow
  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      LayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''	).to(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      self.default_image_processor
       UpperCamelCase__										=      prepare_img()
       UpperCamelCase__										=      image_processor(images=SCREAMING_SNAKE_CASE_   ,				return_tensors='''pt'''	).pixel_values.to(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      torch.tensor([[1, 2]]	)
       UpperCamelCase__										=      torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]	).unsqueeze(0	)

       # forward pass
       UpperCamelCase__										=      model(
           input_ids=input_ids.to(SCREAMING_SNAKE_CASE_	)   ,				bbox=bbox.to(SCREAMING_SNAKE_CASE_	)   ,				pixel_values=pixel_values.to(SCREAMING_SNAKE_CASE_	)   ,				)

       # verify the logits
       UpperCamelCase__										=      torch.Size((1, 199, 768)	)
       self.assertEqual(outputs.last_hidden_state.shape   ,				SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      torch.tensor(
           [[-0.0529, 0.3618, 0.1632], [-0.1587, -0.1667, -0.0400], [-0.1557, -0.1671, -0.0505]]	).to(SCREAMING_SNAKE_CASE_	)

       self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]   ,				SCREAMING_SNAKE_CASE_   ,				atol=1E-4	)	)




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   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	"DiagonalGaussianDistribution"



class     _A   (    __UpperCamelCase	, __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	List[str]						=True






  @register_to_config
  def __init__(self   ,				SCREAMING_SNAKE_CASE_ = 3   ,				SCREAMING_SNAKE_CASE_ = 3   ,				SCREAMING_SNAKE_CASE_ = ("DownEncoderBlock2D",)   ,				SCREAMING_SNAKE_CASE_ = ("UpDecoderBlock2D",)   ,				SCREAMING_SNAKE_CASE_ = (64,)   ,				SCREAMING_SNAKE_CASE_ = 1   ,				SCREAMING_SNAKE_CASE_ = "silu"   ,				SCREAMING_SNAKE_CASE_ = 4   ,				SCREAMING_SNAKE_CASE_ = 32   ,				SCREAMING_SNAKE_CASE_ = 32   ,				SCREAMING_SNAKE_CASE_ = 0.18215   ,				)		->      Optional[int]:





       '''simple docstring'''





       super().__init__()

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

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

       UpperCamelCase__										=      nn.Convad(2 * latent_channels   ,				2 * latent_channels   ,				1	)
       UpperCamelCase__										=      nn.Convad(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				1	)

       UpperCamelCase__										=      False
       UpperCamelCase__										=      False

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






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=False	)		->      Dict:





       '''simple docstring'''





       if isinstance(SCREAMING_SNAKE_CASE_   ,				(Encoder, Decoder)	):
            UpperCamelCase__										=      value






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_ = True	)		->      Union[str, Any]:





       '''simple docstring'''





       UpperCamelCase__										=      use_tiling






  def 		_a					(self	)		->      int:





       '''simple docstring'''





       self.enable_tiling(SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      True






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





       '''simple docstring'''





       UpperCamelCase__										=      False






  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def 		_a					(self	)		->      Dict[str, AttentionProcessor]:





       '''simple docstring'''





       UpperCamelCase__										=      {}

       def fn_recursive_add_processors(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
            if hasattr(SCREAMING_SNAKE_CASE_   ,				'''set_processor'''	):
                 UpperCamelCase__										=      module.processor

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

            return processors

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

       return processors






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      List[Any]:





       '''simple docstring'''





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

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

       def fn_recursive_attn_processor(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
            if hasattr(SCREAMING_SNAKE_CASE_   ,				'''set_processor'''	):
                 if not isinstance(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
                      module.set_processor(SCREAMING_SNAKE_CASE_	)
                 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}"   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)

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






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       self.set_attn_processor(AttnProcessor()	)






  @apply_forward_hook
  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = True	)		->      AutoencoderKLOutput:





       '''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(SCREAMING_SNAKE_CASE_   ,				return_dict=SCREAMING_SNAKE_CASE_	)

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

       UpperCamelCase__										=      self.quant_conv(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      DiagonalGaussianDistribution(SCREAMING_SNAKE_CASE_	)

       if not return_dict:
            return (posterior,)

       return AutoencoderKLOutput(latent_dist=SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = True	)		->      Union[DecoderOutput, torch.FloatTensor]:





       '''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(SCREAMING_SNAKE_CASE_   ,				return_dict=SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      self.post_quant_conv(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      self.decoder(SCREAMING_SNAKE_CASE_	)

       if not return_dict:
            return (dec,)

       return DecoderOutput(sample=SCREAMING_SNAKE_CASE_	)






  @apply_forward_hook
  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = True	)		->      Union[DecoderOutput, torch.FloatTensor]:





       '''simple docstring'''





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

       if not return_dict:
            return (decoded,)

       return DecoderOutput(sample=SCREAMING_SNAKE_CASE_	)






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





       '''simple docstring'''





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






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      List[Any]:





       '''simple docstring'''





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






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = True	)		->      AutoencoderKLOutput:





       '''simple docstring'''





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

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

       UpperCamelCase__										=      torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=2	)
       UpperCamelCase__										=      DiagonalGaussianDistribution(SCREAMING_SNAKE_CASE_	)

       if not return_dict:
            return (posterior,)

       return AutoencoderKLOutput(latent_dist=SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = True	)		->      Union[DecoderOutput, torch.FloatTensor]:





       '''simple docstring'''





       UpperCamelCase__										=      int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)	)
       UpperCamelCase__										=      int(self.tile_sample_min_size * self.tile_overlap_factor	)
       UpperCamelCase__										=      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.
       UpperCamelCase__										=      []
       for i in range(0   ,				z.shape[2]   ,				SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            for j in range(0   ,				z.shape[3]   ,				SCREAMING_SNAKE_CASE_	):
                 UpperCamelCase__										=      z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                 UpperCamelCase__										=      self.post_quant_conv(SCREAMING_SNAKE_CASE_	)
                 UpperCamelCase__										=      self.decoder(SCREAMING_SNAKE_CASE_	)
                 row.append(SCREAMING_SNAKE_CASE_	)
            rows.append(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      []
       for i, row in enumerate(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            for j, tile in enumerate(SCREAMING_SNAKE_CASE_	):
                 # blend the above tile and the left tile
                 # to the current tile and add the current tile to the result row
                 if i > 0:
                      UpperCamelCase__										=      self.blend_v(rows[i - 1][j]   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
                 if j > 0:
                      UpperCamelCase__										=      self.blend_h(row[j - 1]   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
                 result_row.append(tile[:, :, :row_limit, :row_limit]	)
            result_rows.append(torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=3	)	)

       UpperCamelCase__										=      torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=2	)
       if not return_dict:
            return (dec,)

       return DecoderOutput(sample=SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = False   ,				SCREAMING_SNAKE_CASE_ = True   ,				SCREAMING_SNAKE_CASE_ = None   ,				)		->      Union[DecoderOutput, torch.FloatTensor]:





       '''simple docstring'''





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

       if not return_dict:
            return (dec,)

       return DecoderOutput(sample=SCREAMING_SNAKE_CASE_	)



