code: 
  
  
  
  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 )
  
  
  
  
style_context: 
  
  
  
  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_ )
  
  
  
  