






import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
     import torch
     from torch import nn

     from transformers import ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
     from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image






class 						lowercase_   :






      '''simple docstring'''




      def __init__( self :		str					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Optional[Any]=13					,   __UpperCAmelCase :		Dict=64					,   __UpperCAmelCase :		Tuple=2					,   __UpperCAmelCase :		Optional[int]=3					,   __UpperCAmelCase :		Dict=True					,   __UpperCAmelCase :		Dict=True					,   __UpperCAmelCase :		Tuple=32					,   __UpperCAmelCase :		List[str]=5					,   __UpperCAmelCase :		Union[str, Any]=4					,   __UpperCAmelCase :		Dict=37					,   __UpperCAmelCase :		List[str]="gelu"					,   __UpperCAmelCase :		Any=0.1					,   __UpperCAmelCase :		str=0.1					,   __UpperCAmelCase :		Dict=10					,   __UpperCAmelCase :		Optional[int]=0.02					,   __UpperCAmelCase :		Dict=[1, 16, 4, 4]					,   __UpperCAmelCase :		Union[str, Any]=None					,   )       ->Optional[Any]:




       """simple docstring"""


       a						    = parent
       a						    = batch_size
       a						    = image_size
       a						    = patch_size
       a						    = num_channels
       a						    = is_training
       a						    = use_labels
       a						    = hidden_size
       a						    = num_hidden_layers
       a						    = num_attention_heads
       a						    = intermediate_size
       a						    = hidden_act
       a						    = hidden_dropout_prob
       a						    = attention_probs_dropout_prob
       a						    = type_sequence_label_size
       a						    = initializer_range
       a						    = scope
       a						    = backbone_featmap_shape

       # in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       # the number of patches is based on the feature map of the backbone, which by default uses an output stride
       # of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
       a						    = (self.image_size // 32) ** 2
       a						    = num_patches + 1







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




       """simple docstring"""


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

       a						    = None
       if self.use_labels:
        a						    = ids_tensor([self.batch_size]					,   self.type_sequence_label_size    )

       a						    = self.get_config()

       return config, pixel_values, labels







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




       """simple docstring"""


       a						    = {
           '''global_padding''': '''same''',
           '''layer_type''': '''bottleneck''',
           '''depths''': [3, 4, 9],
           '''out_features''': ['''stage1''', '''stage2''', '''stage3'''],
           '''embedding_dynamic_padding''': True,
           '''hidden_sizes''': [4, 8, 16, 32],
           '''num_groups''': 2,
       }

       return ViTHybridConfig(
           image_size=self.image_size					,   patch_size=self.patch_size					,   num_channels=self.num_channels					,   hidden_size=self.hidden_size					,   num_hidden_layers=self.num_hidden_layers					,   num_attention_heads=self.num_attention_heads					,   intermediate_size=self.intermediate_size					,   hidden_act=self.hidden_act					,   hidden_dropout_prob=self.hidden_dropout_prob					,   attention_probs_dropout_prob=self.attention_probs_dropout_prob					,   is_decoder=__UpperCAmelCase					,   initializer_range=self.initializer_range					,   backbone_featmap_shape=self.backbone_featmap_shape					,   backbone_config=__UpperCAmelCase					,   )







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		Union[str, Any]					,   __UpperCAmelCase :		List[Any]					,   __UpperCAmelCase :		Optional[int]    )       ->str:




       """simple docstring"""


       a						    = ViTHybridModel(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)    )







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		List[Any]    )       ->List[str]:




       """simple docstring"""


       a						    = self.type_sequence_label_size
       a						    = ViTHybridForImageClassification(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   labels=__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.type_sequence_label_size)    )







      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       a						    = self.prepare_config_and_inputs()
       a    ,	a    ,	a						    = config_and_inputs
       a						    = {'''pixel_values''': pixel_values}
       return config, inputs_dict






@require_torch
class 						lowercase_   (      lowercase					,    lowercase					,    unittest.TestCase  ):






      '''simple docstring'''




      __snake_case	    =    (ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
      __snake_case	    =    (
          {'''feature-extraction''': ViTHybridModel, '''image-classification''': ViTHybridForImageClassification}
          if is_torch_available()
          else {}
      )
      __snake_case	    =    False
      __snake_case	    =    False
      __snake_case	    =    False







      def  __lowerCAmelCase   ( self :		str    )       ->List[Any]:




       """simple docstring"""


       a						    = ViTHybridModelTester(self    )
       a						    = ConfigTester(self					,   config_class=__UpperCAmelCase					,   has_text_modality=__UpperCAmelCase					,   hidden_size=37    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


       self.config_tester.run_common_tests()







      @unittest.skip(reason='''ViT does not use inputs_embeds'''    )
      def  __lowerCAmelCase   ( self :		List[str]    )       ->int:




       """simple docstring"""


       pass







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        a						    = model_class(__UpperCAmelCase    )
        self.assertIsInstance(model.get_input_embeddings()					,   (nn.Module)    )
        a						    = model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(__UpperCAmelCase					,   nn.Linear    )    )







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




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()

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

        a						    = ['''pixel_values''']
        self.assertListEqual(arg_names[:1]					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		List[Any]    )       ->Dict:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->Any:




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()

       a						    = _config_zero_init(__UpperCAmelCase    )
       for model_class in self.all_model_classes:
        a						    = model_class(config=__UpperCAmelCase    )
        # Skip the check for the backbone
        for name, module in model.named_modules():
         if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
          a						    = [F"""{name}.{key}""" for key in module.state_dict().keys()]
          break

        for name, param in model.named_parameters():
         if param.requires_grad:
          if name in backbone_params:
           continue
          self.assertIn(
              ((param.data.mean() * 1e9).round() / 1e9).item()					,   [0.0, 1.0]					,   msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""					,   )







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




       """simple docstring"""


       for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        a						    = ViTHybridModel.from_pretrained(__UpperCAmelCase    )
        self.assertIsNotNone(__UpperCAmelCase    )
def 		_a       ( )				->     Tuple:
 a						    = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
 return image






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






      '''simple docstring'''




      @cached_property
      def  __lowerCAmelCase   ( self :		str    )       ->Union[str, Any]:




       """simple docstring"""


       return (
           ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )
           if is_vision_available()
           else None
       )







      @slow
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Union[str, Any]:




       """simple docstring"""


       a						    = ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]    ).to(
           __UpperCAmelCase    )

       a						    = self.default_image_processor
       a						    = prepare_img()
       a						    = image_processor(images=__UpperCAmelCase					,   return_tensors='''pt'''    ).to(__UpperCAmelCase    )

       # forward pass
       with torch.no_grad():
        a						    = model(**__UpperCAmelCase    )

       # verify the logits
       a						    = torch.Size((1, 1_000)    )
       self.assertEqual(outputs.logits.shape					,   __UpperCAmelCase    )

       a						    = torch.tensor([-1.9090, -0.4993, -0.2389]    ).to(__UpperCAmelCase    )

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







      @slow
      @require_accelerate
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->str:




       """simple docstring"""


       a						    = ViTHybridImageProcessor.from_pretrained('''google/vit-hybrid-base-bit-384'''    )
       a						    = ViTHybridForImageClassification.from_pretrained('''google/vit-hybrid-base-bit-384'''					,   device_map='''auto'''    )

       a						    = prepare_img()

       a						    = image_processor(images=__UpperCAmelCase					,   return_tensors='''pt'''    )
       a						    = model(**__UpperCAmelCase    )
       a						    = outputs.logits
       # model predicts one of the 1000 ImageNet classes
       a						    = logits.argmax(-1    ).item()

       self.assertTrue(model.config.idalabel[predicted_class_idx]					,   '''tabby, tabby cat'''    )










# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")






class 						lowercase_   :






      '''simple docstring'''




      def __init__( self :		Optional[int]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		bool = True					,   __UpperCAmelCase :		bool = False    )       ->Tuple:




       """simple docstring"""


       a						    = scheduler
       a						    = optimizers if isinstance(__UpperCAmelCase					,   (list, tuple)    ) else [optimizers]
       a						    = split_batches
       a						    = step_with_optimizer
       a						    = GradientState()







      def  __lowerCAmelCase   ( self :		Tuple					,   *__UpperCAmelCase :		List[Any]					,   **__UpperCAmelCase :		Optional[Any]    )       ->int:




       """simple docstring"""


       if not self.step_with_optimizer:
        # No link between scheduler and optimizer -> just step
        self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
        return

       # Otherwise, first make sure the optimizer was stepped.
       if not self.gradient_state.sync_gradients:
        if self.gradient_state.adjust_scheduler:
         self.scheduler._step_count += 1
        return

       for opt in self.optimizers:
        if opt.step_was_skipped:
         return
       if self.split_batches:
        # Split batches -> the training dataloader batch size is not changed so one step per training step
        self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
       else:
        # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
        # num_processes steps per training step
        a						    = AcceleratorState().num_processes
        for _ in range(__UpperCAmelCase    ):
         # Special case when using OneCycle and `drop_last` was not used
         if hasattr(self.scheduler					,   '''total_steps'''    ):
          if self.scheduler._step_count <= self.scheduler.total_steps:
           self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
         else:
          self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->List[str]:




       """simple docstring"""


       return self.scheduler.get_last_lr()







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




       """simple docstring"""


       return self.scheduler.state_dict()







      def  __lowerCAmelCase   ( self :		Union[str, Any]					,   __UpperCAmelCase :		Any    )       ->Tuple:




       """simple docstring"""


       self.scheduler.load_state_dict(__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->Optional[int]:




       """simple docstring"""


       return self.scheduler.get_lr()







      def  __lowerCAmelCase   ( self :		int					,   *__UpperCAmelCase :		List[str]					,   **__UpperCAmelCase :		Optional[Any]    )       ->List[str]:




       """simple docstring"""


       return self.scheduler.print_lr(*__UpperCAmelCase					,   **__UpperCAmelCase    )


