



'''simple docstring'''






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        SCREAMING_SNAKE_CASE(							A__							,	A__							,	unittest.TestCase    ):







	"""simple docstring"""






	lowerCamelCase__    				= StableDiffusionXLImgaImgPipeline
	lowerCamelCase__    				= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
	lowerCamelCase__    				= PipelineTesterMixin.required_optional_params - {"""latents"""}
	lowerCamelCase__    				= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	lowerCamelCase__    				= IMAGE_TO_IMAGE_IMAGE_PARAMS
	lowerCamelCase__    				= IMAGE_TO_IMAGE_IMAGE_PARAMS







	def   A			(		self    :  List[str]	)    ->	Optional[int]:
								torch.manual_seed(0	)
								UpperCAmelCase				: Optional[int]											=	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.0_00_85	,    beta_end=0.0_12	,    steps_offset=1	,    beta_schedule='''scaled_linear'''	,    timestep_spacing='''leading'''	,    )
								torch.manual_seed(0	)
								UpperCAmelCase				: str											=	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				: Any											=	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=1000	,    hidden_act='''gelu'''	,    projection_dim=32	,    )
								UpperCAmelCase				: Dict											=	CLIPTextModel(__snake_case	)
								UpperCAmelCase				: Any											=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	,    local_files_only=__snake_case	)

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

								UpperCAmelCase				: str											=	{
								    '''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    :  str	,    __snake_case    :  Dict	,    __snake_case    :  Union[str, Any]=0	)    ->	int:
								UpperCAmelCase				: Dict											=	floats_tensor((1, 3, 32, 32)	,    rng=random.Random(__snake_case	)	).to(__snake_case	)
								UpperCAmelCase				: Any											=	image / 2 + 0.5
								if str(__snake_case	).startswith('''mps'''	):
															UpperCAmelCase				: List[str]											=	torch.manual_seed(__snake_case	)
								else:
															UpperCAmelCase				: List[str]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: int											=	{
								    '''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   A			(		self    :  str	)    ->	List[str]:
								UpperCAmelCase				: Tuple											=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase				: Union[str, Any]											=	self.get_dummy_components()
								UpperCAmelCase				: str											=	StableDiffusionXLImgaImgPipeline(**__snake_case	)
								UpperCAmelCase				: List[Any]											=	sd_pipe.to(__snake_case	)
								sd_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Any											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: Tuple											=	sd_pipe(**__snake_case	).images
								UpperCAmelCase				: Union[str, Any]											=	image[0, -3:, -3:, -1]

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

								UpperCAmelCase				: Any											=	np.array([0.46_56, 0.48_40, 0.44_39, 0.66_98, 0.55_74, 0.45_24, 0.57_99, 0.59_43, 0.51_65]	)

								assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2







	def   A			(		self    :  Optional[Any]	)    ->	Union[str, Any]:
								super().test_attention_slicing_forward_pass(expected_max_diff=3E-3	)







	def   A			(		self    :  Optional[Any]	)    ->	Tuple:
								super().test_inference_batch_single_identical(expected_max_diff=3E-3	)







	def   A			(		self    :  Optional[int]	)    ->	int:
								pass







	def   A			(		self    :  Dict	)    ->	str:
								UpperCAmelCase				: Any											=	self.get_dummy_components()
								UpperCAmelCase				: Tuple											=	StableDiffusionXLImgaImgPipeline(**__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	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				: Any											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: List[Any]											=	3 * ['''this is a negative prompt''']
								UpperCAmelCase				: List[str]											=	negative_prompt
								UpperCAmelCase				: List[Any]											=	3 * [inputs['''prompt''']]

								UpperCAmelCase				: List[str]											=	sd_pipe(**__snake_case	)
								UpperCAmelCase				: Tuple											=	output.images[0, -3:, -3:, -1]

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

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

								UpperCAmelCase				: Dict											=	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				: str											=	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        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	def   A			(		self    :  List[str]	)    ->	List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   A			(		self    :  List[Any]	,    __snake_case    :  Union[str, Any]	,    __snake_case    :  List[str]="cpu"	,    __snake_case    :  Tuple=torch.floataa	,    __snake_case    :  str=0	)    ->	Optional[int]:
								UpperCAmelCase				: List[str]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: List[Any]											=	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				: List[str]											=	{
								    '''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]	)    ->	List[Any]:
								UpperCAmelCase				: Optional[Any]											=	DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''	)
								pipe.to(__snake_case	)
								pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: int											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: Optional[int]											=	pipe(**__snake_case	).images
								UpperCAmelCase				: Tuple											=	image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase				: List[Any]											=	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'''






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


UpperCamelCase__:     Dict										=    logging.get_logger(__name__)

UpperCamelCase__:     int										=    {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}

class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= """megatron-bert"""







	def __init__(		self    :  int	,    __snake_case    :  List[Any]=29056	,    __snake_case    :  Dict=1024	,    __snake_case    :  Tuple=24	,    __snake_case    :  Any=16	,    __snake_case    :  Tuple=4096	,    __snake_case    :  int="gelu"	,    __snake_case    :  Dict=0.1	,    __snake_case    :  str=0.1	,    __snake_case    :  Optional[Any]=512	,    __snake_case    :  Tuple=2	,    __snake_case    :  Tuple=0.02	,    __snake_case    :  int=1E-12	,    __snake_case    :  Union[str, Any]=0	,    __snake_case    :  int="absolute"	,    __snake_case    :  str=True	,    **__snake_case    :  Tuple	,    )    ->	Union[str, Any]:
								super().__init__(pad_token_id=__snake_case	,    **__snake_case	)

								UpperCAmelCase				: Any											=	vocab_size
								UpperCAmelCase				: Tuple											=	hidden_size
								UpperCAmelCase				: List[Any]											=	num_hidden_layers
								UpperCAmelCase				: str											=	num_attention_heads
								UpperCAmelCase				: Optional[int]											=	hidden_act
								UpperCAmelCase				: Optional[int]											=	intermediate_size
								UpperCAmelCase				: Dict											=	hidden_dropout_prob
								UpperCAmelCase				: List[Any]											=	attention_probs_dropout_prob
								UpperCAmelCase				: Dict											=	max_position_embeddings
								UpperCAmelCase				: Any											=	type_vocab_size
								UpperCAmelCase				: str											=	initializer_range
								UpperCAmelCase				: str											=	layer_norm_eps
								UpperCAmelCase				: Optional[Any]											=	position_embedding_type
								UpperCAmelCase				: str											=	use_cache





