






import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import 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, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import YolosForObjectDetection, YolosModel
    from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor


class       __lowercase      :







     '''simple docstring'''




     def __init__(				self       :			Any ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Optional[int]=13 ,	UpperCamelCase_       :			Any=[30, 30] ,	UpperCamelCase_       :			Optional[Any]=2 ,	UpperCamelCase_       :			List[Any]=3 ,	UpperCamelCase_       :			Union[str, Any]=True ,	UpperCamelCase_       :			str=True ,	UpperCamelCase_       :			int=32 ,	UpperCamelCase_       :			Optional[Any]=5 ,	UpperCamelCase_       :			List[str]=4 ,	UpperCamelCase_       :			str=37 ,	UpperCamelCase_       :			Any="gelu" ,	UpperCamelCase_       :			str=0.1 ,	UpperCamelCase_       :			Optional[int]=0.1 ,	UpperCamelCase_       :			Dict=10 ,	UpperCamelCase_       :			str=0.02 ,	UpperCamelCase_       :			Optional[Any]=3 ,	UpperCamelCase_       :			List[str]=None ,	UpperCamelCase_       :			Dict=8 ,	UpperCamelCase_       :			List[str]=10 ,	):


         """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         =	num_labels
         __A         =	scope
         __A         =	n_targets
         __A         =	num_detection_tokens
         # we set the expected sequence length (which is used in several tests)
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
         __A         =	(image_size[1] // patch_size) * (image_size[0] // patch_size)
         __A         =	num_patches + 1 + self.num_detection_tokens




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



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

         __A         =	None
         if self.use_labels:
             # labels is a list of Dict (each Dict being the labels for a given example in the batch)
             __A         =	[]
             for i in range(self.batch_size				):
                 __A         =	{}
                 __A         =	torch.randint(
                     high=self.num_labels ,	size=(self.n_targets,) ,	device=UpperCamelCase_				)
                 __A         =	torch.rand(self.n_targets ,	4 ,	device=UpperCamelCase_				)
                 labels.append(UpperCamelCase_				)

         __A         =	self.get_config()

         return config, pixel_values, labels




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         return YolosConfig(
             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 ,	num_detection_tokens=self.num_detection_tokens ,	num_labels=self.num_labels ,	)




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


         """simple docstring"""



         __A         =	YolosModel(config=UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()
         __A         =	model(UpperCamelCase_				)
         self.parent.assertEqual(
             result.last_hidden_state.shape ,	(self.batch_size, self.expected_seq_len, self.hidden_size)				)




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


         """simple docstring"""



         __A         =	YolosForObjectDetection(UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()

         __A         =	model(pixel_values=UpperCamelCase_				)
         __A         =	model(UpperCamelCase_				)

         self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)				)
         self.parent.assertEqual(result.pred_boxes.shape ,	(self.batch_size, self.num_detection_tokens, 4)				)

         __A         =	model(pixel_values=UpperCamelCase_ ,	labels=UpperCamelCase_				)

         self.parent.assertEqual(result.loss.shape ,	()				)
         self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)				)
         self.parent.assertEqual(result.pred_boxes.shape ,	(self.batch_size, self.num_detection_tokens, 4)				)





     def 		lowerCAmelCase_     (				self       :			Dict				):


         """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'''
     SCREAMING_SNAKE_CASE   						=				(YolosModel, YolosForObjectDetection) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE   						=				(
         {"feature-extraction": YolosModel, "object-detection": YolosForObjectDetection} if is_torch_available() else {}
     )

     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False




     def 		lowerCAmelCase_     (				self       :			Dict ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Any=False				):


         """simple docstring"""



         __A         =	super()._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_ ,	return_labels=UpperCamelCase_				)

         if return_labels:
             if model_class.__name__ == "YolosForObjectDetection":
                 __A         =	[]
                 for i in range(self.model_tester.batch_size				):
                     __A         =	{}
                     __A         =	torch.ones(
                         size=(self.model_tester.n_targets,) ,	device=UpperCamelCase_ ,	dtype=torch.long				)
                     __A         =	torch.ones(
                         self.model_tester.n_targets ,	4 ,	device=UpperCamelCase_ ,	dtype=torch.float				)
                     labels.append(UpperCamelCase_				)
                 __A         =	labels

         return inputs_dict




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         __A         =	YolosModelTester(self				)
         __A         =	ConfigTester(self ,	config_class=UpperCamelCase_ ,	has_text_modality=UpperCamelCase_ ,	hidden_size=37				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         self.config_tester.run_common_tests()




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         pass




     def 		lowerCAmelCase_     (				self       :			Optional[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_				)
             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       :			Dict				):


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


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



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

         # in YOLOS, the seq_len is different
         __A         =	self.model_tester.expected_seq_len
         for model_class in self.all_model_classes:
             __A         =	True
             __A         =	False
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)
             __A         =	outputs.attentions
             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)
             __A         =	outputs.attentions
             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)

             self.assertListEqual(
                 list(attentions[0].shape[-3:]				) ,	[self.model_tester.num_attention_heads, seq_len, seq_len] ,	)
             __A         =	len(UpperCamelCase_				)

             # Check attention is always last and order is fine
             __A         =	True
             __A         =	True
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()
             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)

             __A         =	1
             self.assertEqual(out_len + added_hidden_states ,	len(UpperCamelCase_				)				)

             __A         =	outputs.attentions

             self.assertEqual(len(UpperCamelCase_				) ,	self.model_tester.num_hidden_layers				)
             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]				) ,	[self.model_tester.num_attention_heads, seq_len, seq_len] ,	)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         def check_hidden_states_output(UpperCamelCase_       :			int ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Tuple				):
             __A         =	model_class(UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()

             with torch.no_grad():
                 __A         =	model(**self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)				)

             __A         =	outputs.hidden_states

             __A         =	getattr(
                 self.model_tester ,	"""expected_num_hidden_layers""" ,	self.model_tester.num_hidden_layers + 1				)
             self.assertEqual(len(UpperCamelCase_				) ,	UpperCamelCase_				)

             # YOLOS has a different seq_length
             __A         =	self.model_tester.expected_seq_len

             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]				) ,	[seq_length, self.model_tester.hidden_size] ,	)

         __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       :			Union[str, Any]				):


         """simple docstring"""



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





     @slow
     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A         =	YolosModel.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_torch
@require_vision
class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




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


         """simple docstring"""



         return AutoImageProcessor.from_pretrained("""hustvl/yolos-small"""				) if is_vision_available() else None





     @slow
     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         __A         =	YolosForObjectDetection.from_pretrained("""hustvl/yolos-small"""				).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(inputs.pixel_values				)

         # verify outputs
         __A         =	torch.Size((1, 100, 92)				)
         self.assertEqual(outputs.logits.shape ,	UpperCamelCase_				)

         __A         =	torch.tensor(
             [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]] ,	device=UpperCamelCase_ ,	)
         __A         =	torch.tensor(
             [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]] ,	device=UpperCamelCase_				)
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3] ,	UpperCamelCase_ ,	atol=1e-4				)				)
         self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3] ,	UpperCamelCase_ ,	atol=1e-4				)				)

         # verify postprocessing
         __A         =	image_processor.post_process_object_detection(
             UpperCamelCase_ ,	threshold=0.3 ,	target_sizes=[image.size[::-1]]				)[0]
         __A         =	torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]				).to(UpperCamelCase_				)
         __A         =	[75, 75, 17, 63, 17]
         __A         =	torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]				).to(UpperCamelCase_				)

         self.assertEqual(len(results["""scores"""]				) ,	5				)
         self.assertTrue(torch.allclose(results["""scores"""] ,	UpperCamelCase_ ,	atol=1e-4				)				)
         self.assertSequenceEqual(results["""labels"""].tolist() ,	UpperCamelCase_				)
         self.assertTrue(torch.allclose(results["""boxes"""][0, :] ,	UpperCamelCase_				)				)












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_				)




