






from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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 tensorflow as tf

  from transformers import TFCvtForImageClassification, TFCvtModel
  from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor


class       __lowercase      (	lowercase_  ):







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	self.config_class(**self.inputs_dict				)
         self.parent.assertTrue(hasattr(UpperCamelCase_ ,	"""embed_dim"""				)				)
         self.parent.assertTrue(hasattr(UpperCamelCase_ ,	"""num_heads"""				)				)


class       __lowercase      :







     '''simple docstring'''




     def __init__(				self       :			Tuple ,	UpperCamelCase_       :			int ,	UpperCamelCase_       :			Dict=13 ,	UpperCamelCase_       :			int=64 ,	UpperCamelCase_       :			Dict=3 ,	UpperCamelCase_       :			List[str]=[16, 48, 96] ,	UpperCamelCase_       :			Tuple=[1, 3, 6] ,	UpperCamelCase_       :			Optional[int]=[1, 2, 10] ,	UpperCamelCase_       :			List[str]=[7, 3, 3] ,	UpperCamelCase_       :			Any=[4, 2, 2] ,	UpperCamelCase_       :			str=[2, 1, 1] ,	UpperCamelCase_       :			List[str]=[2, 2, 2] ,	UpperCamelCase_       :			int=[False, False, True] ,	UpperCamelCase_       :			List[Any]=[0.0, 0.0, 0.0] ,	UpperCamelCase_       :			Tuple=0.02 ,	UpperCamelCase_       :			Dict=1e-12 ,	UpperCamelCase_       :			Any=True ,	UpperCamelCase_       :			int=True ,	UpperCamelCase_       :			str=2 ,	):


         """simple docstring"""



         __A         =	parent
         __A         =	batch_size
         __A         =	image_size
         __A         =	patch_sizes
         __A         =	patch_stride
         __A         =	patch_padding
         __A         =	is_training
         __A         =	use_labels
         __A         =	num_labels
         __A         =	num_channels
         __A         =	embed_dim
         __A         =	num_heads
         __A         =	stride_kv
         __A         =	depth
         __A         =	cls_token
         __A         =	attention_drop_rate
         __A         =	initializer_range
         __A         =	layer_norm_eps




     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



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

         __A         =	None
         if self.use_labels:
             # create a random int32 tensor of given shape
             __A         =	ids_tensor([self.batch_size] ,	self.num_labels				)

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




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         return CvtConfig(
             image_size=self.image_size ,	num_labels=self.num_labels ,	num_channels=self.num_channels ,	embed_dim=self.embed_dim ,	num_heads=self.num_heads ,	patch_sizes=self.patch_sizes ,	patch_padding=self.patch_padding ,	patch_stride=self.patch_stride ,	stride_kv=self.stride_kv ,	depth=self.depth ,	cls_token=self.cls_token ,	attention_drop_rate=self.attention_drop_rate ,	initializer_range=self.initializer_range ,	)




     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			List[Any] ,	UpperCamelCase_       :			Any				):


         """simple docstring"""



         __A         =	TFCvtModel(config=UpperCamelCase_				)
         __A         =	model(UpperCamelCase_ ,	training=UpperCamelCase_				)
         __A         =	(self.image_size, self.image_size)
         __A    ,	__A         =	image_size[0], image_size[1]
         for i in range(len(self.depth				)				):
             __A         =	floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1				)
             __A         =	floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1				)
         self.parent.assertEqual(result.last_hidden_state.shape ,	(self.batch_size, self.embed_dim[-1], height, width)				)




     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			Union[str, Any]				):


         """simple docstring"""



         __A         =	self.num_labels
         __A         =	TFCvtForImageClassification(UpperCamelCase_				)
         __A         =	model(UpperCamelCase_ ,	labels=UpperCamelCase_ ,	training=UpperCamelCase_				)
         self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_labels)				)





     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	self.prepare_config_and_inputs()
         __A    ,	__A    ,	__A         =	config_and_inputs
         __A         =	{"""pixel_values""": pixel_values}
         return config, inputs_dict


@require_tf
class       __lowercase      (	lowercase_			,      lowercase_			,      unittest.TestCase  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE   						=				(
         {"feature-extraction": TFCvtModel, "image-classification": TFCvtForImageClassification}
         if is_tf_available()
         else {}
     )
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	TFCvtModelTester(self				)
         __A         =	TFCvtConfigTester(self ,	config_class=UpperCamelCase_ ,	has_text_modality=UpperCamelCase_ ,	hidden_size=37				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         self.config_tester.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()




     @unittest.skip(reason="""Cvt does not output attentions"""				)
     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         pass




     @unittest.skip(reason="""Cvt does not use inputs_embeds"""				)
     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         pass




     @unittest.skip(reason="""Cvt does not support input and output embeddings"""				)
     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         pass




     @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 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         super().test_dataset_conversion()




     @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 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         super().test_keras_fit()




     @unittest.skip(reason="""Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"""				)
     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	tf.keras.mixed_precision.Policy("""mixed_float16"""				)
         tf.keras.mixed_precision.set_global_policy(UpperCamelCase_				)
         super().test_keras_fit()
         tf.keras.mixed_precision.set_global_policy("""float32"""				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """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.call				)
             # 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       :			int				):


         """simple docstring"""



         def check_hidden_states_output(UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Optional[int]				):
             __A         =	model_class(UpperCamelCase_				)

             __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)
             __A         =	outputs.hidden_states

             __A         =	len(self.model_tester.depth				)
             self.assertEqual(len(UpperCamelCase_				) ,	UpperCamelCase_				)

             # verify the first hidden states (first block)
             self.assertListEqual(
                 list(hidden_states[0].shape[-3:]				) ,	[
                     self.model_tester.embed_dim[0],
                     self.model_tester.image_size // 4,
                     self.model_tester.image_size // 4,
                 ] ,	)

         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             __A         =	True
             check_hidden_states_output(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)

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

             check_hidden_states_output(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*UpperCamelCase_				)





     @slow
     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A         =	TFCvtModel.from_pretrained(UpperCamelCase_				)
             self.assertIsNotNone(UpperCamelCase_				)



def 			_SCREAMING_SNAKE_CASE						( )						->      str:





    """simple docstring"""

    __A         =	Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
    return image






@require_tf
@require_vision
class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




     @cached_property
     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)





     @slow
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)

         __A         =	self.default_image_processor
         __A         =	prepare_img()
         __A         =	image_processor(images=UpperCamelCase_ ,	return_tensors="""tf"""				)

         # forward pass
         __A         =	model(**UpperCamelCase_				)

         # verify the logits
         __A         =	tf.TensorShape((1, 1_000)				)
         self.assertEqual(outputs.logits.shape ,	UpperCamelCase_				)

         __A         =	tf.constant([0.9285, 0.9015, -0.3150]				)
         self.assertTrue(np.allclose(outputs.logits[0, :3].numpy() ,	UpperCamelCase_ ,	atol=1e-4				)				)












from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer


@dataclass
class       __lowercase      (	lowercase_  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				42


class       __lowercase      (	lowercase_			,      lowercase_  ):







     '''simple docstring'''




     @register_to_config
     def __init__(				self       :			List[Any] ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			Tuple[str] = ("DownEncoderBlock2D",) ,	UpperCamelCase_       :			Tuple[str] = ("UpDecoderBlock2D",) ,	UpperCamelCase_       :			Tuple[int] = (64,) ,	UpperCamelCase_       :			int = 1 ,	UpperCamelCase_       :			str = "silu" ,	UpperCamelCase_       :			int = 3 ,	UpperCamelCase_       :			int = 32 ,	UpperCamelCase_       :			int = 256 ,	UpperCamelCase_       :			int = 32 ,	UpperCamelCase_       :			Optional[int] = None ,	UpperCamelCase_       :			float = 0.18215 ,	UpperCamelCase_       :			str = "group" ,	):


         """simple docstring"""



         super().__init__()

         # pass init params to Encoder
         __A         =	Encoder(
             in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	down_block_types=UpperCamelCase_ ,	block_out_channels=UpperCamelCase_ ,	layers_per_block=UpperCamelCase_ ,	act_fn=UpperCamelCase_ ,	norm_num_groups=UpperCamelCase_ ,	double_z=UpperCamelCase_ ,	)

         __A         =	vq_embed_dim if vq_embed_dim is not None else latent_channels

         __A         =	nn.Convad(UpperCamelCase_ ,	UpperCamelCase_ ,	1				)
         __A         =	VectorQuantizer(UpperCamelCase_ ,	UpperCamelCase_ ,	beta=0.25 ,	remap=UpperCamelCase_ ,	sane_index_shape=UpperCamelCase_				)
         __A         =	nn.Convad(UpperCamelCase_ ,	UpperCamelCase_ ,	1				)

         # pass init params to Decoder
         __A         =	Decoder(
             in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	up_block_types=UpperCamelCase_ ,	block_out_channels=UpperCamelCase_ ,	layers_per_block=UpperCamelCase_ ,	act_fn=UpperCamelCase_ ,	norm_num_groups=UpperCamelCase_ ,	norm_type=UpperCamelCase_ ,	)




     @apply_forward_hook
     def 		lowerCAmelCase_     (				self       :			List[str] ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         __A         =	self.encoder(UpperCamelCase_				)
         __A         =	self.quant_conv(UpperCamelCase_				)

         if not return_dict:
             return (h,)

         return VQEncoderOutput(latents=UpperCamelCase_				)




     @apply_forward_hook
     def 		lowerCAmelCase_     (				self       :			List[str] ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = False ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         if not force_not_quantize:
             __A    ,	__A    ,	__A         =	self.quantize(UpperCamelCase_				)
         else:
             __A         =	h
         __A         =	self.post_quant_conv(UpperCamelCase_				)
         __A         =	self.decoder(UpperCamelCase_ ,	quant if self.config.norm_type == """spatial""" else None				)

         if not return_dict:
             return (dec,)

         return DecoderOutput(sample=UpperCamelCase_				)





     def 		lowerCAmelCase_     (				self       :			Any ,	UpperCamelCase_       :			torch.FloatTensor ,	UpperCamelCase_       :			bool = True				):


         """simple docstring"""



         __A         =	sample
         __A         =	self.encode(UpperCamelCase_				).latents
         __A         =	self.decode(UpperCamelCase_				).sample

         if not return_dict:
             return (dec,)

         return DecoderOutput(sample=UpperCamelCase_				)




