




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    lowerCamelCase    (		_lowerCamelCase				,_lowerCamelCase				,unittest.TestCase						):






			'''simple docstring'''





			UpperCamelCase__								=StableDiffusionXLImgaImgPipeline
			UpperCamelCase__								=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
			UpperCamelCase__								=PipelineTesterMixin.required_optional_params - {'''latents'''}
			UpperCamelCase__								=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			UpperCamelCase__								=IMAGE_TO_IMAGE_IMAGE_PARAMS
			UpperCamelCase__								=IMAGE_TO_IMAGE_IMAGE_PARAMS

			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Optional[int]:
								torch.manual_seed(0  )
								__magic_name__							:	Union[str, Any]    							=							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=lowerCamelCase_       ,						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       ,						)
								__magic_name__							:	int    							=							EulerDiscreteScheduler(
								    beta_start=0.0_0_0_8_5       ,						beta_end=0.0_1_2       ,						steps_offset=1       ,						beta_schedule='''scaled_linear'''       ,						timestep_spacing='''leading'''       ,						)
								torch.manual_seed(0  )
								__magic_name__							:	Tuple    							=							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  )
								__magic_name__							:	List[Any]    							=							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       ,						)
								__magic_name__							:	Union[str, Any]    							=							CLIPTextModel(lowerCamelCase_  )
								__magic_name__							:	Dict    							=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,						local_files_only=lowerCamelCase_  )

								__magic_name__							:	List[Any]    							=							CLIPTextModelWithProjection(lowerCamelCase_  )
								__magic_name__							:	List[Any]    							=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,						local_files_only=lowerCamelCase_  )

								__magic_name__							:	Union[str, Any]    							=							{
								    '''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]       ,						lowerCamelCase_			:   Tuple       ,						lowerCamelCase_			:   Union[str, Any]=0  )						->   Optional[Any]:
								__magic_name__							:	int    							=							floats_tensor((1, 3, 32, 32)       ,						rng=random.Random(lowerCamelCase_  )  ).to(lowerCamelCase_  )
								__magic_name__							:	Tuple    							=							image / 2 + 0.5
								if str(lowerCamelCase_  ).startswith('''mps'''  ):
													__magic_name__							:	List[Any]    							=							torch.manual_seed(lowerCamelCase_  )
								else:
													__magic_name__							:	List[Any]    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
								__magic_name__							:	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.7_5,
								}
								return inputs

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   str:
								__magic_name__							:	Optional[int]    							=							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								__magic_name__							:	Tuple    							=							self.get_dummy_components()
								__magic_name__							:	List[Any]    							=							StableDiffusionXLImgaImgPipeline(**lowerCamelCase_  )
								__magic_name__							:	int    							=							sd_pipe.to(lowerCamelCase_  )
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

								__magic_name__							:	Tuple    							=							self.get_dummy_inputs(lowerCamelCase_  )
								__magic_name__							:	Tuple    							=							sd_pipe(**lowerCamelCase_  ).images
								__magic_name__							:	List[Any]    							=							image[0, -3:, -3:, -1]

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

								__magic_name__							:	int    							=							np.array([0.4_6_5_6, 0.4_8_4_0, 0.4_4_3_9, 0.6_6_9_8, 0.5_5_7_4, 0.4_5_2_4, 0.5_7_9_9, 0.5_9_4_3, 0.5_1_6_5]  )

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

			def 							UpperCAmelCase__       ( self			:   Tuple  )						->   Any:
								super().test_attention_slicing_forward_pass(expected_max_diff=3E-3  )

			def 							UpperCAmelCase__       ( self			:   Dict  )						->   List[str]:
								super().test_inference_batch_single_identical(expected_max_diff=3E-3  )

			def 							UpperCAmelCase__       ( self			:   Any  )						->   int:
								pass


			def 							UpperCAmelCase__       ( self			:   Tuple  )						->   List[Any]:
								__magic_name__							:	Dict    							=							self.get_dummy_components()
								__magic_name__							:	Dict    							=							StableDiffusionXLImgaImgPipeline(**lowerCamelCase_  )
								__magic_name__							:	Tuple    							=							sd_pipe.to(lowerCamelCase_  )
								__magic_name__							:	Optional[int]    							=							sd_pipe.to(lowerCamelCase_  )
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

								# forward without prompt embeds
								__magic_name__							:	Tuple    							=							self.get_dummy_inputs(lowerCamelCase_  )
								__magic_name__							:	Any    							=							3 * ['''this is a negative prompt''']
								__magic_name__							:	Tuple    							=							negative_prompt
								__magic_name__							:	Optional[Any]    							=							3 * [inputs['''prompt''']]

								__magic_name__							:	Optional[Any]    							=							sd_pipe(**lowerCamelCase_  )
								__magic_name__							:	Optional[Any]    							=							output.images[0, -3:, -3:, -1]

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

								(
								    __magic_name__    
								)							:	List[str]    							=							sd_pipe.encode_prompt(lowerCamelCase_       ,						negative_prompt=lowerCamelCase_  )

								__magic_name__							:	str    							=							sd_pipe(
								    **lowerCamelCase_       ,						prompt_embeds=lowerCamelCase_       ,						negative_prompt_embeds=lowerCamelCase_       ,						pooled_prompt_embeds=lowerCamelCase_       ,						negative_pooled_prompt_embeds=lowerCamelCase_       ,						)
								__magic_name__							:	Dict    							=							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    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			def 							UpperCAmelCase__       ( self			:   Dict  )						->   Any:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()

			def 							UpperCAmelCase__       ( self			:   Tuple       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Optional[Any]="cpu"       ,						lowerCamelCase_			:   Any=torch.floataa       ,						lowerCamelCase_			:   List[Any]=0  )						->   List[str]:
								__magic_name__							:	Dict    							=							torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
								__magic_name__							:	Optional[Any]    							=							np.random.RandomState(lowerCamelCase_  ).standard_normal((1, 4, 64, 64)  )
								__magic_name__							:	Optional[Any]    							=							torch.from_numpy(lowerCamelCase_  ).to(device=lowerCamelCase_       ,						dtype=lowerCamelCase_  )
								__magic_name__							:	Dict    							=							{
								    '''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			:   List[Any]  )						->   List[Any]:
								__magic_name__							:	str    							=							DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''  )
								pipe.to(lowerCamelCase_  )
								pipe.set_progress_bar_config(disable=lowerCamelCase_  )

								__magic_name__							:	Tuple    							=							self.get_inputs(lowerCamelCase_  )
								__magic_name__							:	List[Any]    							=							pipe(**lowerCamelCase_  ).images
								__magic_name__							:	Dict    							=							image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								__magic_name__							:	Tuple    							=							np.array([0.4_9_4_9_3, 0.4_7_8_9_6, 0.4_0_7_9_8, 0.5_4_2_1_4, 0.5_3_2_1_2, 0.4_8_2_0_2, 0.4_7_6_5_6, 0.4_6_3_2_9, 0.4_8_5_0_6]  )
								assert np.abs(image_slice - expected_slice  ).max() < 7E-3









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase   :						List[Any]				=			{
    '''configuration_speech_to_text''': ['''SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Speech2TextConfig'''],
    '''processing_speech_to_text''': ['''Speech2TextProcessor'''],
}

try:
	if not is_sentencepiece_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			['''Speech2TextTokenizer''']

try:
	if not is_speech_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Optional[Any]				=			['''Speech2TextFeatureExtractor''']

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFSpeech2TextForConditionalGeneration''',
	    '''TFSpeech2TextModel''',
	    '''TFSpeech2TextPreTrainedModel''',
	]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''Speech2TextForConditionalGeneration''',
	    '''Speech2TextModel''',
	    '''Speech2TextPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
	from .processing_speech_to_text import SpeechaTextProcessor

	try:
		if not is_sentencepiece_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_speech_to_text import SpeechaTextTokenizer

	try:
		if not is_speech_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_speech_to_text import (
		    TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFSpeechaTextForConditionalGeneration,
		    TFSpeechaTextModel,
		    TFSpeechaTextPreTrainedModel,
		)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_speech_to_text import (
		    SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    SpeechaTextForConditionalGeneration,
		    SpeechaTextModel,
		    SpeechaTextPreTrainedModel,
		)

else:
	import sys

	__lowerCamelCase   :						List[str]				=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



