


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import 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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


_UpperCAmelCase	:     Optional[Any]   =False







@skip_mps
class  snake_case__(							snake_case_,  snake_case_,  snake_case_,  unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__					:     List[Any]    =							StableDiffusionAttendAndExcitePipeline
						SCREAMING_SNAKE_CASE__					:     int    =							False
						SCREAMING_SNAKE_CASE__					:     Optional[int]    =							TEXT_TO_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE__					:     Dict    =							TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""}							)
						SCREAMING_SNAKE_CASE__					:     str    =							TEXT_TO_IMAGE_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE__					:     str    =							TEXT_TO_IMAGE_IMAGE_PARAMS




						@classmethod
						def   lowercase_       ( cls							)    ->			List[str]:
										super().setUpClass()
										torch.use_deterministic_algorithms(UpperCamelCase__							)




						@classmethod
						def   lowercase_       ( cls							)    ->			Union[str, Any]:
										super().tearDownClass()
										torch.use_deterministic_algorithms(UpperCamelCase__							)




						def   lowercase_       ( self							)    ->			Tuple:
										torch.manual_seed(0							)
										lowerCAmelCase_       :					Any			 = UNetaDConditionModel(
										    block_out_channels=(3_2, 6_4)      ,  layers_per_block=1      ,  sample_size=3_2      ,  in_channels=4      ,  out_channels=4      ,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')      ,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,  cross_attention_dim=3_2      ,  attention_head_dim=(2, 4)      ,  use_linear_projection=UpperCamelCase__      ,  )
										lowerCAmelCase_       :					Optional[int]			 = DDIMScheduler(
										    beta_start=0.0_00_85      ,  beta_end=0.0_12      ,  beta_schedule='''scaled_linear'''      ,  clip_sample=UpperCamelCase__      ,  set_alpha_to_one=UpperCamelCase__      ,  )
										torch.manual_seed(0							)
										lowerCAmelCase_       :					Union[str, Any]			 = 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							)
										lowerCAmelCase_       :					Optional[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=5_1_2      ,  )
										lowerCAmelCase_       :					Dict			 = CLIPTextModel(UpperCamelCase__							)
										lowerCAmelCase_       :					List[Any]			 = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

										lowerCAmelCase_       :					str			 = {
										    "unet": unet,
										    "scheduler": scheduler,
										    "vae": vae,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "safety_checker": None,
										    "feature_extractor": None,
										}

										return components




						def   lowercase_       ( self      ,  __lowercase      ,  __lowercase=0							)    ->			str:




										if str(UpperCamelCase__							).startswith('''mps'''							):
														lowerCAmelCase_       :					int			 = torch.manual_seed(UpperCamelCase__							)
										else:
														lowerCAmelCase_       :					List[Any]			 = torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)
										lowerCAmelCase_       :					Optional[Any]			 = {
										    "prompt": "a cat and a frog",
										    "token_indices": [2, 5],
										    "generator": generator,
										    "num_inference_steps": 1,
										    "guidance_scale": 6.0,
										    "output_type": "numpy",
										    "max_iter_to_alter": 2,
										    "thresholds": {0: 0.7},
										}
										return inputs




						def   lowercase_       ( self							)    ->			str:
										lowerCAmelCase_       :					Union[str, Any]			 = "cpu"

										lowerCAmelCase_       :					List[str]			 = self.get_dummy_components()
										lowerCAmelCase_       :					Union[str, Any]			 = self.pipeline_class(**UpperCamelCase__							)
										pipe.to(UpperCamelCase__							)
										pipe.set_progress_bar_config(disable=UpperCamelCase__							)

										lowerCAmelCase_       :					Any			 = self.get_dummy_inputs(UpperCamelCase__							)
										lowerCAmelCase_       :					Dict			 = pipe(**UpperCamelCase__							).images
										lowerCAmelCase_       :					Dict			 = image[0, -3:, -3:, -1]

										self.assertEqual(image.shape      ,  (1, 6_4, 6_4, 3)							)
										lowerCAmelCase_       :					Any			 = np.array(
										    [0.63_90_53_64, 0.62_89_73_07, 0.48_59_90_17, 0.5_13_36_24, 0.5_55_00_48, 0.45_76_95_16, 0.50_32_69_73, 0.5_02_31_39, 0.45_38_44_96]							)
										lowerCAmelCase_       :					Optional[Any]			 = np.abs(image_slice.flatten() - expected_slice							).max()
										self.assertLessEqual(UpperCamelCase__      ,  1e-3							)




						def   lowercase_       ( self							)    ->			Optional[int]:
										super().test_cpu_offload_forward_pass(expected_max_diff=5e-4							)




						def   lowercase_       ( self							)    ->			List[str]:
										self._test_inference_batch_consistent(batch_sizes=[1, 2]							)




						def   lowercase_       ( self							)    ->			Optional[Any]:
										self._test_inference_batch_single_identical(batch_size=2      ,  expected_max_diff=7e-4							)




						def   lowercase_       ( self							)    ->			Optional[int]:
										super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3							)




						def   lowercase_       ( self							)    ->			Any:
										super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4							)




						def   lowercase_       ( self							)    ->			Tuple:
										super().test_save_load_local(expected_max_difference=5e-4							)




						def   lowercase_       ( self							)    ->			int:
										super().test_save_load_optional_components(expected_max_difference=4e-4							)







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







						'''simple docstring'''




						@classmethod
						def   lowercase_       ( cls							)    ->			List[Any]:
										super().setUpClass()
										torch.use_deterministic_algorithms(UpperCamelCase__							)




						@classmethod
						def   lowercase_       ( cls							)    ->			str:
										super().tearDownClass()
										torch.use_deterministic_algorithms(UpperCamelCase__							)




						def   lowercase_       ( self							)    ->			Optional[Any]:
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def   lowercase_       ( self							)    ->			Optional[int]:
										lowerCAmelCase_       :					int			 = torch.manual_seed(5_1							)

										lowerCAmelCase_       :					Any			 = StableDiffusionAttendAndExcitePipeline.from_pretrained(
										    '''CompVis/stable-diffusion-v1-4'''      ,  safety_checker=UpperCamelCase__      ,  torch_dtype=torch.floataa							)
										pipe.to('''cuda'''							)

										lowerCAmelCase_       :					Union[str, Any]			 = "a painting of an elephant with glasses"
										lowerCAmelCase_       :					Optional[Any]			 = [5, 7]

										lowerCAmelCase_       :					int			 = pipe(
										    prompt=UpperCamelCase__      ,  token_indices=UpperCamelCase__      ,  guidance_scale=7.5      ,  generator=UpperCamelCase__      ,  num_inference_steps=5      ,  max_iter_to_alter=5      ,  output_type='''numpy'''      ,  ).images[0]

										lowerCAmelCase_       :					Tuple			 = load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'''							)
										assert np.abs((expected_image - image).max()							) < 5e-1



"""simple docstring"""



import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch







class _lowerCAmelCase				(      unittest.TestCase	):
   __UpperCAmelCase :       Union[str, Any]							  =  JukeboxTokenizer
   __UpperCAmelCase :       Union[str, Any]							  =  {
       '''artist''': '''Zac Brown Band''',
       '''genres''': '''Country''',
       '''lyrics''': '''I met a traveller from an antique land,
        Who said "Two vast and trunkless legs of stone
        Stand in the desert. . . . Near them, on the sand,
        Half sunk a shattered visage lies, whose frown,
        And wrinkled lip, and sneer of cold command,
        Tell that its sculptor well those passions read
        Which yet survive, stamped on these lifeless things,
        The hand that mocked them, and the heart that fed;
        And on the pedestal, these words appear:
        My name is Ozymandias, King of Kings;
        Look on my Works, ye Mighty, and despair!
        Nothing beside remains. Round the decay
        Of that colossal Wreck, boundless and bare
        The lone and level sands stretch far away
        ''',
   }







   @require_torch
   def  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       import torch

       snake_case    :      Any      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-1b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      Optional[int]      					=    [
           torch.tensor([[
               0, 0, 0, 7169, 507, 9, 76, 39, 31, 46, 76, 27,
               76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
               44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
               47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
               76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
               30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 76,
               27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45,
               45, 76, 38, 31, 33, 45, 76, 41, 32, 76, 45, 46,
               41, 40, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
               76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
               76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
               64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
               30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
               27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
               34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
               27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
               41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
               44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
               76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
               32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
               40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
               45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
               31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
               45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
               34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
               31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
               40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
               38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
               76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
               41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
               27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
               46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
               41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
               46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
               41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
               40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
               27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
               76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
               41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
               76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
               27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
               34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
               44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
               40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
               46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
               38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
               40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
               27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
               76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
               76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
               76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
               76, 76]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)







   @require_torch
   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       import torch

       snake_case    :      Tuple      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-5b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      List[Any]      					=    [
           torch.tensor([[
               0, 0, 0, 1069, 11, -1, -1, -1, -1, 9, 77, 39,
               31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
               31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
               40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
               79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
               77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
               27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
               37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
               32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
               77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
               77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
               46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
               77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
               77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
               77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
               77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
               64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
               40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
               40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
               38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
               31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
               41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
               46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
               41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
               31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
               31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
               44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
               31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
               38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
               40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
               27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
               31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
               34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
               31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
               31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
               45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
               31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
               15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
               11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
               45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
               41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
               44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
               46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
               27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
               35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
               77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
               31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
               41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
               77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
               40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
               77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
               27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
               77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
               77, 77, 77, 77, 77, 77]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)





