



from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import numpy as np
    import tensorflow as tf

    from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
    from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import MobileViTImageProcessor



class        lowerCAmelCase_	(						UpperCamelCase_							):
       def       snake_case_ (    self )  ->	List[Any]:
          UpperCamelCase					:						List[str]					=			self.config_class(**self.inputs_dict )
          self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE,    'hidden_sizes' ) )
          self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE,    'neck_hidden_sizes' ) )
          self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE,    'num_attention_heads' ) )



class        lowerCAmelCase_	:
       def __init__(    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_=13,    SCREAMING_SNAKE_CASE_=32,    SCREAMING_SNAKE_CASE_=2,    SCREAMING_SNAKE_CASE_=3,    SCREAMING_SNAKE_CASE_=640,    SCREAMING_SNAKE_CASE_=4,    SCREAMING_SNAKE_CASE_="silu",    SCREAMING_SNAKE_CASE_=3,    SCREAMING_SNAKE_CASE_=32,    SCREAMING_SNAKE_CASE_=0.1,    SCREAMING_SNAKE_CASE_=0.1,    SCREAMING_SNAKE_CASE_=0.1,    SCREAMING_SNAKE_CASE_=0.02,    SCREAMING_SNAKE_CASE_=True,    SCREAMING_SNAKE_CASE_=True,    SCREAMING_SNAKE_CASE_=10,    SCREAMING_SNAKE_CASE_=None,    )  ->	Optional[int]:
          UpperCamelCase					:						Any					=			parent
          UpperCamelCase					:						int					=			batch_size
          UpperCamelCase					:						Tuple					=			image_size
          UpperCamelCase					:						Optional[int]					=			patch_size
          UpperCamelCase					:						Any					=			num_channels
          UpperCamelCase					:						Tuple					=			last_hidden_size
          UpperCamelCase					:						Union[str, Any]					=			num_attention_heads
          UpperCamelCase					:						Any					=			hidden_act
          UpperCamelCase					:						int					=			conv_kernel_size
          UpperCamelCase					:						str					=			output_stride
          UpperCamelCase					:						str					=			hidden_dropout_prob
          UpperCamelCase					:						str					=			attention_probs_dropout_prob
          UpperCamelCase					:						Optional[Any]					=			classifier_dropout_prob
          UpperCamelCase					:						List[str]					=			use_labels
          UpperCamelCase					:						int					=			is_training
          UpperCamelCase					:						Dict					=			num_labels
          UpperCamelCase					:						Optional[Any]					=			initializer_range
          UpperCamelCase					:						Optional[int]					=			scope
       def       snake_case_ (    self )  ->	List[Any]:
          UpperCamelCase					:						str					=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

          UpperCamelCase					:						int					=			None
          UpperCamelCase					:						Any					=			None
          if self.use_labels:
             UpperCamelCase					:						List[Any]					=			ids_tensor([self.batch_size],    self.num_labels )
             UpperCamelCase					:						Dict					=			ids_tensor([self.batch_size, self.image_size, self.image_size],    self.num_labels )

          UpperCamelCase					:						Any					=			self.get_config()

          return config, pixel_values, labels, pixel_labels
       def       snake_case_ (    self )  ->	Dict:
          return MobileViTConfig(
              image_size=self.image_size,    patch_size=self.patch_size,    num_channels=self.num_channels,    num_attention_heads=self.num_attention_heads,    hidden_act=self.hidden_act,    conv_kernel_size=self.conv_kernel_size,    output_stride=self.output_stride,    hidden_dropout_prob=self.hidden_dropout_prob,    attention_probs_dropout_prob=self.attention_probs_dropout_prob,    classifier_dropout_prob=self.classifier_dropout_prob,    initializer_range=self.initializer_range,    )
       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )  ->	Optional[Any]:
          UpperCamelCase					:						Optional[int]					=			TFMobileViTModel(config=__SCREAMING_SNAKE_CASE )
          UpperCamelCase					:						Optional[int]					=			model(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE )
          UpperCamelCase					:						Union[str, Any]					=			self.image_size // self.output_stride
          self.parent.assertEqual(
              result.last_hidden_state.shape,    (self.batch_size, self.last_hidden_size, expected_height, expected_width) )
       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )  ->	Optional[int]:
          UpperCamelCase					:						List[str]					=			self.num_labels
          UpperCamelCase					:						int					=			TFMobileViTForImageClassification(__SCREAMING_SNAKE_CASE )
          UpperCamelCase					:						str					=			model(__SCREAMING_SNAKE_CASE,    labels=__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE )
          self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.num_labels) )
       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )  ->	Union[str, Any]:
          UpperCamelCase					:						Optional[Any]					=			self.num_labels
          UpperCamelCase					:						Union[str, Any]					=			TFMobileViTForSemanticSegmentation(__SCREAMING_SNAKE_CASE )
          UpperCamelCase					:						str					=			self.image_size // self.output_stride

          UpperCamelCase					:						List[str]					=			model(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE )
          self.parent.assertEqual(
              result.logits.shape,    (self.batch_size, self.num_labels, expected_height, expected_width) )

          UpperCamelCase					:						Union[str, Any]					=			model(__SCREAMING_SNAKE_CASE,    labels=__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE )
          self.parent.assertEqual(
              result.logits.shape,    (self.batch_size, self.num_labels, expected_height, expected_width) )






       def       snake_case_ (    self )  ->	List[str]:
          UpperCamelCase					:						int					=			self.prepare_config_and_inputs()
          UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase					:						str					=			config_and_inputs
          UpperCamelCase					:						int					=			{'pixel_values': pixel_values}
          return config, inputs_dict



@require_tf
class        lowerCAmelCase_	(						UpperCamelCase_ ,					UpperCamelCase_ ,					unittest.TestCase							):
       UpperCAmelCase__       :      str       =						(
           (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
           if is_tf_available()
           else ()
       )
       UpperCAmelCase__       :      Dict       =						(
           {"""feature-extraction""": TFMobileViTModel, """image-classification""": TFMobileViTForImageClassification}
           if is_tf_available()
           else {}
       )

       UpperCAmelCase__       :      Optional[int]       =						False
       UpperCAmelCase__       :      Optional[Any]       =						False
       UpperCAmelCase__       :      str       =						False
       UpperCAmelCase__       :      Optional[int]       =						False
       UpperCAmelCase__       :      Union[str, Any]       =						False
       def       snake_case_ (    self )  ->	Optional[int]:
          UpperCamelCase					:						Any					=			TFMobileViTModelTester(self )
          UpperCamelCase					:						Tuple					=			TFMobileViTConfigTester(self,    config_class=__SCREAMING_SNAKE_CASE,    has_text_modality=__SCREAMING_SNAKE_CASE )
       def       snake_case_ (    self )  ->	List[str]:
          self.config_tester.run_common_tests()
       @unittest.skip(reason='MobileViT does not use inputs_embeds' )
       def       snake_case_ (    self )  ->	int:
          pass
       @unittest.skip(reason='MobileViT does not support input and output embeddings' )
       def       snake_case_ (    self )  ->	Any:
          pass
       @unittest.skip(reason='MobileViT does not output attentions' )
       def       snake_case_ (    self )  ->	Optional[int]:
          pass
       def       snake_case_ (    self )  ->	int:
          UpperCamelCase			,  UpperCamelCase					:						Optional[Any]					=			self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase					:						str					=			model_class(__SCREAMING_SNAKE_CASE )
             UpperCamelCase					:						Any					=			inspect.signature(model.call )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase					:						Optional[Any]					=			[*signature.parameters.keys()]

             UpperCamelCase					:						Optional[int]					=			['pixel_values']
             self.assertListEqual(arg_names[:1],    __SCREAMING_SNAKE_CASE )
       def       snake_case_ (    self )  ->	Optional[int]:
          UpperCamelCase					:						Union[str, Any]					=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE )
       def       snake_case_ (    self )  ->	int:
          def check_hidden_states_output(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ ):
             UpperCamelCase					:						Optional[int]					=			model_class(__SCREAMING_SNAKE_CASE )

             UpperCamelCase					:						List[Any]					=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE ) )

             UpperCamelCase					:						Any					=			outputs.hidden_states

             UpperCamelCase					:						List[Any]					=			5
             self.assertEqual(len(__SCREAMING_SNAKE_CASE ),    __SCREAMING_SNAKE_CASE )

             # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
             # with the width and height being successively divided by 2.
             UpperCamelCase					:						List[str]					=			2
             for i in range(len(__SCREAMING_SNAKE_CASE ) ):
                self.assertListEqual(
                    list(hidden_states[i].shape[-2:] ),    [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor],    )
                divisor *= 2

             self.assertEqual(self.model_tester.output_stride,    divisor // 2 )

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

          for model_class in self.all_model_classes:
             UpperCamelCase					:						str					=			True
             check_hidden_states_output(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE )

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

             check_hidden_states_output(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE )
       def       snake_case_ (    self )  ->	int:
          UpperCamelCase					:						List[str]					=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE )
       def       snake_case_ (    self )  ->	Dict:
          UpperCamelCase					:						int					=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_semantic_segmentation(*__SCREAMING_SNAKE_CASE )
       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices('GPU' ) ) == 0,    reason='TF does not support backprop for grouped convolutions on CPU.',    )
       def       snake_case_ (    self )  ->	Dict:
          super().test_dataset_conversion()
       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_=2e-1,    SCREAMING_SNAKE_CASE_=2e-1 )  ->	List[str]:
          self.assertTrue(np.allclose(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE,    atol=__SCREAMING_SNAKE_CASE,    rtol=__SCREAMING_SNAKE_CASE ) )
       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices('GPU' ) ) == 0,    reason='TF does not support backprop for grouped convolutions on CPU.',    )
       @slow
       def       snake_case_ (    self )  ->	Union[str, Any]:
          UpperCamelCase			,  UpperCamelCase					:						Dict					=			self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             # Since `TFMobileViTModel` cannot operate with the default `fit()` method.
             if model_class.__name__ != "TFMobileViTModel":
                UpperCamelCase					:						Union[str, Any]					=			model_class(__SCREAMING_SNAKE_CASE )
                if getattr(__SCREAMING_SNAKE_CASE,    'hf_compute_loss',    __SCREAMING_SNAKE_CASE ):
                   super().test_keras_fit()
       def       snake_case_ (    self )  ->	Dict:
          UpperCamelCase			,  UpperCamelCase					:						Optional[Any]					=			self.model_tester.prepare_config_and_inputs_for_common()
          for model_class in self.all_model_classes:
             # set an ignore index to correctly test the masked loss used in
             # `TFMobileViTForSemanticSegmentation`.
             if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
                UpperCamelCase					:						Optional[Any]					=			5

             UpperCamelCase					:						Union[str, Any]					=			model_class(__SCREAMING_SNAKE_CASE )
             if getattr(__SCREAMING_SNAKE_CASE,    'hf_compute_loss',    __SCREAMING_SNAKE_CASE ):
                # The number of elements in the loss should be the same as the number of elements in the label
                UpperCamelCase					:						List[Any]					=			self._prepare_for_class(inputs_dict.copy(),    __SCREAMING_SNAKE_CASE,    return_labels=__SCREAMING_SNAKE_CASE )
                UpperCamelCase					:						Union[str, Any]					=			prepared_for_class[
                    sorted(prepared_for_class.keys() - inputs_dict.keys(),    reverse=__SCREAMING_SNAKE_CASE )[0]
                ]
                UpperCamelCase					:						List[str]					=			added_label.shape.as_list()[:1]

                # Test that model correctly compute the loss with kwargs
                UpperCamelCase					:						Optional[int]					=			self._prepare_for_class(inputs_dict.copy(),    __SCREAMING_SNAKE_CASE,    return_labels=__SCREAMING_SNAKE_CASE )
                UpperCamelCase					:						Tuple					=			{'input_ids', 'pixel_values', 'input_features'}
                UpperCamelCase					:						List[Any]					=			possible_input_names.intersection(set(__SCREAMING_SNAKE_CASE ) ).pop()
                UpperCamelCase					:						str					=			prepared_for_class.pop(__SCREAMING_SNAKE_CASE )

                UpperCamelCase					:						Any					=			model(__SCREAMING_SNAKE_CASE,    **__SCREAMING_SNAKE_CASE )[0]
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1] )

                # Test that model correctly compute the loss when we mask some positions
                UpperCamelCase					:						Tuple					=			self._prepare_for_class(inputs_dict.copy(),    __SCREAMING_SNAKE_CASE,    return_labels=__SCREAMING_SNAKE_CASE )
                UpperCamelCase					:						Union[str, Any]					=			{'input_ids', 'pixel_values', 'input_features'}
                UpperCamelCase					:						List[Any]					=			possible_input_names.intersection(set(__SCREAMING_SNAKE_CASE ) ).pop()
                UpperCamelCase					:						Tuple					=			prepared_for_class.pop(__SCREAMING_SNAKE_CASE )
                if "labels" in prepared_for_class:
                   UpperCamelCase					:						List[str]					=			prepared_for_class['labels'].numpy()
                   if len(labels.shape ) > 1 and labels.shape[1] != 1:
                      # labels[0] = -100
                      UpperCamelCase					:						Tuple					=			tf.convert_to_tensor(__SCREAMING_SNAKE_CASE )
                      UpperCamelCase					:						str					=			model(__SCREAMING_SNAKE_CASE,    **__SCREAMING_SNAKE_CASE )[0]
                      self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1] )
                      self.assertTrue(not np.any(np.isnan(loss.numpy() ) ) )

                # Test that model correctly compute the loss with a dict
                UpperCamelCase					:						List[Any]					=			self._prepare_for_class(inputs_dict.copy(),    __SCREAMING_SNAKE_CASE,    return_labels=__SCREAMING_SNAKE_CASE )
                UpperCamelCase					:						Union[str, Any]					=			model(__SCREAMING_SNAKE_CASE )[0]
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1] )

                # Test that model correctly compute the loss with a tuple
                UpperCamelCase					:						Any					=			self._prepare_for_class(inputs_dict.copy(),    __SCREAMING_SNAKE_CASE,    return_labels=__SCREAMING_SNAKE_CASE )

                # Get keys that were added with the _prepare_for_class function
                UpperCamelCase					:						Optional[int]					=			prepared_for_class.keys() - inputs_dict.keys()
                UpperCamelCase					:						int					=			inspect.signature(model.call ).parameters
                UpperCamelCase					:						Optional[Any]					=			list(signature.keys() )

                # Create a dictionary holding the location of the tensors in the tuple
                UpperCamelCase					:						Tuple					=			{0: input_name}
                for label_key in label_keys:
                   UpperCamelCase					:						Optional[int]					=			signature_names.index(__SCREAMING_SNAKE_CASE )
                   UpperCamelCase					:						Optional[Any]					=			label_key
                UpperCamelCase					:						int					=			sorted(tuple_index_mapping.items() )
                # Initialize a list with their default values, update the values and convert to a tuple
                UpperCamelCase					:						Optional[int]					=			[]

                for name in signature_names:
                   if name != "kwargs":
                      list_input.append(signature[name].default )

                for index, value in sorted_tuple_index_mapping:
                   UpperCamelCase					:						Optional[Any]					=			prepared_for_class[value]

                UpperCamelCase					:						Dict					=			tuple(__SCREAMING_SNAKE_CASE )

                # Send to model
                UpperCamelCase					:						List[str]					=			model(tuple_input[:-1] )[0]

                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1] )






       @slow
       def       snake_case_ (    self )  ->	List[Any]:
          for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCamelCase					:						List[str]					=			TFMobileViTModel.from_pretrained(__SCREAMING_SNAKE_CASE )
             self.assertIsNotNone(__SCREAMING_SNAKE_CASE )



def 	UpperCamelCase			(     )			->  List[str]:
   UpperCamelCase					:						Optional[Any]					=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'			)
   return image



@require_tf
class        lowerCAmelCase_	(						unittest.TestCase							):
       @slow
       def       snake_case_ (    self )  ->	Dict:
          UpperCamelCase					:						Union[str, Any]					=			TFMobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small' )

          UpperCamelCase					:						Optional[Any]					=			MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small' )
          UpperCamelCase					:						Dict					=			prepare_img()
          UpperCamelCase					:						List[Any]					=			image_processor(images=__SCREAMING_SNAKE_CASE,    return_tensors='tf' )

          # forward pass
          UpperCamelCase					:						Union[str, Any]					=			model(**__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE )

          # verify the logits
          UpperCamelCase					:						Optional[Any]					=			tf.TensorShape((1, 1000) )
          self.assertEqual(outputs.logits.shape,    __SCREAMING_SNAKE_CASE )

          UpperCamelCase					:						Tuple					=			tf.constant([-1.93_64, -1.23_27, -0.46_53] )

          tf.debugging.assert_near(outputs.logits[0, :3],    __SCREAMING_SNAKE_CASE,    atol=1e-4,    rtol=1e-04 )






       @slow
       def       snake_case_ (    self )  ->	str:
          # `from_pt` will be removed
          UpperCamelCase					:						Tuple					=			TFMobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small' )

          UpperCamelCase					:						str					=			MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small' )

          UpperCamelCase					:						Optional[Any]					=			prepare_img()
          UpperCamelCase					:						int					=			image_processor(images=__SCREAMING_SNAKE_CASE,    return_tensors='tf' )

          # forward pass
          UpperCamelCase					:						int					=			model(inputs.pixel_values,    training=__SCREAMING_SNAKE_CASE )
          UpperCamelCase					:						Any					=			outputs.logits

          # verify the logits
          UpperCamelCase					:						List[str]					=			tf.TensorShape((1, 21, 32, 32) )
          self.assertEqual(logits.shape,    __SCREAMING_SNAKE_CASE )

          UpperCamelCase					:						int					=			tf.constant(
              [
                  [[6.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
                  [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.98_68, -9.71_32], [-11.0405, -11.0221, -10.7318]],
                  [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
              ] )

          tf.debugging.assert_near(logits[0, :3, :3, :3],    __SCREAMING_SNAKE_CASE,    rtol=1e-4,    atol=1e-4 )








from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin



@dataclass
class        lowerCAmelCase_	(						a__							):
       UpperCAmelCase__       :      torch.FloatTensor






class        lowerCAmelCase_	(						a__ ,					a__							):
       @register_to_config
       def __init__(    self,    SCREAMING_SNAKE_CASE_ = 16,    SCREAMING_SNAKE_CASE_ = 88,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = 1,    SCREAMING_SNAKE_CASE_ = 0.0,    SCREAMING_SNAKE_CASE_ = 32,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = False,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = "geglu",    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = False,    SCREAMING_SNAKE_CASE_ = False,    SCREAMING_SNAKE_CASE_ = False,    SCREAMING_SNAKE_CASE_ = "layer_norm",    SCREAMING_SNAKE_CASE_ = True,    )  ->	Union[str, Any]:
          super().__init__()
          UpperCamelCase					:						Tuple					=			use_linear_projection
          UpperCamelCase					:						int					=			num_attention_heads
          UpperCamelCase					:						Dict					=			attention_head_dim
          UpperCamelCase					:						Optional[int]					=			num_attention_heads * attention_head_dim

          # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
          # Define whether input is continuous or discrete depending on configuration
          UpperCamelCase					:						Optional[Any]					=			(in_channels is not None) and (patch_size is None)
          UpperCamelCase					:						Tuple					=			num_vector_embeds is not None
          UpperCamelCase					:						Optional[int]					=			in_channels is not None and patch_size is not None

          if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
             UpperCamelCase					:						Any					=			(
                 F"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
                 ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
                 ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
                 ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
                 ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
             )
             deprecate('norm_type!=num_embeds_ada_norm',    '1.0.0',    SCREAMING_SNAKE_CASE_,    standard_warn=SCREAMING_SNAKE_CASE_ )
             UpperCamelCase					:						Any					=			'ada_norm'

          if self.is_input_continuous and self.is_input_vectorized:
             raise ValueError(
                 F"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
                 ' sure that either `in_channels` or `num_vector_embeds` is None.' )
          elif self.is_input_vectorized and self.is_input_patches:
             raise ValueError(
                 F"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
                 ' sure that either `num_vector_embeds` or `num_patches` is None.' )
          elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
             raise ValueError(
                 F"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
                 F""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.""" )

          # 2. Define input layers
          if self.is_input_continuous:
             UpperCamelCase					:						Optional[int]					=			in_channels

             UpperCamelCase					:						str					=			torch.nn.GroupNorm(num_groups=SCREAMING_SNAKE_CASE_,    num_channels=SCREAMING_SNAKE_CASE_,    eps=1e-6,    affine=SCREAMING_SNAKE_CASE_ )
             if use_linear_projection:
                UpperCamelCase					:						Dict					=			nn.Linear(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )
             else:
                UpperCamelCase					:						Any					=			nn.Convad(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    kernel_size=1,    stride=1,    padding=0 )
          elif self.is_input_vectorized:
             assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
             assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

             UpperCamelCase					:						Union[str, Any]					=			sample_size
             UpperCamelCase					:						int					=			sample_size
             UpperCamelCase					:						Optional[Any]					=			num_vector_embeds
             UpperCamelCase					:						str					=			self.height * self.width

             UpperCamelCase					:						Union[str, Any]					=			ImagePositionalEmbeddings(
                 num_embed=SCREAMING_SNAKE_CASE_,    embed_dim=SCREAMING_SNAKE_CASE_,    height=self.height,    width=self.width )
          elif self.is_input_patches:
             assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

             UpperCamelCase					:						int					=			sample_size
             UpperCamelCase					:						Optional[Any]					=			sample_size

             UpperCamelCase					:						List[str]					=			patch_size
             UpperCamelCase					:						Tuple					=			PatchEmbed(
                 height=SCREAMING_SNAKE_CASE_,    width=SCREAMING_SNAKE_CASE_,    patch_size=SCREAMING_SNAKE_CASE_,    in_channels=SCREAMING_SNAKE_CASE_,    embed_dim=SCREAMING_SNAKE_CASE_,    )

          # 3. Define transformers blocks
          UpperCamelCase					:						Any					=			nn.ModuleList(
              [
                  BasicTransformerBlock(
                      SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    dropout=SCREAMING_SNAKE_CASE_,    cross_attention_dim=SCREAMING_SNAKE_CASE_,    activation_fn=SCREAMING_SNAKE_CASE_,    num_embeds_ada_norm=SCREAMING_SNAKE_CASE_,    attention_bias=SCREAMING_SNAKE_CASE_,    only_cross_attention=SCREAMING_SNAKE_CASE_,    upcast_attention=SCREAMING_SNAKE_CASE_,    norm_type=SCREAMING_SNAKE_CASE_,    norm_elementwise_affine=SCREAMING_SNAKE_CASE_,    )
                  for d in range(SCREAMING_SNAKE_CASE_ )
              ] )

          # 4. Define output layers
          UpperCamelCase					:						Optional[int]					=			in_channels if out_channels is None else out_channels
          if self.is_input_continuous:
             # TODO: should use out_channels for continuous projections
             if use_linear_projection:
                UpperCamelCase					:						Optional[Any]					=			nn.Linear(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )
             else:
                UpperCamelCase					:						List[str]					=			nn.Convad(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    kernel_size=1,    stride=1,    padding=0 )
          elif self.is_input_vectorized:
             UpperCamelCase					:						Optional[Any]					=			nn.LayerNorm(SCREAMING_SNAKE_CASE_ )
             UpperCamelCase					:						Any					=			nn.Linear(SCREAMING_SNAKE_CASE_,    self.num_vector_embeds - 1 )
          elif self.is_input_patches:
             UpperCamelCase					:						List[str]					=			nn.LayerNorm(SCREAMING_SNAKE_CASE_,    elementwise_affine=SCREAMING_SNAKE_CASE_,    eps=1e-6 )
             UpperCamelCase					:						Tuple					=			nn.Linear(SCREAMING_SNAKE_CASE_,    2 * inner_dim )
             UpperCamelCase					:						Any					=			nn.Linear(SCREAMING_SNAKE_CASE_,    patch_size * patch_size * self.out_channels )






       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = None,    SCREAMING_SNAKE_CASE_ = True,    )  ->	Optional[Any]:
          # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
          #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
          #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
          # expects mask of shape:
          #   [batch, key_tokens]
          # adds singleton query_tokens dimension:
          #   [batch,                    1, key_tokens]
          # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
          #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
          #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
          if attention_mask is not None and attention_mask.ndim == 2:
             # assume that mask is expressed as:
             #   (1 = keep,      0 = discard)
             # convert mask into a bias that can be added to attention scores:
             #       (keep = +0,     discard = -10000.0)
             UpperCamelCase					:						List[str]					=			(1 - attention_mask.to(hidden_states.dtype )) * -1_00_00.0
             UpperCamelCase					:						Tuple					=			attention_mask.unsqueeze(1 )

          # convert encoder_attention_mask to a bias the same way we do for attention_mask
          if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
             UpperCamelCase					:						int					=			(1 - encoder_attention_mask.to(hidden_states.dtype )) * -1_00_00.0
             UpperCamelCase					:						Dict					=			encoder_attention_mask.unsqueeze(1 )

          # 1. Input
          if self.is_input_continuous:
             UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase					:						Union[str, Any]					=			hidden_states.shape
             UpperCamelCase					:						str					=			hidden_states

             UpperCamelCase					:						Union[str, Any]					=			self.norm(SCREAMING_SNAKE_CASE_ )
             if not self.use_linear_projection:
                UpperCamelCase					:						Tuple					=			self.proj_in(SCREAMING_SNAKE_CASE_ )
                UpperCamelCase					:						Optional[int]					=			hidden_states.shape[1]
                UpperCamelCase					:						Tuple					=			hidden_states.permute(0,    2,    3,    1 ).reshape(SCREAMING_SNAKE_CASE_,    height * width,    SCREAMING_SNAKE_CASE_ )
             else:
                UpperCamelCase					:						Optional[int]					=			hidden_states.shape[1]
                UpperCamelCase					:						Dict					=			hidden_states.permute(0,    2,    3,    1 ).reshape(SCREAMING_SNAKE_CASE_,    height * width,    SCREAMING_SNAKE_CASE_ )
                UpperCamelCase					:						Union[str, Any]					=			self.proj_in(SCREAMING_SNAKE_CASE_ )
          elif self.is_input_vectorized:
             UpperCamelCase					:						Union[str, Any]					=			self.latent_image_embedding(SCREAMING_SNAKE_CASE_ )
          elif self.is_input_patches:
             UpperCamelCase					:						Any					=			self.pos_embed(SCREAMING_SNAKE_CASE_ )

          # 2. Blocks
          for block in self.transformer_blocks:
             UpperCamelCase					:						Optional[int]					=			block(
                 SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    encoder_hidden_states=SCREAMING_SNAKE_CASE_,    encoder_attention_mask=SCREAMING_SNAKE_CASE_,    timestep=SCREAMING_SNAKE_CASE_,    cross_attention_kwargs=SCREAMING_SNAKE_CASE_,    class_labels=SCREAMING_SNAKE_CASE_,    )

          # 3. Output
          if self.is_input_continuous:
             if not self.use_linear_projection:
                UpperCamelCase					:						Tuple					=			hidden_states.reshape(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ ).permute(0,    3,    1,    2 ).contiguous()
                UpperCamelCase					:						int					=			self.proj_out(SCREAMING_SNAKE_CASE_ )
             else:
                UpperCamelCase					:						Optional[int]					=			self.proj_out(SCREAMING_SNAKE_CASE_ )
                UpperCamelCase					:						Optional[int]					=			hidden_states.reshape(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ ).permute(0,    3,    1,    2 ).contiguous()

             UpperCamelCase					:						Any					=			hidden_states + residual
          elif self.is_input_vectorized:
             UpperCamelCase					:						List[Any]					=			self.norm_out(SCREAMING_SNAKE_CASE_ )
             UpperCamelCase					:						Optional[int]					=			self.out(SCREAMING_SNAKE_CASE_ )
             # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
             UpperCamelCase					:						Tuple					=			logits.permute(0,    2,    1 )

             # log(p(x_0))
             UpperCamelCase					:						List[str]					=			F.log_softmax(logits.double(),    dim=1 ).float()
          elif self.is_input_patches:
             # TODO: cleanup!
             UpperCamelCase					:						int					=			self.transformer_blocks[0].norma.emb(
                 SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    hidden_dtype=hidden_states.dtype )
             UpperCamelCase			,  UpperCamelCase					:						int					=			self.proj_out_a(F.silu(SCREAMING_SNAKE_CASE_ ) ).chunk(2,    dim=1 )
             UpperCamelCase					:						Dict					=			self.norm_out(SCREAMING_SNAKE_CASE_ ) * (1 + scale[:, None]) + shift[:, None]
             UpperCamelCase					:						Any					=			self.proj_out_a(SCREAMING_SNAKE_CASE_ )

             # unpatchify
             UpperCamelCase					:						Union[str, Any]					=			int(hidden_states.shape[1] ** 0.5 )
             UpperCamelCase					:						Tuple					=			hidden_states.reshape(
                 shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) )
             UpperCamelCase					:						int					=			torch.einsum('nhwpqc->nchpwq',    SCREAMING_SNAKE_CASE_ )
             UpperCamelCase					:						Any					=			hidden_states.reshape(
                 shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) )

          if not return_dict:
             return (output,)

          return TransformeraDModelOutput(sample=SCREAMING_SNAKE_CASE_ )



