





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 	UpperCamelCase_							(  __lowerCamelCase					,	__lowerCamelCase					,	unittest.TestCase     ):
							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   _snake_case	(				self     :Optional[Any]		) -> str:



									"""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=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				,					)
									SCREAMING_SNAKE_CASE__								=   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		)
									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(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				,					local_files_only=lowerCamelCase__		)

									SCREAMING_SNAKE_CASE__								=   CLIPTextModelWithProjection(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				,					local_files_only=lowerCamelCase__		)

									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   _snake_case	(				self     :int				,					__A     :int				,					__A     :Dict=0		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   floats_tensor((1, 3, 32, 32)				,					rng=random.Random(lowerCamelCase__		)		).to(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   image / 2 + 0.5
									if str(lowerCamelCase__		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__								=   torch.manual_seed(lowerCamelCase__		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.Generator(device=lowerCamelCase__		).manual_seed(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   {
									    """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   _snake_case	(				self     :int		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE__								=   self.get_dummy_components()
									SCREAMING_SNAKE_CASE__								=   StableDiffusionXLImgaImgPipeline(**lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(lowerCamelCase__		)
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

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

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

									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :List[str]		) -> Union[str, Any]:



									"""simple docstring"""



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


							def   _snake_case	(				self     :Any		) -> str:



									"""simple docstring"""



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


							def   _snake_case	(				self     :List[Any]		) -> int:



									"""simple docstring"""



									pass


							def   _snake_case	(				self     :Any		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.get_dummy_components()
									SCREAMING_SNAKE_CASE__								=   StableDiffusionXLImgaImgPipeline(**lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(lowerCamelCase__		)
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__		)

									# forward without prompt embeds
									SCREAMING_SNAKE_CASE__								=   self.get_dummy_inputs(lowerCamelCase__		)
									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(**lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   output.images[0, -3:, -3:, -1]

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

									(
									    (
									    SCREAMING_SNAKE_CASE__
									)						,    (
									    SCREAMING_SNAKE_CASE__
									)						,    (
									    SCREAMING_SNAKE_CASE__
									)						,    (
									    SCREAMING_SNAKE_CASE__
									)						,    
									)								=   sd_pipe.encode_prompt(lowerCamelCase__				,					negative_prompt=lowerCamelCase__		)

									SCREAMING_SNAKE_CASE__								=   sd_pipe(
									    **lowerCamelCase__				,					prompt_embeds=lowerCamelCase__				,					negative_prompt_embeds=lowerCamelCase__				,					pooled_prompt_embeds=lowerCamelCase__				,					negative_pooled_prompt_embeds=lowerCamelCase__				,					)
									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




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


							def   _snake_case	(				self     :Any		) -> Any:



									"""simple docstring"""



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


							def   _snake_case	(				self     :str				,					__A     :Dict				,					__A     :Optional[int]="cpu"				,					__A     :Dict=torch.floataa				,					__A     :List[Any]=0		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=lowerCamelCase__		).manual_seed(lowerCamelCase__		)
									SCREAMING_SNAKE_CASE__								=   np.random.RandomState(lowerCamelCase__		).standard_normal((1, 4, 64, 64)		)
									SCREAMING_SNAKE_CASE__								=   torch.from_numpy(lowerCamelCase__		).to(device=lowerCamelCase__				,					dtype=lowerCamelCase__		)
									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   _snake_case	(				self     :List[Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""		)
									pipe.to(lowerCamelCase__		)
									pipe.set_progress_bar_config(disable=lowerCamelCase__		)

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

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   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 __future__ import annotations







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = None		,				SCREAMING_SNAKE_CASE = None   ):
      '''simple docstring'''







      if start is None:
            A_							       =		0

      if end is None:
            A_							       =		len(SCREAMING_SNAKE_CASE   ) - 1

      if start >= end:
            return

      A_							       =		(start + end) // 2

      slowsort(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE   )
      slowsort(SCREAMING_SNAKE_CASE		,				mid + 1		,				SCREAMING_SNAKE_CASE   )

      if sequence[end] < sequence[mid]:
            A_       ,A_							       =		sequence[mid], sequence[end]

      slowsort(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				end - 1   )


if __name__ == "__main__":
      from doctest import testmod

      testmod()





