from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import RegNetConfig
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 TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST, TFRegNetForImageClassification, TFRegNetModel


if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor





class      lowerCAmelCase						:



					def __init__(					self				,      lowerCAmelCase__				,      lowerCAmelCase__=3				,      lowerCAmelCase__=32				,      lowerCAmelCase__=3				,      lowerCAmelCase__=10				,      lowerCAmelCase__=[10, 20, 30, 40]				,      lowerCAmelCase__=[1, 1, 2, 1]				,      lowerCAmelCase__=True				,      lowerCAmelCase__=True				,      lowerCAmelCase__="relu"				,      lowerCAmelCase__=3				,      lowerCAmelCase__=None				,      ):
												_A=	parent
												_A=	batch_size
												_A=	image_size
												_A=	num_channels
												_A=	embeddings_size
												_A=	hidden_sizes
												_A=	depths
												_A=	is_training
												_A=	use_labels
												_A=	hidden_act
												_A=	num_labels
												_A=	scope
												_A=	len(_A       )



					def 			a__  (					self       ):
												_A=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

												_A=	None
												if self.use_labels:
																			_A=	ids_tensor([self.batch_size]				,      self.num_labels       )

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



					def 			a__  (					self       ):
												return RegNetConfig(
												    num_channels=self.num_channels				,      embeddings_size=self.embeddings_size				,      hidden_sizes=self.hidden_sizes				,      depths=self.depths				,      hidden_act=self.hidden_act				,      num_labels=self.num_labels				,      )



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
												_A=	TFRegNetModel(config=_A       )
												_A=	model(_A				,      training=_A       )
												# expected last hidden states: B, C, H // 32, W // 32
												self.parent.assertEqual(
												    result.last_hidden_state.shape				,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)				,      )



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
												_A=	self.num_labels
												_A=	TFRegNetForImageClassification(_A       )
												_A=	model(_A				,      labels=_A				,      training=_A       )
												self.parent.assertEqual(result.logits.shape				,      (self.batch_size, self.num_labels)       )





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





@require_tf
class      lowerCAmelCase						(					_a			,		_a			,		unittest.TestCase			):
					_SCREAMING_SNAKE_CASE     :     List[Any]		=(TFRegNetModel, TFRegNetForImageClassification) if is_tf_available() else ()
					_SCREAMING_SNAKE_CASE     :     Optional[Any]		=(
					    {"""feature-extraction""": TFRegNetModel, """image-classification""": TFRegNetForImageClassification}
					    if is_tf_available()
					    else {}
					)

					_SCREAMING_SNAKE_CASE     :     Union[str, Any]		=False
					_SCREAMING_SNAKE_CASE     :     Union[str, Any]		=False
					_SCREAMING_SNAKE_CASE     :     List[str]		=False
					_SCREAMING_SNAKE_CASE     :     List[str]		=False
					_SCREAMING_SNAKE_CASE     :     List[Any]		=False



					def 			a__  (					self       ):
												_A=	TFRegNetModelTester(self       )
												_A=	ConfigTester(self				,      config_class=_A				,      has_text_modality=_A       )



					def 			a__  (					self       ):
												return



					@unittest.skip(reason='RegNet does not use inputs_embeds'       )
					def 			a__  (					self       ):
												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.'				,      )
					@slow
					def 			a__  (					self       ):
												super().test_keras_fit()



					@unittest.skip(reason='RegNet does not support input and output embeddings'       )
					def 			a__  (					self       ):
												pass



					def 			a__  (					self       ):
												_A,			_A=	self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			_A=	model_class(_A       )
																			_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]				,      _A       )



					def 			a__  (					self       ):
												_A=	self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_A       )



					def 			a__  (					self       ):
												def check_hidden_states_output(lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
																			_A=	model_class(_A       )
																			_A=	model(**self._prepare_for_class(_A				,      _A       )				,      training=_A       )

																			_A=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																			_A=	self.model_tester.num_stages
																			self.assertEqual(len(_A       )				,      expected_num_stages + 1       )

																			# RegNet's feature maps are of shape (batch_size, num_channels, height, width)
																			self.assertListEqual(
																			    list(hidden_states[0].shape[-2:]       )				,      [self.model_tester.image_size // 2, self.model_tester.image_size // 2]				,      )

												_A,			_A=	self.model_tester.prepare_config_and_inputs_for_common()
												_A=	['basic', 'bottleneck']
												for model_class in self.all_model_classes:
																			for layer_type in layers_type:
																										_A=	layer_type
																										_A=	True
																										check_hidden_states_output(_A				,      _A				,      _A       )

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

																										check_hidden_states_output(_A				,      _A				,      _A       )



					def 			a__  (					self       ):
												_A,			_A=	self.model_tester.prepare_config_and_inputs_for_common()

												def check_equivalence(lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__={}       ):
																			_A=	model(_A				,      return_dict=_A				,      **_A       )
																			_A=	model(_A				,      return_dict=_A				,      **_A       ).to_tuple()

																			def recursive_check(lowerCAmelCase__				,      lowerCAmelCase__       ):
																										if isinstance(_A				,      (List, Tuple)       ):
																																	for tuple_iterable_value, dict_iterable_value in zip(_A				,      _A       ):
																																								recursive_check(_A				,      _A       )
																										elif tuple_object is None:
																																	return
																										else:
																																	self.assertTrue(
																																	    all(tf.equal(_A				,      _A       )       )				,      msg=(
																																	        'Tuple and dict output are not equal. Difference:'
																																	        f" {tf.math.reduce_max(tf.abs(tuple_object - dict_object       )       )}"
																																	    )				,      )

																										recursive_check(_A				,      _A       )

												for model_class in self.all_model_classes:
																			_A=	model_class(_A       )

																			_A=	self._prepare_for_class(_A				,      _A       )
																			_A=	self._prepare_for_class(_A				,      _A       )
																			check_equivalence(_A				,      _A				,      _A       )

																			_A=	self._prepare_for_class(_A				,      _A				,      return_labels=_A       )
																			_A=	self._prepare_for_class(_A				,      _A				,      return_labels=_A       )
																			check_equivalence(_A				,      _A				,      _A       )

																			_A=	self._prepare_for_class(_A				,      _A       )
																			_A=	self._prepare_for_class(_A				,      _A       )
																			check_equivalence(_A				,      _A				,      _A				,      {'output_hidden_states': True}       )

																			_A=	self._prepare_for_class(_A				,      _A				,      return_labels=_A       )
																			_A=	self._prepare_for_class(_A				,      _A				,      return_labels=_A       )
																			check_equivalence(_A				,      _A				,      _A				,      {'output_hidden_states': True}       )



					def 			a__  (					self       ):
												_A=	self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*_A       )





					@slow
					def 			a__  (					self       ):
												for model_name in TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_A=	TFRegNetModel.from_pretrained(_A       )
																			self.assertIsNotNone(_A       )


def 							UpperCamelCase				(    )		->						str:




							'''simple docstring'''






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





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



					@cached_property
					def 			a__  (					self       ):
												return (
												    AutoImageProcessor.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )
												    if is_vision_available()
												    else None
												)





					@slow
					def 			a__  (					self       ):
												_A=	TFRegNetForImageClassification.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )

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

												# forward pass
												_A=	model(**_A				,      training=_A       )

												# verify the logits
												_A=	tf.TensorShape((1, 1000)       )
												self.assertEqual(outputs.logits.shape				,      _A       )

												_A=	tf.constant([-0.4180, -1.5051, -3.4836]       )

												tf.debugging.assert_near(outputs.logits[0, :3]				,      _A				,      atol=1E-4       )
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class      lowerCAmelCase						(					_a			,		_a			,		_a			,		unittest.TestCase			):
     _SCREAMING_SNAKE_CASE     :     Optional[int]		=AltDiffusionPipeline
     _SCREAMING_SNAKE_CASE     :     int		=TEXT_TO_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     List[Any]		=TEXT_TO_IMAGE_BATCH_PARAMS
     _SCREAMING_SNAKE_CASE     :     Dict		=TEXT_TO_IMAGE_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     Any		=TEXT_TO_IMAGE_IMAGE_PARAMS



     def 			a__  (					self       ):
            torch.manual_seed(0       )
            _A=	UNetaDConditionModel(
                block_out_channels=(32, 64)				,      layers_per_block=2				,      sample_size=32				,      in_channels=4				,      out_channels=4				,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')				,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')				,      cross_attention_dim=32				,      )
            _A=	DDIMScheduler(
                beta_start=0.00085				,      beta_end=0.012				,      beta_schedule='scaled_linear'				,      clip_sample=lowerCAmelCase__				,      set_alpha_to_one=lowerCAmelCase__				,      )
            torch.manual_seed(0       )
            _A=	AutoencoderKL(
                block_out_channels=[32, 64]				,      in_channels=3				,      out_channels=3				,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']				,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']				,      latent_channels=4				,      )

            # TODO: address the non-deterministic text encoder (fails for save-load tests)
            # torch.manual_seed(0)
            # text_encoder_config = RobertaSeriesConfig(
            #     hidden_size=32,
            #     project_dim=32,
            #     intermediate_size=37,
            #     layer_norm_eps=1e-05,
            #     num_attention_heads=4,
            #     num_hidden_layers=5,
            #     vocab_size=5002,
            # )
            # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

            torch.manual_seed(0       )
            _A=	CLIPTextConfig(
                bos_token_id=0				,      eos_token_id=2				,      hidden_size=32				,      projection_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=5002				,      )
            _A=	CLIPTextModel(lowerCAmelCase__       )

            _A=	XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'       )
            _A=	77

            _A=	{
                'unet': unet,
                'scheduler': scheduler,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer,
                'safety_checker': None,
                'feature_extractor': None,
            }
            return components



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0       ):
            if str(lowerCAmelCase__       ).startswith('mps'       ):
                   _A=	torch.manual_seed(lowerCAmelCase__       )
            else:
                   _A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
            _A=	{
                'prompt': 'A painting of a squirrel eating a burger',
                'generator': generator,
                'num_inference_steps': 2,
                'guidance_scale': 6.0,
                'output_type': 'numpy',
            }
            return inputs



     def 			a__  (					self       ):
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            super().test_inference_batch_single_identical(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder

            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	'A photo of an astronaut'
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2





     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            _A=	PNDMScheduler(skip_prk_steps=lowerCAmelCase__       )
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder
            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2





@slow
@require_torch_gpu
class      lowerCAmelCase						(					unittest.TestCase			):



     def 			a__  (					self       ):
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()



     def 			a__  (					self       ):
            # make sure here that pndm scheduler skips prk
            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )
            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      guidance_scale=6.0				,      num_inference_steps=20				,      output_type='np'       )

            _A=	output.images

            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2





     def 			a__  (					self       ):
            _A=	DDIMScheduler.from_pretrained('BAAI/AltDiffusion'				,      subfolder='scheduler'       )

            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      scheduler=lowerCAmelCase__				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )

            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      num_inference_steps=2				,      output_type='numpy'       )
            _A=	output.images

            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2