



import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

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,
    assert_mean_pixel_difference,
)


enable_full_determinism()



class        lowerCAmelCase__					(  _lowerCamelCase ,				_lowerCamelCase ,				_lowerCamelCase ,				unittest.TestCase       ):
				A_   :     Union[str, Any]								=     StableUnCLIPPipeline
				A_   :     Union[str, Any]								=     TEXT_TO_IMAGE_PARAMS
				A_   :     Union[str, Any]								=     TEXT_TO_IMAGE_BATCH_PARAMS
				A_   :     Optional[Any]								=     TEXT_TO_IMAGE_IMAGE_PARAMS
				A_   :     str								=     TEXT_TO_IMAGE_IMAGE_PARAMS

				# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
				A_   :     Optional[Any]								=     False


				def __UpperCamelCase						(							self      :  Tuple       )		->     str:
											A      						=    32
											A      						=    embedder_hidden_size

											# prior components

											torch.manual_seed(0       )
											A      						=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

											torch.manual_seed(0       )
											A      						=    CLIPTextModelWithProjection(
											    CLIPTextConfig(
											        bos_token_id=0						, eos_token_id=2						, hidden_size=__UpperCamelCase						, projection_dim=__UpperCamelCase						, intermediate_size=37						, layer_norm_eps=1e-05						, num_attention_heads=4						, num_hidden_layers=5						, pad_token_id=1						, vocab_size=1_000						, )       )

											torch.manual_seed(0       )
											A      						=    PriorTransformer(
											    num_attention_heads=2						, attention_head_dim=12						, embedding_dim=__UpperCamelCase						, num_layers=1						, )

											torch.manual_seed(0       )
											A      						=    DDPMScheduler(
											    variance_type='fixed_small_log'						, prediction_type='sample'						, num_train_timesteps=1_000						, clip_sample=__UpperCamelCase						, clip_sample_range=5.0						, beta_schedule='squaredcos_cap_v2'						, )

											# regular denoising components

											torch.manual_seed(0       )
											A      						=    StableUnCLIPImageNormalizer(embedding_dim=__UpperCamelCase       )
											A      						=    DDPMScheduler(beta_schedule='squaredcos_cap_v2'       )

											torch.manual_seed(0       )
											A      						=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

											torch.manual_seed(0       )
											A      						=    CLIPTextModel(
											    CLIPTextConfig(
											        bos_token_id=0						, eos_token_id=2						, hidden_size=__UpperCamelCase						, projection_dim=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						, )       )

											torch.manual_seed(0       )
											A      						=    UNetaDConditionModel(
											    sample_size=32						, in_channels=4						, out_channels=4						, down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')						, up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')						, block_out_channels=(32, 64)						, attention_head_dim=(2, 4)						, class_embed_type='projection'						, projection_class_embeddings_input_dim=embedder_projection_dim * 2						, cross_attention_dim=__UpperCamelCase						, layers_per_block=1						, upcast_attention=__UpperCamelCase						, use_linear_projection=__UpperCamelCase						, )

											torch.manual_seed(0       )
											A      						=    DDIMScheduler(
											    beta_schedule='scaled_linear'						, beta_start=0.0_0_0_8_5						, beta_end=0.0_1_2						, prediction_type='v_prediction'						, set_alpha_to_one=__UpperCamelCase						, steps_offset=1						, )

											torch.manual_seed(0       )
											A      						=    AutoencoderKL()

											A      						=    {
											    # prior components
											    'prior_tokenizer': prior_tokenizer,
											    'prior_text_encoder': prior_text_encoder,
											    'prior': prior,
											    'prior_scheduler': prior_scheduler,
											    # image noising components
											    'image_normalizer': image_normalizer,
											    'image_noising_scheduler': image_noising_scheduler,
											    # regular denoising components
											    'tokenizer': tokenizer,
											    'text_encoder': text_encoder,
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											}

											return components


				def __UpperCamelCase						(							self      :  Optional[int]						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  str=0       )		->     Optional[int]:
											if str(__UpperCamelCase       ).startswith('mps'       ):
																		A      						=    torch.manual_seed(__UpperCamelCase       )
											else:
																		A      						=    torch.Generator(device=__UpperCamelCase       ).manual_seed(__UpperCamelCase       )
											A      						=    {
											    'prompt': 'A painting of a squirrel eating a burger',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'prior_num_inference_steps': 2,
											    'output_type': 'numpy',
											}
											return inputs


				def __UpperCamelCase						(							self      :  Union[str, Any]       )		->     Union[str, Any]:
											A      						=    torch_device == 'cpu'

											self._test_attention_slicing_forward_pass(test_max_difference=__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     Tuple:
											A      						=    torch_device in ['cpu', 'mps']

											self._test_inference_batch_single_identical(test_max_difference=__UpperCamelCase       )





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


				def __UpperCamelCase						(							self      :  Any       )		->     Union[str, Any]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


				def __UpperCamelCase						(							self      :  List[Any]       )		->     Tuple:
											A      						=    load_numpy(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'       )

											A      						=    StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'						, torch_dtype=torch.floataa       )
											pipe.to(__UpperCamelCase       )
											pipe.set_progress_bar_config(disable=__UpperCamelCase       )
											# stable unclip will oom when integration tests are run on a V100,
											# so turn on memory savings
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											A      						=    torch.Generator(device='cpu'       ).manual_seed(0       )
											A      						=    pipe('anime turle'						, generator=__UpperCamelCase						, output_type='np'       )

											A      						=    output.images[0]

											assert image.shape == (768, 768, 3)

											assert_mean_pixel_difference(__UpperCamelCase						, __UpperCamelCase       )


				def __UpperCamelCase						(							self      :  List[str]       )		->     Dict:
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

											A      						=    StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'						, torch_dtype=torch.floataa       )
											A      						=    pipe.to(__UpperCamelCase       )
											pipe.set_progress_bar_config(disable=__UpperCamelCase       )
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											A      						=    pipe(
											    'anime turtle'						, prior_num_inference_steps=2						, num_inference_steps=2						, output_type='np'						, )

											A      						=    torch.cuda.max_memory_allocated()
											# make sure that less than 7 GB is allocated
											assert mem_bytes < 7 * 10**9




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


__snake_case     :Optional[Any]			       =logging.get_logger(__name__)

__snake_case     :Tuple			       ={
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}



class        lowerCAmelCase__					(  _lowerCamelCase       ):
				A_   :     Any								=     'transfo-xl'
				A_   :     Dict								=     ['mems']
				A_   :     int								=     {
				    'n_token': 'vocab_size',
				    'hidden_size': 'd_model',
				    'num_attention_heads': 'n_head',
				    'num_hidden_layers': 'n_layer',
				}


				def __init__(							self      :  Optional[Any]						, __UpperCamelCase      :  str=267_735						, __UpperCamelCase      :  int=[20_000, 40_000, 200_000]						, __UpperCamelCase      :  List[str]=1_024						, __UpperCamelCase      :  Union[str, Any]=1_024						, __UpperCamelCase      :  Optional[Any]=16						, __UpperCamelCase      :  Dict=64						, __UpperCamelCase      :  Optional[Any]=4_096						, __UpperCamelCase      :  int=4						, __UpperCamelCase      :  List[str]=False						, __UpperCamelCase      :  Optional[int]=18						, __UpperCamelCase      :  int=1_600						, __UpperCamelCase      :  Dict=1_000						, __UpperCamelCase      :  Optional[Any]=True						, __UpperCamelCase      :  Optional[Any]=True						, __UpperCamelCase      :  Optional[Any]=0						, __UpperCamelCase      :  Any=-1						, __UpperCamelCase      :  Any=True						, __UpperCamelCase      :  List[Any]=0.1						, __UpperCamelCase      :  str=0.0						, __UpperCamelCase      :  Dict=True						, __UpperCamelCase      :  int="normal"						, __UpperCamelCase      :  List[Any]=0.0_1						, __UpperCamelCase      :  str=0.0_1						, __UpperCamelCase      :  Any=0.0_2						, __UpperCamelCase      :  str=1e-5						, __UpperCamelCase      :  Optional[int]=0						, **__UpperCamelCase      :  Any						, )		->     int:
											A      						=    vocab_size
											A      						=    []
											self.cutoffs.extend(__UpperCamelCase       )
											if proj_share_all_but_first:
																		A      						=    [False] + [True] * len(self.cutoffs       )
											else:
																		A      						=    [False] + [False] * len(self.cutoffs       )
											A      						=    d_model
											A      						=    d_embed
											A      						=    d_head
											A      						=    d_inner
											A      						=    div_val
											A      						=    pre_lnorm
											A      						=    n_layer
											A      						=    n_head
											A      						=    mem_len
											A      						=    same_length
											A      						=    attn_type
											A      						=    clamp_len
											A      						=    sample_softmax
											A      						=    adaptive
											A      						=    dropout
											A      						=    dropatt
											A      						=    untie_r
											A      						=    init
											A      						=    init_range
											A      						=    proj_init_std
											A      						=    init_std
											A      						=    layer_norm_epsilon
											super().__init__(eos_token_id=__UpperCamelCase						, **__UpperCamelCase       )


				@property
				def __UpperCamelCase						(							self      :  Any       )		->     Optional[int]:
											# Message copied from Transformer-XL documentation
											logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )
											return -1


				@max_position_embeddings.setter
				def __UpperCamelCase						(							self      :  List[Any]						, __UpperCamelCase      :  List[Any]       )		->     List[Any]:
											# Message copied from Transformer-XL documentation
											raise NotImplementedError(
											    f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )