






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

 def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=13		,    __UpperCAmelCase=[30, 30]		,    __UpperCAmelCase=2		,    __UpperCAmelCase=3		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=32		,    __UpperCAmelCase=5		,    __UpperCAmelCase=4		,    __UpperCAmelCase=37		,    __UpperCAmelCase="gelu"		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=10		,    __UpperCAmelCase=0.02		,    __UpperCAmelCase=3		,    __UpperCAmelCase=None		,    __UpperCAmelCase=8		,    __UpperCAmelCase=10		,    )->					Tuple:







       '''simple docstring'''





       lowerCAmelCase__       =						parent
       lowerCAmelCase__       =						batch_size
       lowerCAmelCase__       =						image_size
       lowerCAmelCase__       =						patch_size
       lowerCAmelCase__       =						num_channels
       lowerCAmelCase__       =						is_training
       lowerCAmelCase__       =						use_labels
       lowerCAmelCase__       =						hidden_size
       lowerCAmelCase__       =						num_hidden_layers
       lowerCAmelCase__       =						num_attention_heads
       lowerCAmelCase__       =						intermediate_size
       lowerCAmelCase__       =						hidden_act
       lowerCAmelCase__       =						hidden_dropout_prob
       lowerCAmelCase__       =						attention_probs_dropout_prob
       lowerCAmelCase__       =						type_sequence_label_size
       lowerCAmelCase__       =						initializer_range
       lowerCAmelCase__       =						num_labels
       lowerCAmelCase__       =						scope
       lowerCAmelCase__       =						n_targets
       lowerCAmelCase__       =						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
       lowerCAmelCase__       =						(image_size[1] // patch_size) * (image_size[0] // patch_size)
       lowerCAmelCase__       =						num_patches + 1 + self.num_detection_tokens

 def 		UpperCAmelCase  (							self       )->					Tuple:







       '''simple docstring'''





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

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

       lowerCAmelCase__       =						self.get_config()

       return config, pixel_values, labels

 def 		UpperCAmelCase  (							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 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Dict:







       '''simple docstring'''





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

 def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







       '''simple docstring'''





       lowerCAmelCase__       =						YolosForObjectDetection(__UpperCAmelCase       )
       model.to(__UpperCAmelCase       )
       model.eval()

       lowerCAmelCase__       =						model(pixel_values=__UpperCAmelCase       )
       lowerCAmelCase__       =						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)       )

       lowerCAmelCase__       =						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 		UpperCAmelCase  (							self       )->					Any:







       '''simple docstring'''





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





@require_torch
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase	):
 a_         =(YolosModel, YolosForObjectDetection) if is_torch_available() else ()
 a_         =(
     {"""feature-extraction""": YolosModel, """object-detection""": YolosForObjectDetection} if is_torch_available() else {}
 )

 a_         =False
 a_         =False
 a_         =False
 a_         =False

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







       '''simple docstring'''





       lowerCAmelCase__       =						super()._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase		,    return_labels=__UpperCAmelCase       )

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

       return inputs_dict

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







       '''simple docstring'''





       lowerCAmelCase__       =						YolosModelTester(self       )
       lowerCAmelCase__       =						ConfigTester(self		,    config_class=__UpperCAmelCase		,    has_text_modality=__UpperCAmelCase		,    hidden_size=37       )

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







       '''simple docstring'''





       self.config_tester.run_common_tests()

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







       '''simple docstring'''





       pass

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







       '''simple docstring'''





       lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()

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

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







       '''simple docstring'''





       lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()

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

             lowerCAmelCase__       =						["pixel_values"]
             self.assertListEqual(arg_names[:1]		,    __UpperCAmelCase       )

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







       '''simple docstring'''





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

 def 		UpperCAmelCase  (							self       )->					str:







       '''simple docstring'''





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

       # in YOLOS, the seq_len is different
       lowerCAmelCase__       =						self.model_tester.expected_seq_len
       for model_class in self.all_model_classes:
             lowerCAmelCase__       =						True
             lowerCAmelCase__       =						False
             lowerCAmelCase__       =						True
             lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
             model.to(__UpperCAmelCase       )
             model.eval()
             with torch.no_grad():
                   lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )
             lowerCAmelCase__       =						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"]
             lowerCAmelCase__       =						True
             lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
             model.to(__UpperCAmelCase       )
             model.eval()
             with torch.no_grad():
                   lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )
             lowerCAmelCase__       =						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]		,    )
             lowerCAmelCase__       =						len(__UpperCAmelCase       )

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

             lowerCAmelCase__       =						1
             self.assertEqual(out_len + added_hidden_states		,    len(__UpperCAmelCase       )       )

             lowerCAmelCase__       =						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 		UpperCAmelCase  (							self       )->					Dict:







       '''simple docstring'''





       def check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       ):
             lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
             model.to(__UpperCAmelCase       )
             model.eval()

             with torch.no_grad():
                   lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )

             lowerCAmelCase__       =						outputs.hidden_states

             lowerCAmelCase__       =						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
             lowerCAmelCase__       =						self.model_tester.expected_seq_len

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

       lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             lowerCAmelCase__       =						True
             check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

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

             check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

 def 		UpperCAmelCase  (							self       )->					int:







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_object_detection(*__UpperCAmelCase       )


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







       '''simple docstring'''





       for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCAmelCase__       =						YolosModel.from_pretrained(__UpperCAmelCase       )
             self.assertIsNotNone(__UpperCAmelCase       )





def        _a  (     ) ->   Union[str, Any]:
      """simple docstring"""







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





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

 @cached_property
 def 		UpperCAmelCase  (							self       )->					Optional[int]:







       '''simple docstring'''





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


 @slow
 def 		UpperCAmelCase  (							self       )->					int:







       '''simple docstring'''





       lowerCAmelCase__       =						YolosForObjectDetection.from_pretrained("hustvl/yolos-small"       ).to(__UpperCAmelCase       )

       lowerCAmelCase__       =						self.default_image_processor
       lowerCAmelCase__       =						prepare_img()
       lowerCAmelCase__       =						image_processor(images=__UpperCAmelCase		,    return_tensors="pt"       ).to(__UpperCAmelCase       )

       # forward pass
       with torch.no_grad():
             lowerCAmelCase__       =						model(inputs.pixel_values       )

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

       lowerCAmelCase__       =						torch.tensor(
           [[-24.0_248, -10.3_024, -14.8_290], [-42.0_392, -16.8_200, -27.4_334], [-27.2_743, -11.8_154, -18.7_148]]		,    device=__UpperCAmelCase		,    )
       lowerCAmelCase__       =						torch.tensor(
           [[0.2_559, 0.5_455, 0.4_706], [0.2_989, 0.7_279, 0.1_875], [0.7_732, 0.4_017, 0.4_462]]		,    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
       lowerCAmelCase__       =						image_processor.post_process_object_detection(
           __UpperCAmelCase		,    threshold=0.3		,    target_sizes=[image.size[::-1]]       )[0]
       lowerCAmelCase__       =						torch.tensor([0.9_994, 0.9_790, 0.9_964, 0.9_972, 0.9_861]       ).to(__UpperCAmelCase       )
       lowerCAmelCase__       =						[75, 75, 17, 63, 17]
       lowerCAmelCase__       =						torch.tensor([335.0_609, 79.3_848, 375.4_216, 187.2_495]       ).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 typing import Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class      lowercase__     (			_UpperCAmelCase	):

	def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase = None		,    )->					List[str]:







							'''simple docstring'''





							super().__init__()
							self.register_modules(transformer=__UpperCAmelCase		,    vae=__UpperCAmelCase		,    scheduler=__UpperCAmelCase       )

							# create a imagenet -> id dictionary for easier use
							lowerCAmelCase__       =						{}
							if idalabel is not None:
													for key, value in idalabel.items():
																			for label in value.split(","       ):
																									lowerCAmelCase__       =						int(__UpperCAmelCase       )
													lowerCAmelCase__       =						dict(sorted(self.labels.items()       )       )

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







							'''simple docstring'''





							if not isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
													lowerCAmelCase__       =						list(__UpperCAmelCase       )

							for l in label:
													if l not in self.labels:
																			raise ValueError(
																			    F"{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}."       )

							return [self.labels[l] for l in label]


	@torch.no_grad()
	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase = 4.0		,    __UpperCAmelCase = None		,    __UpperCAmelCase = 50		,    __UpperCAmelCase = "pil"		,    __UpperCAmelCase = True		,    )->					Union[ImagePipelineOutput, Tuple]:







							'''simple docstring'''





							lowerCAmelCase__       =						len(__UpperCAmelCase       )
							lowerCAmelCase__       =						self.transformer.config.sample_size
							lowerCAmelCase__       =						self.transformer.config.in_channels

							lowerCAmelCase__       =						randn_tensor(
							    shape=(batch_size, latent_channels, latent_size, latent_size)		,    generator=__UpperCAmelCase		,    device=self.device		,    dtype=self.transformer.dtype		,    )
							lowerCAmelCase__       =						torch.cat([latents] * 2       ) if guidance_scale > 1 else latents

							lowerCAmelCase__       =						torch.tensor(__UpperCAmelCase		,    device=self.device       ).reshape(-1       )
							lowerCAmelCase__       =						torch.tensor([1000] * batch_size		,    device=self.device       )
							lowerCAmelCase__       =						torch.cat([class_labels, class_null]		,    0       ) if guidance_scale > 1 else class_labels

							# set step values
							self.scheduler.set_timesteps(__UpperCAmelCase       )

							for t in self.progress_bar(self.scheduler.timesteps       ):
													if guidance_scale > 1:
																			lowerCAmelCase__       =						latent_model_input[: len(__UpperCAmelCase       ) // 2]
																			lowerCAmelCase__       =						torch.cat([half, half]		,    dim=0       )
													lowerCAmelCase__       =						self.scheduler.scale_model_input(__UpperCAmelCase		,    __UpperCAmelCase       )

													lowerCAmelCase__       =						t
													if not torch.is_tensor(__UpperCAmelCase       ):
																			# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
																			# This would be a good case for the `match` statement (Python 3.10+)
																			lowerCAmelCase__       =						latent_model_input.device.type == "mps"
																			if isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
																									lowerCAmelCase__       =						torch.floataa if is_mps else torch.floataa
																			else:
																									lowerCAmelCase__       =						torch.intaa if is_mps else torch.intaa
																			lowerCAmelCase__       =						torch.tensor([timesteps]		,    dtype=__UpperCAmelCase		,    device=latent_model_input.device       )
													elif len(timesteps.shape       ) == 0:
																			lowerCAmelCase__       =						timesteps[None].to(latent_model_input.device       )
													# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
													lowerCAmelCase__       =						timesteps.expand(latent_model_input.shape[0]       )
													# predict noise model_output
													lowerCAmelCase__       =						self.transformer(
													    __UpperCAmelCase		,    timestep=__UpperCAmelCase		,    class_labels=__UpperCAmelCase       ).sample

													# perform guidance
													if guidance_scale > 1:
																			lowerCAmelCase__      ,							lowerCAmelCase__       =						noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
																			lowerCAmelCase__      ,							lowerCAmelCase__       =						torch.split(__UpperCAmelCase		,    len(__UpperCAmelCase       ) // 2		,    dim=0       )

																			lowerCAmelCase__       =						uncond_eps + guidance_scale * (cond_eps - uncond_eps)
																			lowerCAmelCase__       =						torch.cat([half_eps, half_eps]		,    dim=0       )

																			lowerCAmelCase__       =						torch.cat([eps, rest]		,    dim=1       )

													# learned sigma
													if self.transformer.config.out_channels // 2 == latent_channels:
																			lowerCAmelCase__      ,							lowerCAmelCase__       =						torch.split(__UpperCAmelCase		,    __UpperCAmelCase		,    dim=1       )
													else:
																			lowerCAmelCase__       =						noise_pred

													# compute previous image: x_t -> x_t-1
													lowerCAmelCase__       =						self.scheduler.step(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       ).prev_sample

							if guidance_scale > 1:
													lowerCAmelCase__      ,							lowerCAmelCase__       =						latent_model_input.chunk(2		,    dim=0       )
							else:
													lowerCAmelCase__       =						latent_model_input

							lowerCAmelCase__       =						1 / self.vae.config.scaling_factor * latents
							lowerCAmelCase__       =						self.vae.decode(__UpperCAmelCase       ).sample

							lowerCAmelCase__       =						(samples / 2 + 0.5).clamp(0		,    1       )

							# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
							lowerCAmelCase__       =						samples.cpu().permute(0		,    2		,    3		,    1       ).float().numpy()

							if output_type == "pil":
													lowerCAmelCase__       =						self.numpy_to_pil(__UpperCAmelCase       )

							if not return_dict:
													return (samples,)

							return ImagePipelineOutput(images=__UpperCAmelCase       )






