



from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block



@dataclass
class    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




   A     :	torch.FloatTensor



class    UpperCAmelCase__(				nn.Module						):

   '''simple docstring'''




   def __init__(    self	:      Union[str, Any]      ,					lowerCAmelCase	:      Tuple=3      ,					lowerCAmelCase	:      int=3      ,					lowerCAmelCase	:      List[str]=("DownEncoderBlock2D",)      ,					lowerCAmelCase	:      Dict=(64,)      ,					lowerCAmelCase	:      str=2      ,					lowerCAmelCase	:      int=32      ,					lowerCAmelCase	:      Union[str, Any]="silu"      ,					lowerCAmelCase	:      List[Any]=True      ,					)						->					List[str]:




         """simple docstring"""

         super().__init__()
         lowercase__  =	layers_per_block

         lowercase__  =	torch.nn.Convad(
             lowerCAmelCase      ,					block_out_channels[0]      ,					kernel_size=3      ,					stride=1      ,					padding=1      ,					)

         lowercase__  =	None
         lowercase__  =	nn.ModuleList([])

         # down
         lowercase__  =	block_out_channels[0]
         for i, down_block_type in enumerate(lowerCAmelCase):
               lowercase__  =	output_channel
               lowercase__  =	block_out_channels[i]
               lowercase__  =	i == len(lowerCAmelCase) - 1

               lowercase__  =	get_down_block(
                   lowerCAmelCase      ,					num_layers=self.layers_per_block      ,					in_channels=lowerCAmelCase      ,					out_channels=lowerCAmelCase      ,					add_downsample=not is_final_block      ,					resnet_eps=1E-6      ,					downsample_padding=0      ,					resnet_act_fn=lowerCAmelCase      ,					resnet_groups=lowerCAmelCase      ,					attention_head_dim=lowerCAmelCase      ,					temb_channels=lowerCAmelCase      ,					)
               self.down_blocks.append(lowerCAmelCase)

         # mid
         lowercase__  =	UNetMidBlockaD(
             in_channels=block_out_channels[-1]      ,					resnet_eps=1E-6      ,					resnet_act_fn=lowerCAmelCase      ,					output_scale_factor=1      ,					resnet_time_scale_shift='default'      ,					attention_head_dim=block_out_channels[-1]      ,					resnet_groups=lowerCAmelCase      ,					temb_channels=lowerCAmelCase      ,					)

         # out
         lowercase__  =	nn.GroupNorm(num_channels=block_out_channels[-1]      ,					num_groups=lowerCAmelCase      ,					eps=1E-6)
         lowercase__  =	nn.SiLU()

         lowercase__  =	2 * out_channels if double_z else out_channels
         lowercase__  =	nn.Convad(block_out_channels[-1]      ,					lowerCAmelCase      ,					3      ,					padding=1)

         lowercase__  =	False



   def       UpperCAmelCase		(    self	:      Tuple      ,					lowerCAmelCase	:      Any)						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	x
         lowercase__  =	self.conv_in(lowerCAmelCase)

         if self.training and self.gradient_checkpointing:

               def create_custom_forward(lowerCAmelCase	:      Any):
                     def custom_forward(*lowerCAmelCase	:      str):
                           return module(*lowerCAmelCase)

                     return custom_forward

               # down
               if is_torch_version('>='      ,					'1.11.0'):
                     for down_block in self.down_blocks:
                           lowercase__  =	torch.utils.checkpoint.checkpoint(
                               create_custom_forward(lowerCAmelCase)      ,					lowerCAmelCase      ,					use_reentrant=lowerCAmelCase)
                     # middle
                     lowercase__  =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block)      ,					lowerCAmelCase      ,					use_reentrant=lowerCAmelCase)
               else:
                     for down_block in self.down_blocks:
                           lowercase__  =	torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase)      ,					lowerCAmelCase)
                     # middle
                     lowercase__  =	torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block)      ,					lowerCAmelCase)

         else:
               # down
               for down_block in self.down_blocks:
                     lowercase__  =	down_block(lowerCAmelCase)

               # middle
               lowercase__  =	self.mid_block(lowerCAmelCase)

         # post-process
         lowercase__  =	self.conv_norm_out(lowerCAmelCase)
         lowercase__  =	self.conv_act(lowerCAmelCase)
         lowercase__  =	self.conv_out(lowerCAmelCase)

         return sample



class    UpperCAmelCase__(				nn.Module						):

   '''simple docstring'''




   def __init__(    self	:      Union[str, Any]      ,					lowerCAmelCase	:      str=3      ,					lowerCAmelCase	:      int=3      ,					lowerCAmelCase	:      List[str]=("UpDecoderBlock2D",)      ,					lowerCAmelCase	:      str=(64,)      ,					lowerCAmelCase	:      Optional[Any]=2      ,					lowerCAmelCase	:      int=32      ,					lowerCAmelCase	:      Any="silu"      ,					lowerCAmelCase	:      int="group"      ,					)						->					Dict:




         """simple docstring"""

         super().__init__()
         lowercase__  =	layers_per_block

         lowercase__  =	nn.Convad(
             lowerCAmelCase      ,					block_out_channels[-1]      ,					kernel_size=3      ,					stride=1      ,					padding=1      ,					)

         lowercase__  =	None
         lowercase__  =	nn.ModuleList([])

         lowercase__  =	in_channels if norm_type == 'spatial' else None

         # mid
         lowercase__  =	UNetMidBlockaD(
             in_channels=block_out_channels[-1]      ,					resnet_eps=1E-6      ,					resnet_act_fn=lowerCAmelCase      ,					output_scale_factor=1      ,					resnet_time_scale_shift='default' if norm_type == 'group' else norm_type      ,					attention_head_dim=block_out_channels[-1]      ,					resnet_groups=lowerCAmelCase      ,					temb_channels=lowerCAmelCase      ,					)

         # up
         lowercase__  =	list(reversed(lowerCAmelCase))
         lowercase__  =	reversed_block_out_channels[0]
         for i, up_block_type in enumerate(lowerCAmelCase):
               lowercase__  =	output_channel
               lowercase__  =	reversed_block_out_channels[i]

               lowercase__  =	i == len(lowerCAmelCase) - 1

               lowercase__  =	get_up_block(
                   lowerCAmelCase      ,					num_layers=self.layers_per_block + 1      ,					in_channels=lowerCAmelCase      ,					out_channels=lowerCAmelCase      ,					prev_output_channel=lowerCAmelCase      ,					add_upsample=not is_final_block      ,					resnet_eps=1E-6      ,					resnet_act_fn=lowerCAmelCase      ,					resnet_groups=lowerCAmelCase      ,					attention_head_dim=lowerCAmelCase      ,					temb_channels=lowerCAmelCase      ,					resnet_time_scale_shift=lowerCAmelCase      ,					)
               self.up_blocks.append(lowerCAmelCase)
               lowercase__  =	output_channel

         # out
         if norm_type == "spatial":
               lowercase__  =	SpatialNorm(block_out_channels[0]      ,					lowerCAmelCase)
         else:
               lowercase__  =	nn.GroupNorm(num_channels=block_out_channels[0]      ,					num_groups=lowerCAmelCase      ,					eps=1E-6)
         lowercase__  =	nn.SiLU()
         lowercase__  =	nn.Convad(block_out_channels[0]      ,					lowerCAmelCase      ,					3      ,					padding=1)

         lowercase__  =	False



   def       UpperCAmelCase		(    self	:      str      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      Tuple=None)						->					Dict:




         """simple docstring"""

         lowercase__  =	z
         lowercase__  =	self.conv_in(lowerCAmelCase)

         lowercase__  =	next(iter(self.up_blocks.parameters())).dtype
         if self.training and self.gradient_checkpointing:

               def create_custom_forward(lowerCAmelCase	:      List[str]):
                     def custom_forward(*lowerCAmelCase	:      str):
                           return module(*lowerCAmelCase)

                     return custom_forward

               if is_torch_version('>='      ,					'1.11.0'):
                     # middle
                     lowercase__  =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block)      ,					lowerCAmelCase      ,					lowerCAmelCase      ,					use_reentrant=lowerCAmelCase)
                     lowercase__  =	sample.to(lowerCAmelCase)

                     # up
                     for up_block in self.up_blocks:
                           lowercase__  =	torch.utils.checkpoint.checkpoint(
                               create_custom_forward(lowerCAmelCase)      ,					lowerCAmelCase      ,					lowerCAmelCase      ,					use_reentrant=lowerCAmelCase)
               else:
                     # middle
                     lowercase__  =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block)      ,					lowerCAmelCase      ,					lowerCAmelCase)
                     lowercase__  =	sample.to(lowerCAmelCase)

                     # up
                     for up_block in self.up_blocks:
                           lowercase__  =	torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase)      ,					lowerCAmelCase      ,					lowerCAmelCase)
         else:
               # middle
               lowercase__  =	self.mid_block(lowerCAmelCase      ,					lowerCAmelCase)
               lowercase__  =	sample.to(lowerCAmelCase)

               # up
               for up_block in self.up_blocks:
                     lowercase__  =	up_block(lowerCAmelCase      ,					lowerCAmelCase)

        # post-process
         if latent_embeds is None:
               lowercase__  =	self.conv_norm_out(lowerCAmelCase)
         else:
               lowercase__  =	self.conv_norm_out(lowerCAmelCase      ,					lowerCAmelCase)
         lowercase__  =	self.conv_act(lowerCAmelCase)
         lowercase__  =	self.conv_out(lowerCAmelCase)

         return sample



class    UpperCAmelCase__(				nn.Module						):

   '''simple docstring'''




   def __init__(    self	:      List[Any]      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      Any=None      ,					lowerCAmelCase	:      Tuple="random"      ,					lowerCAmelCase	:      int=False      ,					lowerCAmelCase	:      int=True)						->					List[str]:




         """simple docstring"""

         super().__init__()
         lowercase__  =	n_e
         lowercase__  =	vq_embed_dim
         lowercase__  =	beta
         lowercase__  =	legacy

         lowercase__  =	nn.Embedding(self.n_e      ,					self.vq_embed_dim)
         self.embedding.weight.data.uniform_(-1.0 / self.n_e      ,					1.0 / self.n_e)

         lowercase__  =	remap
         if self.remap is not None:
               self.register_buffer('used'      ,					torch.tensor(np.load(self.remap)))
               lowercase__  =	self.used.shape[0]
               lowercase__  =	unknown_index  # "random" or "extra" or integer
               if self.unknown_index == "extra":
                     lowercase__  =	self.re_embed
                     lowercase__  =	self.re_embed + 1
               print(
                   f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
                   f'''Using {self.unknown_index} for unknown indices.''')
         else:
               lowercase__  =	n_e

         lowercase__  =	sane_index_shape


   def       UpperCAmelCase		(    self	:      Optional[Any]      ,					lowerCAmelCase	:      Union[str, Any])						->					Any:




         """simple docstring"""

         lowercase__  =	inds.shape
         assert len(lowerCAmelCase) > 1
         lowercase__  =	inds.reshape(ishape[0]      ,					-1)
         lowercase__  =	self.used.to(lowerCAmelCase)
         lowercase__  =	(inds[:, :, None] == used[None, None, ...]).long()
         lowercase__  =	match.argmax(-1)
         lowercase__  =	match.sum(2) < 1
         if self.unknown_index == "random":
               lowercase__  =	torch.randint(0      ,					self.re_embed      ,					size=new[unknown].shape).to(device=new.device)
         else:
               lowercase__  =	self.unknown_index
         return new.reshape(lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      Optional[int]      ,					lowerCAmelCase	:      List[str])						->					int:




         """simple docstring"""

         lowercase__  =	inds.shape
         assert len(lowerCAmelCase) > 1
         lowercase__  =	inds.reshape(ishape[0]      ,					-1)
         lowercase__  =	self.used.to(lowerCAmelCase)
         if self.re_embed > self.used.shape[0]:  # extra token
               lowercase__  =	0  # simply set to zero
         lowercase__  =	torch.gather(used[None, :][inds.shape[0] * [0], :]      ,					1      ,					lowerCAmelCase)
         return back.reshape(lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      str)						->					Optional[Any]:




         """simple docstring"""

         lowercase__  =	z.permute(0      ,					2      ,					3      ,					1).contiguous()
         lowercase__  =	z.view(-1      ,					self.vq_embed_dim)

         # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
         lowercase__  =	torch.argmin(torch.cdist(lowerCAmelCase      ,					self.embedding.weight)      ,					dim=1)

         lowercase__  =	self.embedding(lowerCAmelCase).view(z.shape)
         lowercase__  =	None
         lowercase__  =	None

         # compute loss for embedding
         if not self.legacy:
               lowercase__  =	self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
         else:
               lowercase__  =	torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)

         # preserve gradients
         lowercase__  =	z + (z_q - z).detach()

         # reshape back to match original input shape
         lowercase__  =	z_q.permute(0      ,					3      ,					1      ,					2).contiguous()

         if self.remap is not None:
               lowercase__  =	min_encoding_indices.reshape(z.shape[0]      ,					-1)  # add batch axis
               lowercase__  =	self.remap_to_used(lowerCAmelCase)
               lowercase__  =	min_encoding_indices.reshape(-1      ,					1)  # flatten

         if self.sane_index_shape:
               lowercase__  =	min_encoding_indices.reshape(z_q.shape[0]      ,					z_q.shape[2]      ,					z_q.shape[3])

         return z_q, loss, (perplexity, min_encodings, min_encoding_indices)



   def       UpperCAmelCase		(    self	:      Optional[Any]      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      List[Any])						->					List[Any]:




         """simple docstring"""

         if self.remap is not None:
               lowercase__  =	indices.reshape(shape[0]      ,					-1)  # add batch axis
               lowercase__  =	self.unmap_to_all(lowerCAmelCase)
               lowercase__  =	indices.reshape(-1)  # flatten again

         # get quantized latent vectors
         lowercase__  =	self.embedding(lowerCAmelCase)

         if shape is not None:
               lowercase__  =	z_q.view(lowerCAmelCase)
               # reshape back to match original input shape
               lowercase__  =	z_q.permute(0      ,					3      ,					1      ,					2).contiguous()

         return z_q



class    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




   def __init__(    self	:      str      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      Dict=False)						->					List[Any]:




         """simple docstring"""

         lowercase__  =	parameters
         lowercase__,						lowercase__  =	torch.chunk(lowerCAmelCase      ,					2      ,					dim=1)
         lowercase__  =	torch.clamp(self.logvar      ,					-30.0      ,					20.0)
         lowercase__  =	deterministic
         lowercase__  =	torch.exp(0.5 * self.logvar)
         lowercase__  =	torch.exp(self.logvar)
         if self.deterministic:
               lowercase__  =	lowercase__  =	torch.zeros_like(
                   self.mean      ,					device=self.parameters.device      ,					dtype=self.parameters.dtype)


   def       UpperCAmelCase		(    self	:      Any      ,					lowerCAmelCase	:      Optional[torch.Generator] = None)						->					torch.FloatTensor:




         """simple docstring"""

         lowercase__  =	randn_tensor(
             self.mean.shape      ,					generator=lowerCAmelCase      ,					device=self.parameters.device      ,					dtype=self.parameters.dtype)
         lowercase__  =	self.mean + self.std * sample
         return x


   def       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      List[Any]=None)						->					Any:




         """simple docstring"""

         if self.deterministic:
               return torch.Tensor([0.0])
         else:
               if other is None:
                     return 0.5 * torch.sum(torch.pow(self.mean      ,					2) + self.var - 1.0 - self.logvar      ,					dim=[1, 2, 3])
               else:
                     return 0.5 * torch.sum(
                         torch.pow(self.mean - other.mean      ,					2) / other.var
                         + self.var / other.var
                         - 1.0
                         - self.logvar
                         + other.logvar      ,					dim=[1, 2, 3]      ,					)


   def       UpperCAmelCase		(    self	:      str      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      List[Any]=[1, 2, 3])						->					List[str]:




         """simple docstring"""

         if self.deterministic:
               return torch.Tensor([0.0])
         lowercase__  =	np.log(2.0 * np.pi)
         return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean      ,					2) / self.var      ,					dim=lowerCAmelCase)



   def       UpperCAmelCase		(    self	:      Optional[Any])						->					Tuple:




         """simple docstring"""

         return self.mean





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

   '''simple docstring'''




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




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Dict)						->					Any:




         """simple docstring"""

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

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

         lowercase__  =	self.get_config()

         return config, pixel_values, labels


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




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Optional[int]      ,					lowerCAmelCase	:      Any      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      Union[str, Any])						->					int:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Dict      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      List[str])						->					Tuple:




         """simple docstring"""

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



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




         """simple docstring"""

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



@require_torch
class    UpperCAmelCase__(				lowerCamelCase  ,				lowerCamelCase  ,				unittest.TestCase						):

   '''simple docstring'''




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


   def       UpperCAmelCase		(    self	:      Optional[int])						->					Union[str, Any]:




         """simple docstring"""

         lowercase__  =	ViTHybridModelTester(self)
         lowercase__  =	ConfigTester(self      ,					config_class=lowerCAmelCase      ,					has_text_modality=lowerCAmelCase      ,					hidden_size=37)


   def       UpperCAmelCase		(    self	:      Dict)						->					Any:




         """simple docstring"""

         self.config_tester.run_common_tests()


   @unittest.skip(reason='ViT does not use inputs_embeds')
   def       UpperCAmelCase		(    self	:      Dict)						->					Dict:




         """simple docstring"""

         pass


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




         """simple docstring"""

         lowercase__,						lowercase__  =	self.model_tester.prepare_config_and_inputs_for_common()

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


   def       UpperCAmelCase		(    self	:      Tuple)						->					Optional[Any]:




         """simple docstring"""

         lowercase__,						lowercase__  =	self.model_tester.prepare_config_and_inputs_for_common()

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

               lowercase__  =	['pixel_values']
               self.assertListEqual(arg_names[:1]      ,					lowerCAmelCase)


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




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      int)						->					List[Any]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      str)						->					Tuple:




         """simple docstring"""

         lowercase__,						lowercase__  =	self.model_tester.prepare_config_and_inputs_for_common()

         lowercase__  =	_config_zero_init(lowerCAmelCase)
         for model_class in self.all_model_classes:
               lowercase__  =	model_class(config=lowerCAmelCase)
               # Skip the check for the backbone
               for name, module in model.named_modules():
                     if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
                           lowercase__  =	[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       UpperCAmelCase		(    self	:      Any)						->					Optional[int]:




         """simple docstring"""

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

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



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

   '''simple docstring'''




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




         """simple docstring"""

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


   @slow
   def       UpperCAmelCase		(    self	:      Union[str, Any])						->					Optional[Any]:




         """simple docstring"""

         lowercase__  =	ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(
             lowerCAmelCase)

         lowercase__  =	self.default_image_processor
         lowercase__  =	prepare_img()
         lowercase__  =	image_processor(images=lowerCAmelCase      ,					return_tensors='pt').to(lowerCAmelCase)

         # forward pass
         with torch.no_grad():
               lowercase__  =	model(**lowerCAmelCase)

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

         lowercase__  =	torch.tensor([-1.90_90, -0.49_93, -0.23_89]).to(lowerCAmelCase)

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



   @slow
   @require_accelerate
   def       UpperCAmelCase		(    self	:      int)						->					Dict:




         """simple docstring"""

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

         lowercase__  =	prepare_img()

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

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