


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class   A__							(						snake_case__				,       snake_case__				,       unittest.TestCase					):






						"""simple docstring"""




						UpperCAmelCase     =  StableDiffusionXLPipeline
						UpperCAmelCase     =  TEXT_TO_IMAGE_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_IMAGE_PARAMS
						UpperCAmelCase     =  TEXT_TO_IMAGE_IMAGE_PARAMS




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


										"""simple docstring"""





										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=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''')			,      attention_head_dim=(2, 4)			,      use_linear_projection=_a			,      addition_embed_type='''text_time'''			,      addition_time_embed_dim=8			,      transformer_layers_per_block=(1, 2)			,      projection_class_embeddings_input_dim=80			,      cross_attention_dim=64			,      )
										_SCREAMING_SNAKE_CASE     					=EulerDiscreteScheduler(
										    beta_start=0.0_00_85			,      beta_end=0.0_12			,      steps_offset=1			,      beta_schedule='''scaled_linear'''			,      timestep_spacing='''leading'''			,      )
										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=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			,      sample_size=128			,      )
										torch.manual_seed(0					)
										_SCREAMING_SNAKE_CASE     					=CLIPTextConfig(
										    bos_token_id=0			,      eos_token_id=2			,      hidden_size=32			,      intermediate_size=37			,      layer_norm_eps=1E-05			,      num_attention_heads=4			,      num_hidden_layers=5			,      pad_token_id=1			,      vocab_size=1000			,      hidden_act='''gelu'''			,      projection_dim=32			,      )
										_SCREAMING_SNAKE_CASE     					=CLIPTextModel(_a					)
										_SCREAMING_SNAKE_CASE     					=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			,      local_files_only=_a					)

										_SCREAMING_SNAKE_CASE     					=CLIPTextModelWithProjection(_a					)
										_SCREAMING_SNAKE_CASE     					=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			,      local_files_only=_a					)

										_SCREAMING_SNAKE_CASE     					={
										    "unet": unet,
										    "scheduler": scheduler,
										    "vae": vae,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "text_encoder_2": text_encoder_a,
										    "tokenizer_2": tokenizer_a,
										    # "safety_checker": None,
										    # "feature_extractor": None,
										}
										return components




						def 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : List[Any]			,      _a   : str=0					)							->       int:


										"""simple docstring"""





										if str(_a					).startswith('''mps'''					):
														_SCREAMING_SNAKE_CASE     					=torch.manual_seed(_a					)
										else:
														_SCREAMING_SNAKE_CASE     					=torch.Generator(device=_a					).manual_seed(_a					)
										_SCREAMING_SNAKE_CASE     					={
										    "prompt": "A painting of a squirrel eating a burger",
										    "generator": generator,
										    "num_inference_steps": 2,
										    "guidance_scale": 5.0,
										    "output_type": "numpy",
										}
										return inputs




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					="cpu"  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE     					=self.get_dummy_components()
										_SCREAMING_SNAKE_CASE     					=StableDiffusionXLPipeline(**_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										sd_pipe.set_progress_bar_config(disable=_a					)

										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe(**_a					).images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1]

										assert image.shape == (1, 64, 64, 3)
										_SCREAMING_SNAKE_CASE     					=np.array([0.58_73, 0.61_28, 0.47_97, 0.51_22, 0.56_74, 0.46_39, 0.52_27, 0.51_49, 0.47_47]					)

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




						def 							__UpperCamelCase			(    self   : Tuple					)							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.get_dummy_components()
										_SCREAMING_SNAKE_CASE     					=StableDiffusionXLPipeline(**_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										_SCREAMING_SNAKE_CASE     					=sd_pipe.to(_a					)
										sd_pipe.set_progress_bar_config(disable=_a					)

										# forward without prompt embeds
										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=3 * ["this is a negative prompt"]
										_SCREAMING_SNAKE_CASE     					=negative_prompt
										_SCREAMING_SNAKE_CASE     					=3 * [inputs["prompt"]]

										_SCREAMING_SNAKE_CASE     					=sd_pipe(**_a					)
										_SCREAMING_SNAKE_CASE     					=output.images[0, -3:, -3:, -1]

										# forward with prompt embeds
										_SCREAMING_SNAKE_CASE     					=self.get_dummy_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=3 * ["this is a negative prompt"]
										_SCREAMING_SNAKE_CASE     					=3 * [inputs.pop('''prompt'''					)]

										(
										    _SCREAMING_SNAKE_CASE  
										)     					=sd_pipe.encode_prompt(_a			,      negative_prompt=_a					)

										_SCREAMING_SNAKE_CASE     					=sd_pipe(
										    **_a			,      prompt_embeds=_a			,      negative_prompt_embeds=_a			,      pooled_prompt_embeds=_a			,      negative_pooled_prompt_embeds=_a			,      )
										_SCREAMING_SNAKE_CASE     					=output.images[0, -3:, -3:, -1]

										# make sure that it's equal
										assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1E-4




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


										"""simple docstring"""





										super().test_attention_slicing_forward_pass(expected_max_diff=3E-3					)




						def 							__UpperCamelCase			(    self   : Tuple					)							->       Dict:


										"""simple docstring"""





										super().test_inference_batch_single_identical(expected_max_diff=3E-3					)

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






						"""simple docstring"""




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


										"""simple docstring"""





										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def 							__UpperCamelCase			(    self   : Union[str, Any]			,      _a   : Tuple			,      _a   : List[Any]="cpu"			,      _a   : str=torch.floataa			,      _a   : Any=0					)							->       Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=torch.Generator(device=_a					).manual_seed(_a					)
										_SCREAMING_SNAKE_CASE     					=np.random.RandomState(_a					).standard_normal((1, 4, 64, 64)					)
										_SCREAMING_SNAKE_CASE     					=torch.from_numpy(_a					).to(device=_a			,      dtype=_a					)
										_SCREAMING_SNAKE_CASE     					={
										    "prompt": "a photograph of an astronaut riding a horse",
										    "latents": latents,
										    "generator": generator,
										    "num_inference_steps": 3,
										    "guidance_scale": 7.5,
										    "output_type": "numpy",
										}
										return inputs




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''					)
										pipe.to(_a					)
										pipe.set_progress_bar_config(disable=_a					)

										_SCREAMING_SNAKE_CASE     					=self.get_inputs(_a					)
										_SCREAMING_SNAKE_CASE     					=pipe(**_a					).images
										_SCREAMING_SNAKE_CASE     					=image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 512, 512, 3)
										_SCREAMING_SNAKE_CASE     					=np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]					)
										assert np.abs(image_slice - expected_slice					).max() < 7E-3



import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
							import jax
							import jax.numpy as jnp

							from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor

class   A__							(						unittest.TestCase					):




						def __init__(    self   : List[Any]			,      _a   : Optional[Any]			,      _a   : int=3			,      _a   : int=32			,      _a   : str=3			,      _a   : List[str]=10			,      _a   : Optional[Any]=[10, 20, 30, 40]			,      _a   : Tuple=[1, 1, 2, 1]			,      _a   : Optional[Any]=True			,      _a   : Optional[int]=True			,      _a   : Any="relu"			,      _a   : Any=3			,      _a   : Tuple=None			,      )							->       Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=parent
										_SCREAMING_SNAKE_CASE     					=batch_size
										_SCREAMING_SNAKE_CASE     					=image_size
										_SCREAMING_SNAKE_CASE     					=num_channels
										_SCREAMING_SNAKE_CASE     					=embeddings_size
										_SCREAMING_SNAKE_CASE     					=hidden_sizes
										_SCREAMING_SNAKE_CASE     					=depths
										_SCREAMING_SNAKE_CASE     					=is_training
										_SCREAMING_SNAKE_CASE     					=use_labels
										_SCREAMING_SNAKE_CASE     					=hidden_act
										_SCREAMING_SNAKE_CASE     					=num_labels
										_SCREAMING_SNAKE_CASE     					=scope
										_SCREAMING_SNAKE_CASE     					=len(_a					)




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


										"""simple docstring"""





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

										_SCREAMING_SNAKE_CASE     					=self.get_config()

										return config, pixel_values




						def 							__UpperCamelCase			(    self   : List[str]					)							->       Dict:


										"""simple docstring"""





										return ResNetConfig(
										    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			,      image_size=self.image_size			,      )




						def 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : Union[str, Any]			,      _a   : List[Any]					)							->       Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetModel(config=_a					)
										_SCREAMING_SNAKE_CASE     					=model(_a					)

										# Output shape (b, c, h, w)
										self.parent.assertEqual(
										    result.last_hidden_state.shape			,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)			,      )




						def 							__UpperCamelCase			(    self   : int			,      _a   : str			,      _a   : str					)							->       Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.num_labels
										_SCREAMING_SNAKE_CASE     					=FlaxResNetForImageClassification(config=_a					)
										_SCREAMING_SNAKE_CASE     					=model(_a					)
										self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.num_labels)					)




						def 							__UpperCamelCase			(    self   : Tuple					)							->       Optional[Any]:


										"""simple docstring"""





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

@require_flax
class   A__							(						UpperCamelCase__				,       unittest.TestCase					):
						UpperCAmelCase     =  (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

						UpperCAmelCase     =  False
						UpperCAmelCase     =  False
						UpperCAmelCase     =  False




						def 							__UpperCamelCase			(    self   : str					)							->       None:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetModelTester(self					)
										_SCREAMING_SNAKE_CASE     					=ConfigTester(self			,      config_class=_a			,      has_text_modality=_a					)




						def 							__UpperCamelCase			(    self   : Optional[Any]					)							->       int:


										"""simple docstring"""





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




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


										"""simple docstring"""





										return




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_a					)




						def 							__UpperCamelCase			(    self   : Optional[int]					)							->       List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_a					)




						@unittest.skip(reason='''ResNet does not use inputs_embeds'''					)
						def 							__UpperCamelCase			(    self   : List[Any]					)							->       List[Any]:


										"""simple docstring"""





										pass




						@unittest.skip(reason='''ResNet does not support input and output embeddings'''					)
						def 							__UpperCamelCase			(    self   : Dict					)							->       Dict:


										"""simple docstring"""





										pass




						def 							__UpperCamelCase			(    self   : Any					)							->       int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs_for_common()

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

														_SCREAMING_SNAKE_CASE     					=['''pixel_values''']
														self.assertListEqual(arg_names[:1]			,      _a					)




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


										"""simple docstring"""





										def check_hidden_states_output(_a   : Optional[Any]			,      _a   : Tuple			,      _a   : Optional[Any]					):
														_SCREAMING_SNAKE_CASE     					=model_class(_a					)

														_SCREAMING_SNAKE_CASE     					=model(**self._prepare_for_class(_a			,      _a					)					)

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

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




						@unittest.skip(reason='''ResNet does not use feedforward chunking'''					)
						def 							__UpperCamelCase			(    self   : Optional[int]					)							->       Dict:


										"""simple docstring"""





										pass




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__					):
																		_SCREAMING_SNAKE_CASE     					=self._prepare_for_class(_a			,      _a					)
																		_SCREAMING_SNAKE_CASE     					=model_class(_a					)

																		@jax.jit
																		def model_jitted(_a   : Optional[Any]			,      **_a   : Any					):
																						return model(pixel_values=_a			,      **_a					)

																		with self.subTest('''JIT Enabled'''					):
																						_SCREAMING_SNAKE_CASE     					=model_jitted(**_a					).to_tuple()

																		with self.subTest('''JIT Disabled'''					):
																						with jax.disable_jit():
																										_SCREAMING_SNAKE_CASE     					=model_jitted(**_a					).to_tuple()

																		self.assertEqual(len(_a					)			,      len(_a					)					)
																		for jitted_output, output in zip(_a			,      _a					):
																						self.assertEqual(jitted_output.shape			,      output.shape					)







def   lowerCamelCase(				):
				_SCREAMING_SNAKE_CASE     					=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
				return image


@require_flax
class   A__							(						unittest.TestCase					):




						@cached_property
						def 							__UpperCamelCase			(    self   : List[Any]					)							->       Any:


										"""simple docstring"""





										return AutoImageProcessor.from_pretrained('''microsoft/resnet-50'''					) if is_vision_available() else None




						@slow
						def 							__UpperCamelCase			(    self   : Union[str, Any]					)							->       Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''					)

										_SCREAMING_SNAKE_CASE     					=self.default_image_processor
										_SCREAMING_SNAKE_CASE     					=prepare_img()
										_SCREAMING_SNAKE_CASE     					=image_processor(images=_a			,      return_tensors='''np'''					)

										_SCREAMING_SNAKE_CASE     					=model(**_a					)

										# verify the logits
										_SCREAMING_SNAKE_CASE     					=(1, 1000)
										self.assertEqual(outputs.logits.shape			,      _a					)

										_SCREAMING_SNAKE_CASE     					=jnp.array([-11.10_69, -9.78_77, -8.37_77]					)

										self.assertTrue(jnp.allclose(outputs.logits[0, :3]			,      _a			,      atol=1E-4					)					)