





import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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


if is_torch_available():
   import torch

   from transformers import (
       MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
       MODEL_MAPPING,
       LevitForImageClassification,
       LevitForImageClassificationWithTeacher,
       LevitModel,
   )
   from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import LevitImageProcessor



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  def     lowercase		(    self			:		List[Any]					)							->						List[str]:
     __lowerCAmelCase							     =     self.config_class(**self.inputs_dict					)
     self.parent.assertTrue(hasattr(lowerCAmelCase_ ,  'hidden_sizes'					)					)
     self.parent.assertTrue(hasattr(lowerCAmelCase_ ,  'num_attention_heads'					)					)





class     _UpperCAmelCase   :

  """simple docstring"""


  def __init__(    self			:		List[Any] ,  lowerCAmelCase_			:		Dict ,  lowerCAmelCase_			:		Optional[Any]=1_3 ,  lowerCAmelCase_			:		Tuple=6_4 ,  lowerCAmelCase_			:		Any=3 ,  lowerCAmelCase_			:		Tuple=3 ,  lowerCAmelCase_			:		str=2 ,  lowerCAmelCase_			:		Dict=1 ,  lowerCAmelCase_			:		List[Any]=1_6 ,  lowerCAmelCase_			:		str=[1_2_8, 2_5_6, 3_8_4] ,  lowerCAmelCase_			:		Any=[4, 6, 8] ,  lowerCAmelCase_			:		Optional[int]=[2, 3, 4] ,  lowerCAmelCase_			:		List[str]=[1_6, 1_6, 1_6] ,  lowerCAmelCase_			:		List[Any]=0 ,  lowerCAmelCase_			:		Union[str, Any]=[2, 2, 2] ,  lowerCAmelCase_			:		Optional[Any]=[2, 2, 2] ,  lowerCAmelCase_			:		int=0.02 ,  lowerCAmelCase_			:		Dict=True ,  lowerCAmelCase_			:		List[Any]=True ,  lowerCAmelCase_			:		List[Any]=2 ,  )							->						Union[str, Any]:
     __lowerCAmelCase							     =     parent
     __lowerCAmelCase							     =     batch_size
     __lowerCAmelCase							     =     image_size
     __lowerCAmelCase							     =     num_channels
     __lowerCAmelCase							     =     kernel_size
     __lowerCAmelCase							     =     stride
     __lowerCAmelCase							     =     padding
     __lowerCAmelCase							     =     hidden_sizes
     __lowerCAmelCase							     =     num_attention_heads
     __lowerCAmelCase							     =     depths
     __lowerCAmelCase							     =     key_dim
     __lowerCAmelCase							     =     drop_path_rate
     __lowerCAmelCase							     =     patch_size
     __lowerCAmelCase							     =     attention_ratio
     __lowerCAmelCase							     =     mlp_ratio
     __lowerCAmelCase							     =     initializer_range
     __lowerCAmelCase							     =     [
         ['Subsample', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
         ['Subsample', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
     ]
     __lowerCAmelCase							     =     is_training
     __lowerCAmelCase							     =     use_labels
     __lowerCAmelCase							     =     num_labels
     __lowerCAmelCase							     =     initializer_range






  def     lowercase		(    self			:		List[str]					)							->						str:
     __lowerCAmelCase							     =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

     __lowerCAmelCase							     =     None
     if self.use_labels:
        __lowerCAmelCase							     =     ids_tensor([self.batch_size] ,  self.num_labels					)

     __lowerCAmelCase							     =     self.get_config()
     return config, pixel_values, labels






  def     lowercase		(    self			:		List[str]					)							->						int:
     return LevitConfig(
         image_size=self.image_size ,  num_channels=self.num_channels ,  kernel_size=self.kernel_size ,  stride=self.stride ,  padding=self.padding ,  patch_size=self.patch_size ,  hidden_sizes=self.hidden_sizes ,  num_attention_heads=self.num_attention_heads ,  depths=self.depths ,  key_dim=self.key_dim ,  drop_path_rate=self.drop_path_rate ,  mlp_ratio=self.mlp_ratio ,  attention_ratio=self.attention_ratio ,  initializer_range=self.initializer_range ,  down_ops=self.down_ops ,  )






  def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		Dict ,  lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		str					)							->						Optional[int]:
     __lowerCAmelCase							     =     LevitModel(config=lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()
     __lowerCAmelCase							     =     model(lowerCAmelCase_					)
     __lowerCAmelCase							     =     (self.image_size, self.image_size)
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     image_size[0], image_size[1]
     for _ in range(4					):
        __lowerCAmelCase							     =     floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1					)
        __lowerCAmelCase							     =     floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1					)
     self.parent.assertEqual(
         result.last_hidden_state.shape ,  (self.batch_size, ceil(height / 4					) * ceil(width / 4					), self.hidden_sizes[-1]) ,  )






  def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		Optional[int] ,  lowerCAmelCase_			:		Optional[Any] ,  lowerCAmelCase_			:		List[str]					)							->						Optional[int]:
     __lowerCAmelCase							     =     self.num_labels
     __lowerCAmelCase							     =     LevitForImageClassification(lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()
     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  labels=lowerCAmelCase_					)
     self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_labels)					)






  def     lowercase		(    self			:		Any					)							->						List[str]:
     __lowerCAmelCase							     =     self.prepare_config_and_inputs()
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     config_and_inputs
     __lowerCAmelCase							     =     {'pixel_values': pixel_values}
     return config, inputs_dict





@require_torch
class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase		,   unittest.TestCase   ):

  """simple docstring"""


  a_      =			(
      (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
      if is_torch_available()
      else ()
  )
  a_      =			(
      {
          """feature-extraction""": LevitModel,
          """image-classification""": (LevitForImageClassification, LevitForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  a_      =			False
  a_      =			False
  a_      =			False
  a_      =			False
  a_      =			False






  def     lowercase		(    self			:		Tuple					)							->						int:
     __lowerCAmelCase							     =     LevitModelTester(self					)
     __lowerCAmelCase							     =     ConfigTester(self ,  config_class=lowerCAmelCase_ ,  has_text_modality=lowerCAmelCase_ ,  hidden_size=3_7					)






  def     lowercase		(    self			:		List[str]					)							->						Dict:
     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()






  def     lowercase		(    self			:		str					)							->						Optional[Any]:
     return






  @unittest.skip(reason='Levit does not use inputs_embeds'					)
  def     lowercase		(    self			:		Optional[int]					)							->						List[Any]:
     pass






  @unittest.skip(reason='Levit does not support input and output embeddings'					)
  def     lowercase		(    self			:		Dict					)							->						str:
     pass






  @unittest.skip(reason='Levit does not output attentions'					)
  def     lowercase		(    self			:		List[str]					)							->						Optional[int]:
     pass






  def     lowercase		(    self			:		Dict					)							->						Dict:
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_common()

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

        __lowerCAmelCase							     =     ['pixel_values']
        self.assertListEqual(arg_names[:1] ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Optional[Any]					)							->						Dict:
     def check_hidden_states_output(lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		Optional[Any] ,  lowerCAmelCase_			:		Optional[int]					):
        __lowerCAmelCase							     =     model_class(lowerCAmelCase_					)
        model.to(lowerCAmelCase_					)
        model.eval()

        with torch.no_grad():
           __lowerCAmelCase							     =     model(**self._prepare_for_class(lowerCAmelCase_ ,  lowerCAmelCase_					)					)

        __lowerCAmelCase							     =     outputs.hidden_states

        __lowerCAmelCase							     =     len(self.model_tester.depths					) + 1
        self.assertEqual(len(lowerCAmelCase_					) ,  lowerCAmelCase_					)

        __lowerCAmelCase							     =     (self.model_tester.image_size, self.model_tester.image_size)
        __lowerCAmelCase     ,     __lowerCAmelCase							     =     image_size[0], image_size[1]
        for _ in range(4					):
           __lowerCAmelCase							     =     floor(
               (
                   (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                   / self.model_tester.stride
               )
               + 1					)
           __lowerCAmelCase							     =     floor(
               (
                   (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                   / self.model_tester.stride
               )
               + 1					)
        # verify the first hidden states (first block)
        self.assertListEqual(
            list(hidden_states[0].shape[-2:]					) ,  [
                height * width,
                self.model_tester.hidden_sizes[0],
            ] ,  )

     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        __lowerCAmelCase							     =     True
        check_hidden_states_output(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        __lowerCAmelCase							     =     True

        check_hidden_states_output(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)






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






  def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		List[str] ,  lowerCAmelCase_			:		Any=False					)							->						int:
     __lowerCAmelCase							     =     super()._prepare_for_class(lowerCAmelCase_ ,  lowerCAmelCase_ ,  return_labels=lowerCAmelCase_					)

     if return_labels:
        if model_class.__name__ == "LevitForImageClassificationWithTeacher":
           del inputs_dict["labels"]

     return inputs_dict






  def     lowercase		(    self			:		Any					)							->						List[Any]:
     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowerCAmelCase_					)






  def     lowercase		(    self			:		Union[str, Any]					)							->						str:
     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase_					)






  def     lowercase		(    self			:		int					)							->						Optional[int]:
     if not self.model_tester.is_training:
        return

     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_common()
     __lowerCAmelCase							     =     True

     for model_class in self.all_model_classes:
        # LevitForImageClassificationWithTeacher supports inference-only
        if (
            model_class in get_values(lowerCAmelCase_					)
            or model_class.__name__ == "LevitForImageClassificationWithTeacher"
        ):
           continue
        __lowerCAmelCase							     =     model_class(lowerCAmelCase_					)
        model.to(lowerCAmelCase_					)
        model.train()
        __lowerCAmelCase							     =     self._prepare_for_class(lowerCAmelCase_ ,  lowerCAmelCase_ ,  return_labels=lowerCAmelCase_					)
        __lowerCAmelCase							     =     model(**lowerCAmelCase_					).loss
        loss.backward()






  def     lowercase		(    self			:		Dict					)							->						Tuple:
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_common()
     if not self.model_tester.is_training:
        return

     __lowerCAmelCase							     =     False
     __lowerCAmelCase							     =     True

     for model_class in self.all_model_classes:
        if model_class in get_values(lowerCAmelCase_					) or not model_class.supports_gradient_checkpointing:
           continue
        # LevitForImageClassificationWithTeacher supports inference-only
        if model_class.__name__ == "LevitForImageClassificationWithTeacher":
           continue
        __lowerCAmelCase							     =     model_class(lowerCAmelCase_					)
        model.gradient_checkpointing_enable()
        model.to(lowerCAmelCase_					)
        model.train()
        __lowerCAmelCase							     =     self._prepare_for_class(lowerCAmelCase_ ,  lowerCAmelCase_ ,  return_labels=lowerCAmelCase_					)
        __lowerCAmelCase							     =     model(**lowerCAmelCase_					).loss
        loss.backward()






  def     lowercase		(    self			:		str					)							->						List[Any]:
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_common()

     __lowerCAmelCase							     =     [
         {'title': 'multi_label_classification', 'num_labels': 2, 'dtype': torch.float},
         {'title': 'single_label_classification', 'num_labels': 1, 'dtype': torch.long},
         {'title': 'regression', 'num_labels': 1, 'dtype': torch.float},
     ]

     for model_class in self.all_model_classes:
        if (
            model_class
            not in [
                *get_values(lowerCAmelCase_					),
            ]
            or model_class.__name__ == "LevitForImageClassificationWithTeacher"
        ):
           continue

        for problem_type in problem_types:
           with self.subTest(msg=f"""Testing {model_class} with {problem_type["title"]}"""					):
              __lowerCAmelCase							     =     problem_type['title']
              __lowerCAmelCase							     =     problem_type['num_labels']

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

              __lowerCAmelCase							     =     self._prepare_for_class(lowerCAmelCase_ ,  lowerCAmelCase_ ,  return_labels=lowerCAmelCase_					)

              if problem_type["num_labels"] > 1:
                 __lowerCAmelCase							     =     inputs['labels'].unsqueeze(1					).repeat(1 ,  problem_type['num_labels']					)

              __lowerCAmelCase							     =     inputs['labels'].to(problem_type['dtype']					)

              # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
              # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
              # they have the same size." which is a symptom something in wrong for the regression problem.
              # See https://github.com/huggingface/transformers/issues/11780
              with warnings.catch_warnings(record=lowerCAmelCase_					) as warning_list:
                 __lowerCAmelCase							     =     model(**lowerCAmelCase_					).loss
              for w in warning_list:
                 if "Using a target size that is different to the input size" in str(w.message					):
                    raise ValueError(
                        f"""Something is going wrong in the regression problem: intercepted {w.message}"""					)

              loss.backward()






  @slow
  def     lowercase		(    self			:		str					)							->						Optional[int]:
     for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase							     =     LevitModel.from_pretrained(lowerCAmelCase_					)
        self.assertIsNotNone(lowerCAmelCase_					)





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




@require_torch
@require_vision
class     _UpperCAmelCase   (   unittest.TestCase   ):

  """simple docstring"""


  @cached_property
  def     lowercase		(    self			:		Dict					)							->						Tuple:
     return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]					)






  @slow
  def     lowercase		(    self			:		Tuple					)							->						Dict:
     __lowerCAmelCase							     =     LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]					).to(
         lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.default_image_processor
     __lowerCAmelCase							     =     prepare_img()
     __lowerCAmelCase							     =     image_processor(images=lowerCAmelCase_ ,  return_tensors='pt'					).to(lowerCAmelCase_					)

     # forward pass
     with torch.no_grad():
        __lowerCAmelCase							     =     model(**lowerCAmelCase_					)

     # verify the logits
     __lowerCAmelCase							     =     torch.Size((1, 1_0_0_0)					)
     self.assertEqual(outputs.logits.shape ,  lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.tensor([1.04_48, -0.37_45, -1.83_17]					).to(lowerCAmelCase_					)

     self.assertTrue(torch.allclose(outputs.logits[0, :3] ,  lowerCAmelCase_ ,  atol=1e-4					)					)








from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin



@dataclass
class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			42





class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase   ):

  """simple docstring"""


  @register_to_config
  def __init__(    self			:		Optional[int] ,  lowerCAmelCase_			:		int = 1_6 ,  lowerCAmelCase_			:		int = 8_8 ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		int = 1 ,  lowerCAmelCase_			:		float = 0.0 ,  lowerCAmelCase_			:		int = 3_2 ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		bool = False ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		str = "geglu" ,  lowerCAmelCase_			:		bool = True ,  lowerCAmelCase_			:		bool = True ,  )							->						Dict:
     super().__init__()
     __lowerCAmelCase							     =     num_attention_heads
     __lowerCAmelCase							     =     attention_head_dim
     __lowerCAmelCase							     =     num_attention_heads * attention_head_dim

     __lowerCAmelCase							     =     in_channels

     __lowerCAmelCase							     =     torch.nn.GroupNorm(num_groups=lowerCAmelCase_ ,  num_channels=lowerCAmelCase_ ,  eps=1e-6 ,  affine=lowerCAmelCase_					)
     __lowerCAmelCase							     =     nn.Linear(lowerCAmelCase_ ,  lowerCAmelCase_					)

     # 3. Define transformers blocks
     __lowerCAmelCase							     =     nn.ModuleList(
         [
             BasicTransformerBlock(
                 lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  dropout=lowerCAmelCase_ ,  cross_attention_dim=lowerCAmelCase_ ,  activation_fn=lowerCAmelCase_ ,  attention_bias=lowerCAmelCase_ ,  double_self_attention=lowerCAmelCase_ ,  norm_elementwise_affine=lowerCAmelCase_ ,  )
             for d in range(lowerCAmelCase_					)
         ]					)

     __lowerCAmelCase							     =     nn.Linear(lowerCAmelCase_ ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		List[Any]=None ,  lowerCAmelCase_			:		Optional[int]=None ,  lowerCAmelCase_			:		Union[str, Any]=None ,  lowerCAmelCase_			:		Optional[int]=1 ,  lowerCAmelCase_			:		Optional[int]=None ,  lowerCAmelCase_			:		bool = True ,  )							->						str:
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     hidden_states.shape
     __lowerCAmelCase							     =     batch_frames // num_frames

     __lowerCAmelCase							     =     hidden_states

     __lowerCAmelCase							     =     hidden_states[None, :].reshape(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)
     __lowerCAmelCase							     =     hidden_states.permute(0 ,  2 ,  1 ,  3 ,  4					)

     __lowerCAmelCase							     =     self.norm(lowerCAmelCase_					)
     __lowerCAmelCase							     =     hidden_states.permute(0 ,  3 ,  4 ,  2 ,  1					).reshape(batch_size * height * width ,  lowerCAmelCase_ ,  lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.proj_in(lowerCAmelCase_					)

     # 2. Blocks
     for block in self.transformer_blocks:
        __lowerCAmelCase							     =     block(
            lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  timestep=lowerCAmelCase_ ,  cross_attention_kwargs=lowerCAmelCase_ ,  class_labels=lowerCAmelCase_ ,  )

     # 3. Output
     __lowerCAmelCase							     =     self.proj_out(lowerCAmelCase_					)
     __lowerCAmelCase							     =     (
         hidden_states[None, None, :]
         .reshape(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)
         .permute(0 ,  3 ,  4 ,  1 ,  2					)
         .contiguous()
     )
     __lowerCAmelCase							     =     hidden_states.reshape(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)

     __lowerCAmelCase							     =     hidden_states + residual

     if not return_dict:
        return (output,)

     return TransformerTemporalModelOutput(sample=lowerCAmelCase_					)

