


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__(    snake_case__    ,    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Optional[int]			= StableDiffusionXLImgaImgPipeline
    A_  :     Optional[int]			= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    A_  :     Dict			= PipelineTesterMixin.required_optional_params - {'latents'}
    A_  :     List[str]			= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    A_  :     Tuple			= IMAGE_TO_IMAGE_IMAGE_PARAMS
    A_  :     Optional[int]			= IMAGE_TO_IMAGE_IMAGE_PARAMS





    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       torch.manual_seed(0	)
       UpperCAmelCase_							:      Optional[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=__snake_case       ,   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       ,   )
       UpperCAmelCase_							:      List[str]						=   EulerDiscreteScheduler(
           beta_start=0.00_085       ,   beta_end=0.012       ,   steps_offset=1       ,   beta_schedule='''scaled_linear'''       ,   timestep_spacing='''leading'''       ,   )
       torch.manual_seed(0	)
       UpperCAmelCase_							:      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	)
       UpperCAmelCase_							:      Dict						=   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=1_000       ,   hidden_act='''gelu'''       ,   projection_dim=32       ,   )
       UpperCAmelCase_							:      Any						=   CLIPTextModel(__snake_case	)
       UpperCAmelCase_							:      Dict						=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,   local_files_only=__snake_case	)

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

       UpperCAmelCase_							:      Tuple						=   {
           '''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				:   Union[str, Any]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   Any=0	):





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





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase_							:      List[Any]						=   self.get_dummy_components()
       UpperCAmelCase_							:      Optional[Any]						=   StableDiffusionXLImgaImgPipeline(**__snake_case	)
       UpperCAmelCase_							:      Dict						=   sd_pipe.to(__snake_case	)
       sd_pipe.set_progress_bar_config(disable=__snake_case	)

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

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

       UpperCAmelCase_							:      Union[str, 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				:   str	):





       '''simple docstring'''
       super().test_attention_slicing_forward_pass(expected_max_diff=3E-3	)





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       super().test_inference_batch_single_identical(expected_max_diff=3E-3	)





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       pass





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   self.get_dummy_components()
       UpperCAmelCase_							:      List[str]						=   StableDiffusionXLImgaImgPipeline(**__snake_case	)
       UpperCAmelCase_							:      List[str]						=   sd_pipe.to(__snake_case	)
       UpperCAmelCase_							:      List[str]						=   sd_pipe.to(__snake_case	)
       sd_pipe.set_progress_bar_config(disable=__snake_case	)

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

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

       # forward with prompt embeds
       UpperCAmelCase_							:      Optional[int]						=   self.get_dummy_inputs(__snake_case	)
       UpperCAmelCase_							:      str						=   3 * ['''this is a negative prompt''']
       UpperCAmelCase_							:      str						=   3 * [inputs.pop('''prompt'''	)]

       (
           (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      (
           UpperCAmelCase_
       )		,      
       )							:      Dict						=   sd_pipe.encode_prompt(__snake_case       ,   negative_prompt=__snake_case	)

       UpperCAmelCase_							:      Any						=   sd_pipe(
           **__snake_case       ,   prompt_embeds=__snake_case       ,   negative_prompt_embeds=__snake_case       ,   pooled_prompt_embeds=__snake_case       ,   negative_pooled_prompt_embeds=__snake_case       ,   )
       UpperCAmelCase_							:      Tuple						=   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 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





    def 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   Tuple       ,   __snake_case				:   List[str]="cpu"       ,   __snake_case				:   List[Any]=torch.floataa       ,   __snake_case				:   Tuple=0	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
       UpperCAmelCase_							:      Tuple						=   np.random.RandomState(__snake_case	).standard_normal((1, 4, 64, 64)	)
       UpperCAmelCase_							:      List[str]						=   torch.from_numpy(__snake_case	).to(device=__snake_case       ,   dtype=__snake_case	)
       UpperCAmelCase_							:      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 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      str						=   DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''	)
       pipe.to(__snake_case	)
       pipe.set_progress_bar_config(disable=__snake_case	)

       UpperCAmelCase_							:      Dict						=   self.get_inputs(__snake_case	)
       UpperCAmelCase_							:      Dict						=   pipe(**__snake_case	).images
       UpperCAmelCase_							:      Optional[Any]						=   image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase_							:      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



from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase				:  str      					=							logging.get_logger(__name__)

__UpperCamelCase				:  Any      					=							{
    'tiiuae/falcon-40b': 'https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json',
    'tiiuae/falcon-7b': 'https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json',
}





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     int			= 'falcon'
    A_  :     int			= ['past_key_values']





    def __init__(	self				:   Optional[Any]       ,   __snake_case				:   Tuple=65_024       ,   __snake_case				:   List[str]=4_544       ,   __snake_case				:   Optional[Any]=32       ,   __snake_case				:   Any=71       ,   __snake_case				:   str=1E-5       ,   __snake_case				:   List[str]=0.02       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   Dict=0.0       ,   __snake_case				:   Optional[Any]=0.0       ,   __snake_case				:   Any=None       ,   __snake_case				:   List[Any]=False       ,   __snake_case				:   Dict=False       ,   __snake_case				:   Optional[int]=True       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   Optional[Any]=False       ,   __snake_case				:   Dict=11       ,   __snake_case				:   List[str]=11       ,   **__snake_case				:   int       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   vocab_size
       # Backward compatibility with n_embed kwarg
       UpperCAmelCase_							:      Union[str, Any]						=   kwargs.pop('''n_embed'''       ,   __snake_case	)
       UpperCAmelCase_							:      str						=   hidden_size if n_embed is None else n_embed
       UpperCAmelCase_							:      Tuple						=   num_hidden_layers
       UpperCAmelCase_							:      Tuple						=   num_attention_heads
       UpperCAmelCase_							:      Optional[int]						=   layer_norm_epsilon
       UpperCAmelCase_							:      int						=   initializer_range
       UpperCAmelCase_							:      Optional[int]						=   use_cache
       UpperCAmelCase_							:      List[Any]						=   hidden_dropout
       UpperCAmelCase_							:      Any						=   attention_dropout

       UpperCAmelCase_							:      Tuple						=   bos_token_id
       UpperCAmelCase_							:      List[Any]						=   eos_token_id
       UpperCAmelCase_							:      Any						=   num_attention_heads if num_kv_heads is None else num_kv_heads
       UpperCAmelCase_							:      Optional[int]						=   alibi
       UpperCAmelCase_							:      Dict						=   new_decoder_architecture
       UpperCAmelCase_							:      List[Any]						=   multi_query  # Ignored when new_decoder_architecture is True
       UpperCAmelCase_							:      Tuple						=   parallel_attn
       UpperCAmelCase_							:      List[Any]						=   bias

       super().__init__(bos_token_id=__snake_case       ,   eos_token_id=__snake_case       ,   **__snake_case	)





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return self.hidden_size // self.num_attention_heads





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return not self.alibi