






import gc
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class  lowerCAmelCase_       (				snake_case__     ,    unittest.TestCase					):
							UpperCamelCase_		:List[str]		      =    LDMTextToImagePipeline
							UpperCamelCase_		:List[Any]		      =    TEXT_TO_IMAGE_PARAMS - {
							    'negative_prompt',
							    'negative_prompt_embeds',
							    'cross_attention_kwargs',
							    'prompt_embeds',
							}
							UpperCamelCase_		:Tuple		      =    PipelineTesterMixin.required_optional_params - {
							    'num_images_per_prompt',
							    'callback',
							    'callback_steps',
							}
							UpperCamelCase_		:Union[str, Any]		      =    TEXT_TO_IMAGE_BATCH_PARAMS
							UpperCamelCase_		:int		      =    False




							def 						__snake_case					(   self     :   Tuple	):
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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''') ,							cross_attention_dim=32 ,							)
											lowerCAmelCase__      			=						DDIMScheduler(
											    beta_start=0.00_085 ,							beta_end=0.012 ,							beta_schedule='''scaled_linear''' ,							clip_sample=SCREAMING_SNAKE_CASE_ ,							set_alpha_to_one=SCREAMING_SNAKE_CASE_ ,							)
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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 ,							)
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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 ,							)
											lowerCAmelCase__      			=						CLIPTextModel(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

											lowerCAmelCase__      			=						{
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vqvae''': vae,
											    '''bert''': text_encoder,
											    '''tokenizer''': tokenizer,
											}
											return components




							def 						__snake_case					(   self     :   int ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   int=0	):
											if str(SCREAMING_SNAKE_CASE_	).startswith('''mps'''	):
															lowerCAmelCase__      			=						torch.manual_seed(SCREAMING_SNAKE_CASE_	)
											else:
															lowerCAmelCase__      			=						torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{
											    '''prompt''': '''A painting of a squirrel eating a burger''',
											    '''generator''': generator,
											    '''num_inference_steps''': 2,
											    '''guidance_scale''': 6.0,
											    '''output_type''': '''numpy''',
											}
											return inputs




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator

											lowerCAmelCase__      			=						self.get_dummy_components()
											lowerCAmelCase__      			=						LDMTextToImagePipeline(**SCREAMING_SNAKE_CASE_	)
											pipe.to(SCREAMING_SNAKE_CASE_	)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						pipe(**SCREAMING_SNAKE_CASE_	).images
											lowerCAmelCase__      			=						image[0, -3:, -3:, -1]

											assert image.shape == (1, 16, 16, 3)
											lowerCAmelCase__      			=						np.array([0.6_101, 0.6_156, 0.5_622, 0.4_895, 0.6_661, 0.3_804, 0.5_748, 0.6_136, 0.5_014]	)

											assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3



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




							def 						__snake_case					(   self     :   str	):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




							def 						__snake_case					(   self     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Dict=torch.floataa ,							SCREAMING_SNAKE_CASE_     :   int=0	):
											lowerCAmelCase__      			=						torch.manual_seed(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						np.random.RandomState(SCREAMING_SNAKE_CASE_	).standard_normal((1, 4, 32, 32)	)
											lowerCAmelCase__      			=						torch.from_numpy(SCREAMING_SNAKE_CASE_	).to(device=SCREAMING_SNAKE_CASE_ ,							dtype=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{
											    '''prompt''': '''A painting of a squirrel eating a burger''',
											    '''latents''': latents,
											    '''generator''': generator,
											    '''num_inference_steps''': 3,
											    '''guidance_scale''': 6.0,
											    '''output_type''': '''numpy''',
											}
											return inputs




							def 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''	).to(SCREAMING_SNAKE_CASE_	)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

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

											assert image.shape == (1, 256, 256, 3)
											lowerCAmelCase__      			=						np.array([0.51_825, 0.52_850, 0.52_543, 0.54_258, 0.52_304, 0.52_569, 0.54_363, 0.55_276, 0.56_878]	)
											lowerCAmelCase__      			=						np.abs(expected_slice - image_slice	).max()
											assert max_diff < 1e-3



@nightly
@require_torch_gpu
class  lowerCAmelCase_       (				unittest.TestCase					):




							def 						__snake_case					(   self     :   Any	):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




							def 						__snake_case					(   self     :   int ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Tuple=torch.floataa ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=0	):
											lowerCAmelCase__      			=						torch.manual_seed(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						np.random.RandomState(SCREAMING_SNAKE_CASE_	).standard_normal((1, 4, 32, 32)	)
											lowerCAmelCase__      			=						torch.from_numpy(SCREAMING_SNAKE_CASE_	).to(device=SCREAMING_SNAKE_CASE_ ,							dtype=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{
											    '''prompt''': '''A painting of a squirrel eating a burger''',
											    '''latents''': latents,
											    '''generator''': generator,
											    '''num_inference_steps''': 50,
											    '''guidance_scale''': 6.0,
											    '''output_type''': '''numpy''',
											}
											return inputs




							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''	).to(SCREAMING_SNAKE_CASE_	)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						self.get_inputs(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						pipe(**SCREAMING_SNAKE_CASE_	).images[0]

											lowerCAmelCase__      			=						load_numpy(
											    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy'''	)
											lowerCAmelCase__      			=						np.abs(expected_image - image	).max()
											assert max_diff < 1e-3













from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


_UpperCAmelCase			:		Tuple   				=    logging.get_logger(__name__)


_UpperCAmelCase			:		int   				=    {
    "Salesforce/codegen-350M-nl": "https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json",
    "Salesforce/codegen-350M-multi": "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json",
    "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json",
    "Salesforce/codegen-2B-nl": "https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json",
    "Salesforce/codegen-2B-multi": "https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json",
    "Salesforce/codegen-2B-mono": "https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json",
    "Salesforce/codegen-6B-nl": "https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json",
    "Salesforce/codegen-6B-multi": "https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json",
    "Salesforce/codegen-6B-mono": "https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json",
    "Salesforce/codegen-16B-nl": "https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json",
    "Salesforce/codegen-16B-multi": "https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json",
    "Salesforce/codegen-16B-mono": "https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json",
}
class  lowerCAmelCase_       (				snake_case__					):
							UpperCamelCase_		:int		      =    'codegen'
							UpperCamelCase_		:int		      =    {
							    'max_position_embeddings': 'n_positions',
							    'hidden_size': 'n_embd',
							    'num_attention_heads': 'n_head',
							    'num_hidden_layers': 'n_layer',
							}




							def __init__(   self     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   str=50_400 ,							SCREAMING_SNAKE_CASE_     :   str=2_048 ,							SCREAMING_SNAKE_CASE_     :   int=2_048 ,							SCREAMING_SNAKE_CASE_     :   Any=4_096 ,							SCREAMING_SNAKE_CASE_     :   List[Any]=28 ,							SCREAMING_SNAKE_CASE_     :   str=16 ,							SCREAMING_SNAKE_CASE_     :   str=64 ,							SCREAMING_SNAKE_CASE_     :   Tuple=None ,							SCREAMING_SNAKE_CASE_     :   Dict="gelu_new" ,							SCREAMING_SNAKE_CASE_     :   str=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=1e-5 ,							SCREAMING_SNAKE_CASE_     :   Any=0.02 ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any]=True ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=50_256 ,							SCREAMING_SNAKE_CASE_     :   Any=50_256 ,							SCREAMING_SNAKE_CASE_     :   List[str]=False ,							**SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							):
											lowerCAmelCase__      			=						vocab_size
											lowerCAmelCase__      			=						n_ctx
											lowerCAmelCase__      			=						n_positions
											lowerCAmelCase__      			=						n_embd
											lowerCAmelCase__      			=						n_layer
											lowerCAmelCase__      			=						n_head
											lowerCAmelCase__      			=						n_inner
											lowerCAmelCase__      			=						rotary_dim
											lowerCAmelCase__      			=						activation_function
											lowerCAmelCase__      			=						resid_pdrop
											lowerCAmelCase__      			=						embd_pdrop
											lowerCAmelCase__      			=						attn_pdrop
											lowerCAmelCase__      			=						layer_norm_epsilon
											lowerCAmelCase__      			=						initializer_range
											lowerCAmelCase__      			=						use_cache

											lowerCAmelCase__      			=						bos_token_id
											lowerCAmelCase__      			=						eos_token_id

											super().__init__(
											    bos_token_id=SCREAMING_SNAKE_CASE_ ,							eos_token_id=SCREAMING_SNAKE_CASE_ ,							tie_word_embeddings=SCREAMING_SNAKE_CASE_ ,							**SCREAMING_SNAKE_CASE_	)



class  lowerCAmelCase_       (				snake_case__					):




							def __init__(   self     :   str ,							SCREAMING_SNAKE_CASE_     :   PretrainedConfig ,							SCREAMING_SNAKE_CASE_     :   str = "default" ,							SCREAMING_SNAKE_CASE_     :   List[PatchingSpec] = None ,							SCREAMING_SNAKE_CASE_     :   bool = False ,							):
											super().__init__(SCREAMING_SNAKE_CASE_ ,							task=SCREAMING_SNAKE_CASE_ ,							patching_specs=SCREAMING_SNAKE_CASE_ ,							use_past=SCREAMING_SNAKE_CASE_	)
											if not getattr(self._config ,							'''pad_token_id''' ,							SCREAMING_SNAKE_CASE_	):
															# TODO: how to do that better?
															lowerCAmelCase__      			=						0




							@property
							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}	)
											if self.use_past:
															self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_ ,							direction='''inputs'''	)
															lowerCAmelCase__      			=						{0: '''batch''', 1: '''past_sequence + sequence'''}
											else:
															lowerCAmelCase__      			=						{0: '''batch''', 1: '''sequence'''}

											return common_inputs




							@property
							def 						__snake_case					(   self     :   Dict	):
											return self._config.n_layer




							@property
							def 						__snake_case					(   self     :   Union[str, Any]	):
											return self._config.n_head




							def 						__snake_case					(   self     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   PreTrainedTokenizer ,							SCREAMING_SNAKE_CASE_     :   int = -1 ,							SCREAMING_SNAKE_CASE_     :   int = -1 ,							SCREAMING_SNAKE_CASE_     :   bool = False ,							SCREAMING_SNAKE_CASE_     :   Optional[TensorType] = None ,							):
											lowerCAmelCase__      			=						super(SCREAMING_SNAKE_CASE_ ,							self	).generate_dummy_inputs(
											    SCREAMING_SNAKE_CASE_ ,							batch_size=SCREAMING_SNAKE_CASE_ ,							seq_length=SCREAMING_SNAKE_CASE_ ,							is_pair=SCREAMING_SNAKE_CASE_ ,							framework=SCREAMING_SNAKE_CASE_	)

											# We need to order the input in the way they appears in the forward()
											lowerCAmelCase__      			=						OrderedDict({'''input_ids''': common_inputs['''input_ids''']}	)

											# Need to add the past_keys
											if self.use_past:
															if not is_torch_available():
																			raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
															else:
																			import torch

																			lowerCAmelCase__			,		lowerCAmelCase__      			=						common_inputs['''input_ids'''].shape
																			# Not using the same length for past_key_values
																			lowerCAmelCase__      			=						seqlen + 2
																			lowerCAmelCase__      			=						(
																			    batch,
																			    self.num_attention_heads,
																			    past_key_values_length,
																			    self._config.hidden_size // self.num_attention_heads,
																			)
																			lowerCAmelCase__      			=						[
																			    (torch.zeros(SCREAMING_SNAKE_CASE_	), torch.zeros(SCREAMING_SNAKE_CASE_	)) for _ in range(self.num_layers	)
																			]

											lowerCAmelCase__      			=						common_inputs['''attention_mask''']
											if self.use_past:
															lowerCAmelCase__      			=						ordered_inputs['''attention_mask'''].dtype
															lowerCAmelCase__      			=						torch.cat(
															    [ordered_inputs['''attention_mask'''], torch.ones(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							dtype=SCREAMING_SNAKE_CASE_	)] ,							dim=1	)

											return ordered_inputs




							@property
							def 						__snake_case					(   self     :   Optional[int]	):
											return 13





