





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 					__snake_case	:


    '''simple docstring'''







    def __init__(       self   ,				A_   ,				A_=13   ,				A_=64   ,				A_=2   ,				A_=3   ,				A_=True   ,				A_=True   ,				A_=32   ,				A_=5   ,				A_=4   ,				A_=37   ,				A_="gelu"   ,				A_=0.1   ,				A_=0.1   ,				A_=10   ,				A_=0.02   ,				A_=[1, 16, 4, 4]   ,				A_=None   ,				):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			parent
        SCREAMING_SNAKE_CASE__											=			batch_size
        SCREAMING_SNAKE_CASE__											=			image_size
        SCREAMING_SNAKE_CASE__											=			patch_size
        SCREAMING_SNAKE_CASE__											=			num_channels
        SCREAMING_SNAKE_CASE__											=			is_training
        SCREAMING_SNAKE_CASE__											=			use_labels
        SCREAMING_SNAKE_CASE__											=			hidden_size
        SCREAMING_SNAKE_CASE__											=			num_hidden_layers
        SCREAMING_SNAKE_CASE__											=			num_attention_heads
        SCREAMING_SNAKE_CASE__											=			intermediate_size
        SCREAMING_SNAKE_CASE__											=			hidden_act
        SCREAMING_SNAKE_CASE__											=			hidden_dropout_prob
        SCREAMING_SNAKE_CASE__											=			attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE__											=			type_sequence_label_size
        SCREAMING_SNAKE_CASE__											=			initializer_range
        SCREAMING_SNAKE_CASE__											=			scope
        SCREAMING_SNAKE_CASE__											=			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
        SCREAMING_SNAKE_CASE__											=			(self.image_size // 32) ** 2
        SCREAMING_SNAKE_CASE__											=			num_patches + 1



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

        SCREAMING_SNAKE_CASE__											=			self.get_config()

        return config, pixel_values, labels



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			{
            '''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=A_   ,				initializer_range=self.initializer_range   ,				backbone_featmap_shape=self.backbone_featmap_shape   ,				backbone_config=A_   ,				)



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

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



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

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




    def  lowercase_					(       self     ):





        '''simple docstring'''

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





@require_torch
class 					__snake_case	(  __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     unittest.TestCase    ):


    '''simple docstring'''







    lowerCamelCase__    :		Tuple				=	(ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
    lowerCamelCase__    :		Any				=	(
        {"""feature-extraction""": ViTHybridModel, """image-classification""": ViTHybridForImageClassification}
        if is_torch_available()
        else {}
    )
    lowerCamelCase__    :		Union[str, Any]				=	False
    lowerCamelCase__    :		Optional[Any]				=	False
    lowerCamelCase__    :		int				=	False



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			ViTHybridModelTester(self     )
        SCREAMING_SNAKE_CASE__											=			ConfigTester(self   ,				config_class=A_   ,				has_text_modality=A_   ,				hidden_size=37     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    @unittest.skip(reason='''ViT does not use inputs_embeds'''     )
    def  lowercase_					(       self     ):





        '''simple docstring'''

        pass



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

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

            SCREAMING_SNAKE_CASE__											=			['''pixel_values''']
            self.assertListEqual(arg_names[:1]   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

        SCREAMING_SNAKE_CASE__											=			_config_zero_init(A_     )
        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__											=			model_class(config=A_     )
            # Skip the check for the backbone
            for name, module in model.named_modules():
                if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
                    SCREAMING_SNAKE_CASE__											=			[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  lowercase_					(       self     ):





        '''simple docstring'''

        for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            SCREAMING_SNAKE_CASE__											=			ViTHybridModel.from_pretrained(A_     )
            self.assertIsNotNone(A_     )





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



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


    '''simple docstring'''







    @cached_property
    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]     ).to(
            A_     )

        SCREAMING_SNAKE_CASE__											=			self.default_image_processor
        SCREAMING_SNAKE_CASE__											=			prepare_img()
        SCREAMING_SNAKE_CASE__											=			image_processor(images=A_   ,				return_tensors='''pt'''     ).to(A_     )

        # forward pass
        with torch.no_grad():
            SCREAMING_SNAKE_CASE__											=			model(**A_     )

        # verify the logits
        SCREAMING_SNAKE_CASE__											=			torch.Size((1, 10_00)     )
        self.assertEqual(outputs.logits.shape   ,				A_     )

        SCREAMING_SNAKE_CASE__											=			torch.tensor([-1.9090, -0.4993, -0.2389]     ).to(A_     )

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




    @slow
    @require_accelerate
    def  lowercase_					(       self     ):





        '''simple docstring'''

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

        SCREAMING_SNAKE_CASE__											=			prepare_img()

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

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







from ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ..utils.dummy_pt_objects import *  # noqa F403
else:
  from .scheduling_consistency_models import CMStochasticIterativeScheduler
  from .scheduling_ddim import DDIMScheduler
  from .scheduling_ddim_inverse import DDIMInverseScheduler
  from .scheduling_ddim_parallel import DDIMParallelScheduler
  from .scheduling_ddpm import DDPMScheduler
  from .scheduling_ddpm_parallel import DDPMParallelScheduler
  from .scheduling_deis_multistep import DEISMultistepScheduler
  from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
  from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
  from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
  from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
  from .scheduling_euler_discrete import EulerDiscreteScheduler
  from .scheduling_heun_discrete import HeunDiscreteScheduler
  from .scheduling_ipndm import IPNDMScheduler
  from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
  from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
  from .scheduling_karras_ve import KarrasVeScheduler
  from .scheduling_pndm import PNDMScheduler
  from .scheduling_repaint import RePaintScheduler
  from .scheduling_sde_ve import ScoreSdeVeScheduler
  from .scheduling_sde_vp import ScoreSdeVpScheduler
  from .scheduling_unclip import UnCLIPScheduler
  from .scheduling_unipc_multistep import UniPCMultistepScheduler
  from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
  from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ..utils.dummy_flax_objects import *  # noqa F403
else:
  from .scheduling_ddim_flax import FlaxDDIMScheduler
  from .scheduling_ddpm_flax import FlaxDDPMScheduler
  from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
  from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
  from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
  from .scheduling_pndm_flax import FlaxPNDMScheduler
  from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
  from .scheduling_utils_flax import (
      FlaxKarrasDiffusionSchedulers,
      FlaxSchedulerMixin,
      FlaxSchedulerOutput,
      broadcast_to_shape_from_left,
  )


try:
  if not (is_torch_available() and is_scipy_available()):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
  from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
  if not (is_torch_available() and is_torchsde_available()):
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
  from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler
