
'''simple docstring'''
import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






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







		"""simple docstring"""

		lowerCamelCase     							=						StableDiffusionXLImgaImgPipeline
		lowerCamelCase     							=						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
		lowerCamelCase     							=						PipelineTesterMixin.required_optional_params - {"""latents"""}
		lowerCamelCase     							=						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
		lowerCamelCase     							=						IMAGE_TO_IMAGE_IMAGE_PARAMS
		lowerCamelCase     							=						IMAGE_TO_IMAGE_IMAGE_PARAMS



		def       lowerCAmelCase				(      self				:			Tuple      )						->			Dict:






				"""simple docstring"""







				torch.manual_seed(0      )
				snake_case  :       List[str]     			=						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=UpperCamelCase__  ,	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  ,	)
				snake_case  :       Any     			=						EulerDiscreteScheduler(
				    beta_start=0.00_085  ,	beta_end=0.012  ,	steps_offset=1  ,	beta_schedule='''scaled_linear'''  ,	timestep_spacing='''leading'''  ,	)
				torch.manual_seed(0      )
				snake_case  :       Optional[Any]     			=						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      )
				snake_case  :       Optional[Any]     			=						CLIPTextConfig(
				    bos_token_id=0  ,	eos_token_id=2  ,	hidden_size=32  ,	intermediate_size=37  ,	layer_norm_eps=1e-0_5  ,	num_attention_heads=4  ,	num_hidden_layers=5  ,	pad_token_id=1  ,	vocab_size=1000  ,	hidden_act='''gelu'''  ,	projection_dim=32  ,	)
				snake_case  :       Tuple     			=						CLIPTextModel(UpperCamelCase__      )
				snake_case  :       Union[str, Any]     			=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  ,	local_files_only=UpperCamelCase__      )

				snake_case  :       str     			=						CLIPTextModelWithProjection(UpperCamelCase__      )
				snake_case  :       Union[str, Any]     			=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  ,	local_files_only=UpperCamelCase__      )

				snake_case  :       Optional[int]     			=						{
				    '''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       lowerCAmelCase				(      self				:			Optional[Any]  ,	UpperCamelCase__				:			Dict  ,	UpperCamelCase__				:			Union[str, Any]=0      )						->			Tuple:






				"""simple docstring"""







				snake_case  :       Optional[int]     			=						floats_tensor((1, 3, 32, 32)  ,	rng=random.Random(UpperCamelCase__      )      ).to(UpperCamelCase__      )
				snake_case  :       Any     			=						image / 2 + 0.5
				if str(UpperCamelCase__      ).startswith('''mps'''      ):
						snake_case  :       str     			=						torch.manual_seed(UpperCamelCase__      )
				else:
						snake_case  :       Optional[int]     			=						torch.Generator(device=UpperCamelCase__      ).manual_seed(UpperCamelCase__      )
				snake_case  :       List[str]     			=						{
				    '''prompt''': '''A painting of a squirrel eating a burger''',
				    '''image''': image,
				    '''generator''': generator,
				    '''num_inference_steps''': 2,
				    '''guidance_scale''': 5.0,
				    '''output_type''': '''numpy''',
				    '''strength''': 0.75,
				}
				return inputs



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






				"""simple docstring"""







				snake_case  :       Dict     			=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
				snake_case  :       Optional[Any]     			=						self.get_dummy_components()
				snake_case  :       Optional[Any]     			=						StableDiffusionXLImgaImgPipeline(**UpperCamelCase__      )
				snake_case  :       Optional[int]     			=						sd_pipe.to(UpperCamelCase__      )
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

				snake_case  :       Optional[int]     			=						self.get_dummy_inputs(UpperCamelCase__      )
				snake_case  :       str     			=						sd_pipe(**UpperCamelCase__      ).images
				snake_case  :       Union[str, Any]     			=						image[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)

				snake_case  :       Optional[Any]     			=						np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]      )

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



		def       lowerCAmelCase				(      self				:			Optional[Any]      )						->			str:






				"""simple docstring"""







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



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






				"""simple docstring"""







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



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






				"""simple docstring"""







				pass



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






				"""simple docstring"""







				snake_case  :       List[Any]     			=						self.get_dummy_components()
				snake_case  :       str     			=						StableDiffusionXLImgaImgPipeline(**UpperCamelCase__      )
				snake_case  :       str     			=						sd_pipe.to(UpperCamelCase__      )
				snake_case  :       int     			=						sd_pipe.to(UpperCamelCase__      )
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

				# forward without prompt embeds
				snake_case  :       Tuple     			=						self.get_dummy_inputs(UpperCamelCase__      )
				snake_case  :       Tuple     			=						3 * ['''this is a negative prompt''']
				snake_case  :       Union[str, Any]     			=						negative_prompt
				snake_case  :       Tuple     			=						3 * [inputs['''prompt''']]

				snake_case  :       int     			=						sd_pipe(**UpperCamelCase__      )
				snake_case  :       List[Any]     			=						output.images[0, -3:, -3:, -1]

				# forward with prompt embeds
				snake_case  :       Dict     			=						self.get_dummy_inputs(UpperCamelCase__      )
				snake_case  :       Dict     			=						3 * ['''this is a negative prompt''']
				snake_case  :       List[str]     			=						3 * [inputs.pop('''prompt'''      )]

				(
				    (
				    snake_case
				)			,(
				    snake_case
				)			,(
				    snake_case
				)			,(
				    snake_case
				)			,
				)  :       Union[str, Any]     			=						sd_pipe.encode_prompt(UpperCamelCase__  ,	negative_prompt=UpperCamelCase__      )

				snake_case  :       List[Any]     			=						sd_pipe(
				    **UpperCamelCase__  ,	prompt_embeds=UpperCamelCase__  ,	negative_prompt_embeds=UpperCamelCase__  ,	pooled_prompt_embeds=UpperCamelCase__  ,	negative_pooled_prompt_embeds=UpperCamelCase__  ,	)
				snake_case  :       str     			=						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







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







		"""simple docstring"""

		def       lowerCAmelCase				(      self				:			str      )						->			Optional[Any]:






				"""simple docstring"""







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



		def       lowerCAmelCase				(      self				:			Any  ,	UpperCamelCase__				:			List[str]  ,	UpperCamelCase__				:			Tuple="cpu"  ,	UpperCamelCase__				:			str=torch.floataa  ,	UpperCamelCase__				:			Optional[int]=0      )						->			Optional[int]:






				"""simple docstring"""







				snake_case  :       List[Any]     			=						torch.Generator(device=UpperCamelCase__      ).manual_seed(UpperCamelCase__      )
				snake_case  :       Optional[Any]     			=						np.random.RandomState(UpperCamelCase__      ).standard_normal((1, 4, 64, 64)      )
				snake_case  :       Tuple     			=						torch.from_numpy(UpperCamelCase__      ).to(device=UpperCamelCase__  ,	dtype=UpperCamelCase__      )
				snake_case  :       List[Any]     			=						{
				    '''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       lowerCAmelCase				(      self				:			Optional[Any]      )						->			Union[str, Any]:






				"""simple docstring"""







				snake_case  :       List[str]     			=						DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''      )
				pipe.to(UpperCamelCase__      )
				pipe.set_progress_bar_config(disable=UpperCamelCase__      )

				snake_case  :       str     			=						self.get_inputs(UpperCamelCase__      )
				snake_case  :       List[Any]     			=						pipe(**UpperCamelCase__      ).images
				snake_case  :       Union[str, Any]     			=						image[0, -3:, -3:, -1].flatten()

				assert image.shape == (1, 512, 512, 3)
				snake_case  :       int     			=						np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]      )
				assert np.abs(image_slice - expected_slice      ).max() < 7e-3




'''simple docstring'''
from .data_collator import (
    DataCollatorForLanguageModeling,
    DataCollatorForPermutationLanguageModeling,
    DataCollatorForSeqaSeq,
    DataCollatorForSOP,
    DataCollatorForTokenClassification,
    DataCollatorForWholeWordMask,
    DataCollatorWithPadding,
    DefaultDataCollator,
    default_data_collator,
)
from .metrics import glue_compute_metrics, xnli_compute_metrics
from .processors import (
    DataProcessor,
    InputExample,
    InputFeatures,
    SingleSentenceClassificationProcessor,
    SquadExample,
    SquadFeatures,
    SquadVaProcessor,
    SquadVaProcessor,
    glue_convert_examples_to_features,
    glue_output_modes,
    glue_processors,
    glue_tasks_num_labels,
    squad_convert_examples_to_features,
    xnli_output_modes,
    xnli_processors,
    xnli_tasks_num_labels,
)


