
"""simple docstring"""





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    _UpperCAmelCase(   lowerCamelCase     ):



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






							'''simple docstring'''
							_UpperCamelCase										=  self.config_class(**self.inputs_dict)
							self.parent.assertTrue(hasattr(__a    ,		'''embed_dim'''))
							self.parent.assertTrue(hasattr(__a    ,		'''num_heads'''))






class    _UpperCAmelCase:



		def __init__( self    ,		__a    ,		__a=13    ,		__a=64    ,		__a=3    ,		__a=[16, 48, 96]    ,		__a=[1, 3, 6]    ,		__a=[1, 2, 10]    ,		__a=[7, 3, 3]    ,		__a=[4, 2, 2]    ,		__a=[2, 1, 1]    ,		__a=[2, 2, 2]    ,		__a=[False, False, True]    ,		__a=[0.0, 0.0, 0.0]    ,		__a=0.02    ,		__a=1e-12    ,		__a=True    ,		__a=True    ,		__a=2    ,		)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  parent
							_UpperCamelCase										=  batch_size
							_UpperCamelCase										=  image_size
							_UpperCamelCase										=  patch_sizes
							_UpperCamelCase										=  patch_stride
							_UpperCamelCase										=  patch_padding
							_UpperCamelCase										=  is_training
							_UpperCamelCase										=  use_labels
							_UpperCamelCase										=  num_labels
							_UpperCamelCase										=  num_channels
							_UpperCamelCase										=  embed_dim
							_UpperCamelCase										=  num_heads
							_UpperCamelCase										=  stride_kv
							_UpperCamelCase										=  depth
							_UpperCamelCase										=  cls_token
							_UpperCamelCase										=  attention_drop_rate
							_UpperCamelCase										=  initializer_range
							_UpperCamelCase										=  layer_norm_eps



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






							'''simple docstring'''
							_UpperCamelCase										=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

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



		def 						UpperCAmelCase							( self)	->	int:






							'''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 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a)	->	List[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  TFCvtModel(config=__a)
							_UpperCamelCase										=  model(__a    ,		training=__a)
							_UpperCamelCase										=  (self.image_size, self.image_size)
							_UpperCamelCase			,  _UpperCamelCase										=  image_size[0], image_size[1]
							for i in range(len(self.depth)):
												_UpperCamelCase										=  floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1)
												_UpperCamelCase										=  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 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a)	->	List[str]:






							'''simple docstring'''
							_UpperCamelCase										=  self.num_labels
							_UpperCamelCase										=  TFCvtForImageClassification(__a)
							_UpperCamelCase										=  model(__a    ,		labels=__a    ,		training=__a)
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_labels))




		def 						UpperCAmelCase							( self)	->	Any:






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






@require_tf
class    _UpperCAmelCase(   lowerCamelCase				,					lowerCamelCase				,					unittest.TestCase     ):
		lowercase__      =				(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
		lowercase__      =				(
		    {'feature-extraction': TFCvtModel, 'image-classification': TFCvtForImageClassification}
		    if is_tf_available()
		    else {}
		)
		lowercase__      =				False
		lowercase__      =				False
		lowercase__      =				False
		lowercase__      =				False
		lowercase__      =				False



		def 						UpperCAmelCase							( self)	->	Any:






							'''simple docstring'''
							_UpperCamelCase										=  TFCvtModelTester(self)
							_UpperCamelCase										=  TFCvtConfigTester(self    ,		config_class=__a    ,		has_text_modality=__a    ,		hidden_size=37)



		def 						UpperCAmelCase							( self)	->	Union[str, 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 						UpperCAmelCase							( self)	->	Union[str, Any]:






							'''simple docstring'''
							pass



		@unittest.skip(reason='''Cvt does not use inputs_embeds''')
		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							pass



		@unittest.skip(reason='''Cvt does not support input and output embeddings''')
		def 						UpperCAmelCase							( self)	->	int:






							'''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 						UpperCAmelCase							( self)	->	Tuple:






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






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






							'''simple docstring'''
							_UpperCamelCase										=  tf.keras.mixed_precision.Policy('''mixed_float16''')
							tf.keras.mixed_precision.set_global_policy(__a)
							super().test_keras_fit()
							tf.keras.mixed_precision.set_global_policy('''float32''')



		def 						UpperCAmelCase							( self)	->	str:






							'''simple docstring'''
							_UpperCamelCase			,  _UpperCamelCase										=  self.model_tester.prepare_config_and_inputs_for_common()

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

												_UpperCamelCase										=  ['''pixel_values''']
												self.assertListEqual(arg_names[:1]    ,		__a)



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






							'''simple docstring'''
							def check_hidden_states_output(__a    ,		__a    ,		__a):
												_UpperCamelCase										=  model_class(__a)

												_UpperCamelCase										=  model(**self._prepare_for_class(__a    ,		__a))
												_UpperCamelCase										=  outputs.hidden_states

												_UpperCamelCase										=  len(self.model_tester.depth)
												self.assertEqual(len(__a)    ,		__a)

												# 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,
												    ]    ,		)

							_UpperCamelCase			,  _UpperCamelCase										=  self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												_UpperCamelCase										=  True
												check_hidden_states_output(__a    ,		__a    ,		__a)

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

												check_hidden_states_output(__a    ,		__a    ,		__a)



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






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__a)



		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*__a)




		@slow
		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_UpperCamelCase										=  TFCvtModel.from_pretrained(__a)
												self.assertIsNotNone(__a)





def 				lowerCamelCase__		(      )							->		Tuple:


					"""simple docstring"""





					_UpperCamelCase										=  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''     )
					return image






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



		@cached_property
		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0])




		@slow
		def 						UpperCAmelCase							( self)	->	Dict:






							'''simple docstring'''
							_UpperCamelCase										=  TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0])

							_UpperCamelCase										=  self.default_image_processor
							_UpperCamelCase										=  prepare_img()
							_UpperCamelCase										=  image_processor(images=__a    ,		return_tensors='''tf''')

							# forward pass
							_UpperCamelCase										=  model(**__a)

							# verify the logits
							_UpperCamelCase										=  tf.TensorShape((1, 10_00))
							self.assertEqual(outputs.logits.shape    ,		__a)

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








"""simple docstring"""





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    _UpperCAmelCase(   lowerCamelCase     ):



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






							'''simple docstring'''
							super().__init__()
							self.register_modules(transformer=__a    ,		vae=__a    ,		scheduler=__a)

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



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






							'''simple docstring'''

							if not isinstance(__a    ,		__a):
												_UpperCamelCase										=  list(__a)

							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    ,		__a    ,		__a = 4.0    ,		__a = None    ,		__a = 50    ,		__a = "pil"    ,		__a = True    ,		)	->	Union[ImagePipelineOutput, Tuple]:






							'''simple docstring'''
							_UpperCamelCase										=  len(__a)
							_UpperCamelCase										=  self.transformer.config.sample_size
							_UpperCamelCase										=  self.transformer.config.in_channels

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

							_UpperCamelCase										=  torch.tensor(__a    ,		device=self.device).reshape(-1)
							_UpperCamelCase										=  torch.tensor([10_00] * batch_size    ,		device=self.device)
							_UpperCamelCase										=  torch.cat([class_labels, class_null]    ,		0) if guidance_scale > 1 else class_labels

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

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

												_UpperCamelCase										=  t
												if not torch.is_tensor(__a):
																	# 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+)
																	_UpperCamelCase										=  latent_model_input.device.type == '''mps'''
																	if isinstance(__a    ,		__a):
																						_UpperCamelCase										=  torch.floataa if is_mps else torch.floataa
																	else:
																						_UpperCamelCase										=  torch.intaa if is_mps else torch.intaa
																	_UpperCamelCase										=  torch.tensor([timesteps]    ,		dtype=__a    ,		device=latent_model_input.device)
												elif len(timesteps.shape) == 0:
																	_UpperCamelCase										=  timesteps[None].to(latent_model_input.device)
												# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
												_UpperCamelCase										=  timesteps.expand(latent_model_input.shape[0])
												# predict noise model_output
												_UpperCamelCase										=  self.transformer(
												    __a    ,		timestep=__a    ,		class_labels=__a).sample

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

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

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

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

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

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

							_UpperCamelCase										=  1 / self.vae.config.scaling_factor * latents
							_UpperCamelCase										=  self.vae.decode(__a).sample

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

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

							if output_type == "pil":
												_UpperCamelCase										=  self.numpy_to_pil(__a)

							if not return_dict:
												return (samples,)

							return ImagePipelineOutput(images=__a)






