




'''simple docstring'''





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     lowercase  ( __lowerCAmelCase					,      __lowerCAmelCase					,      unittest.TestCase):
      __lowercase					=	StableDiffusionXLPipeline
      __lowercase					=	TEXT_TO_IMAGE_PARAMS
      __lowercase					=	TEXT_TO_IMAGE_BATCH_PARAMS
      __lowercase					=	TEXT_TO_IMAGE_IMAGE_PARAMS
      __lowercase					=	TEXT_TO_IMAGE_IMAGE_PARAMS






      def  A__(self       )							-> Optional[Any]:
         torch.manual_seed(0       )
         SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)			,	layers_per_block=2			,	sample_size=3_2			,	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=8_0			,	cross_attention_dim=6_4			,	)
         SCREAMING_SNAKE_CASE__      :					int							    =			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__      :					Optional[int]							    =			AutoencoderKL(
             block_out_channels=[3_2, 6_4]			,	in_channels=3			,	out_channels=3			,	down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']			,	up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']			,	latent_channels=4			,	sample_size=1_2_8			,	)
         torch.manual_seed(0       )
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			CLIPTextConfig(
             bos_token_id=0			,	eos_token_id=2			,	hidden_size=3_2			,	intermediate_size=3_7			,	layer_norm_eps=1E-05			,	num_attention_heads=4			,	num_hidden_layers=5			,	pad_token_id=1			,	vocab_size=1_0_0_0			,	hidden_act='''gelu'''			,	projection_dim=3_2			,	)
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			CLIPTextModel(__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					Dict							    =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			,	local_files_only=__UpperCamelCase       )

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

         SCREAMING_SNAKE_CASE__      :					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  A__(self			,	a__			,	a__=0       )							-> str:






         if str(__UpperCamelCase       ).startswith('''mps'''       ):
            SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			torch.manual_seed(__UpperCamelCase       )
         else:
            SCREAMING_SNAKE_CASE__      :					List[str]							    =			torch.Generator(device=__UpperCamelCase       ).manual_seed(__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			{
             '''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  A__(self       )							-> Tuple:
         SCREAMING_SNAKE_CASE__      :					Tuple							    =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
         SCREAMING_SNAKE_CASE__      :					Any							    =			self.get_dummy_components()
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			StableDiffusionXLPipeline(**__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					Any							    =			sd_pipe.to(__UpperCamelCase       )
         sd_pipe.set_progress_bar_config(disable=__UpperCamelCase       )

         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			self.get_dummy_inputs(__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			sd_pipe(**__UpperCamelCase       ).images
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 6_4, 6_4, 3)
         SCREAMING_SNAKE_CASE__      :					Tuple							    =			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  A__(self       )							-> Any:
         SCREAMING_SNAKE_CASE__      :					List[str]							    =			self.get_dummy_components()
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			StableDiffusionXLPipeline(**__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					str							    =			sd_pipe.to(__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					str							    =			sd_pipe.to(__UpperCamelCase       )
         sd_pipe.set_progress_bar_config(disable=__UpperCamelCase       )

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

         SCREAMING_SNAKE_CASE__      :					Tuple							    =			sd_pipe(**__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			output.images[0, -3:, -3:, -1]

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

         (
             (
             SCREAMING_SNAKE_CASE__			
         )    ,					(
             SCREAMING_SNAKE_CASE__			
         )    ,					(
             SCREAMING_SNAKE_CASE__			
         )    ,					(
             SCREAMING_SNAKE_CASE__			
         )    ,								
         )      :					List[Any]							    =			sd_pipe.encode_prompt(__UpperCamelCase			,	negative_prompt=__UpperCamelCase       )

         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			sd_pipe(
             **__UpperCamelCase			,	prompt_embeds=__UpperCamelCase			,	negative_prompt_embeds=__UpperCamelCase			,	pooled_prompt_embeds=__UpperCamelCase			,	negative_pooled_prompt_embeds=__UpperCamelCase			,	)
         SCREAMING_SNAKE_CASE__      :					Any							    =			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  A__(self       )							-> Optional[Any]:
         super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )






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






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






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






      def  A__(self			,	a__			,	a__="cpu"			,	a__=torch.floataa			,	a__=0       )							-> Optional[Any]:
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			torch.Generator(device=__UpperCamelCase       ).manual_seed(__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			np.random.RandomState(__UpperCamelCase       ).standard_normal((1, 4, 6_4, 6_4)       )
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			torch.from_numpy(__UpperCamelCase       ).to(device=__UpperCamelCase			,	dtype=__UpperCamelCase       )
         SCREAMING_SNAKE_CASE__      :					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  A__(self       )							-> List[str]:
         SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''       )
         pipe.to(__UpperCamelCase       )
         pipe.set_progress_bar_config(disable=__UpperCamelCase       )

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         SCREAMING_SNAKE_CASE__      :					int							    =			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







'''simple docstring'''

import os
import unittest

import pytest
from parameterized import parameterized

from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device


# skipping in unittest tests

a__         =     [(1,)]




def    snake_case__     (  )			->       Tuple:





			'''simple docstring'''






			snake_case__          =						bool(os.getenv("""RUN_SLOW"""       ,      0	)	)
			if run_slow:
						assert True
			else:
						assert False, "should have been skipped"




def    snake_case__     (  )			->       Dict:





			'''simple docstring'''






			snake_case__          =						bool(os.getenv("""RUN_SLOW"""       ,      0	)	)
			if run_slow and torch_device == "cuda":
						assert True
			else:
						assert False, "should have been skipped"





@require_torch
class __magic_name__(						unittest.TestCase      ):






						@slow
						@require_torch_gpu
						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									check_slow_torch_cuda()






						@require_torch_gpu
						@slow
						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									check_slow_torch_cuda()







						@parameterized.expand(__UpperCamelCase       )
						@slow
						def       __lowerCAmelCase(   self						:						Tuple					,   __UpperCamelCase						:						int=None       ):



									'''simple docstring'''






									check_slow()







@slow
@require_torch_gpu
def    snake_case__     (  )			->       Tuple:





			'''simple docstring'''






			check_slow_torch_cuda()




@require_torch_gpu
@slow
def    snake_case__     (  )			->       List[Any]:





			'''simple docstring'''






			check_slow_torch_cuda()




@slow
@pytest.mark.parametrize("""param"""       ,      [1]	)
def    snake_case__     (  a	)			->       Union[str, Any]:





			'''simple docstring'''






			check_slow()




@pytest.mark.parametrize("""param"""       ,      [1]	)
@slow
def    snake_case__     (  a	)			->       Dict:





			'''simple docstring'''






			check_slow()