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     _snake_case  (  snake_case__    ,					snake_case__    ,					unittest.TestCase ):
					lowerCamelCase__:						Optional[int]     =   StableDiffusionXLImgaImgPipeline
					lowerCamelCase__:						Any     =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
					lowerCamelCase__:						str     =   PipelineTesterMixin.required_optional_params - {"latents"}
					lowerCamelCase__:						Union[str, Any]     =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					lowerCamelCase__:						Dict     =   IMAGE_TO_IMAGE_IMAGE_PARAMS
					lowerCamelCase__:						List[str]     =   IMAGE_TO_IMAGE_IMAGE_PARAMS
					def       _lowerCamelCase       (     self: Union[str, Any]   )			-> int:
												torch.manual_seed(0   )
												__UpperCAmelCase  :	List[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=__lowerCamelCase		,     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[Any]   			=      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  :	Tuple   			=      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=1_28		,     )
												torch.manual_seed(0   )
												__UpperCAmelCase  :	List[str]   			=      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=10_00		,     hidden_act="gelu"		,     projection_dim=32		,     )
												__UpperCAmelCase  :	List[str]   			=      CLIPTextModel(__lowerCamelCase   )
												__UpperCAmelCase  :	List[Any]   			=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		,     local_files_only=__lowerCamelCase   )

												__UpperCAmelCase  :	str   			=      CLIPTextModelWithProjection(__lowerCamelCase   )
												__UpperCAmelCase  :	int   			=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		,     local_files_only=__lowerCamelCase   )

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

												__UpperCAmelCase  :	Tuple   			=      self.get_dummy_inputs(__lowerCamelCase   )
												__UpperCAmelCase  :	Optional[int]   			=      sd_pipe(**__lowerCamelCase   ).images
												__UpperCAmelCase  :	Dict   			=      image[0, -3:, -3:, -1]

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

												__UpperCAmelCase  :	Union[str, 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       _lowerCamelCase       (     self: int   )			-> Dict:
												super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )
					def       _lowerCamelCase       (     self: Optional[int]   )			-> List[str]:
												super().test_inference_batch_single_identical(expected_max_diff=3e-3   )
					def       _lowerCamelCase       (     self: str   )			-> int:
												pass





					def       _lowerCamelCase       (     self: Tuple   )			-> Dict:
												__UpperCAmelCase  :	str   			=      self.get_dummy_components()
												__UpperCAmelCase  :	Dict   			=      StableDiffusionXLImgaImgPipeline(**__lowerCamelCase   )
												__UpperCAmelCase  :	List[str]   			=      sd_pipe.to(__lowerCamelCase   )
												__UpperCAmelCase  :	Union[str, Any]   			=      sd_pipe.to(__lowerCamelCase   )
												sd_pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												# forward without prompt embeds
												__UpperCAmelCase  :	int   			=      self.get_dummy_inputs(__lowerCamelCase   )
												__UpperCAmelCase  :	Union[str, Any]   			=      3 * ["this is a negative prompt"]
												__UpperCAmelCase  :	List[Any]   			=      negative_prompt
												__UpperCAmelCase  :	Optional[int]   			=      3 * [inputs["prompt"]]

												__UpperCAmelCase  :	Optional[int]   			=      sd_pipe(**__lowerCamelCase   )
												__UpperCAmelCase  :	Dict   			=      output.images[0, -3:, -3:, -1]

												# forward with prompt embeds
												__UpperCAmelCase  :	Tuple   			=      self.get_dummy_inputs(__lowerCamelCase   )
												__UpperCAmelCase  :	Any   			=      3 * ["this is a negative prompt"]
												__UpperCAmelCase  :	List[Any]   			=      3 * [inputs.pop("prompt"   )]

												(
												    (
												    __UpperCAmelCase 
												)			,		(
												    __UpperCAmelCase 
												)			,		(
												    __UpperCAmelCase 
												)			,		(
												    __UpperCAmelCase 
												)			,		 
												)  :	Tuple   			=      sd_pipe.encode_prompt(__lowerCamelCase		,     negative_prompt=__lowerCamelCase   )

												__UpperCAmelCase  :	List[str]   			=      sd_pipe(
												    **__lowerCamelCase		,     prompt_embeds=__lowerCamelCase		,     negative_prompt_embeds=__lowerCamelCase		,     pooled_prompt_embeds=__lowerCamelCase		,     negative_pooled_prompt_embeds=__lowerCamelCase		,     )
												__UpperCAmelCase  :	Any   			=      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     _snake_case  (  unittest.TestCase ):
					def       _lowerCamelCase       (     self: Any   )			-> List[str]:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					def       _lowerCamelCase       (     self: Dict		,     __lowerCamelCase: Union[str, Any]		,     __lowerCamelCase: int="cpu"		,     __lowerCamelCase: Dict=torch.floataa		,     __lowerCamelCase: List[str]=0   )			-> int:
												__UpperCAmelCase  :	int   			=      torch.Generator(device=__lowerCamelCase   ).manual_seed(__lowerCamelCase   )
												__UpperCAmelCase  :	List[str]   			=      np.random.RandomState(__lowerCamelCase   ).standard_normal((1, 4, 64, 64)   )
												__UpperCAmelCase  :	Dict   			=      torch.from_numpy(__lowerCamelCase   ).to(device=__lowerCamelCase		,     dtype=__lowerCamelCase   )
												__UpperCAmelCase  :	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       _lowerCamelCase       (     self: Union[str, Any]   )			-> Optional[Any]:
												__UpperCAmelCase  :	Any   			=      DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"   )
												pipe.to(__lowerCamelCase   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												__UpperCAmelCase  :	Union[str, Any]   			=      self.get_inputs(__lowerCamelCase   )
												__UpperCAmelCase  :	Dict   			=      pipe(**__lowerCamelCase   ).images
												__UpperCAmelCase  :	Dict   			=      image[0, -3:, -3:, -1].flatten()

												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	List[str]   			=      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












from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_a      	=  logging.get_logger(__name__)

_a      	=  {
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






	UpperCamelCase__								=      "distilbert"
	UpperCamelCase__								=      {
	    "hidden_size": "dim",
	    "num_attention_heads": "n_heads",
	    "num_hidden_layers": "n_layers",
	}

	def __init__(							self       ,    UpperCAmelCase=3_0522       ,    UpperCAmelCase=512       ,    UpperCAmelCase=False       ,    UpperCAmelCase=6       ,    UpperCAmelCase=12       ,    UpperCAmelCase=768       ,    UpperCAmelCase=4 * 768       ,    UpperCAmelCase=0.1       ,    UpperCAmelCase=0.1       ,    UpperCAmelCase="gelu"       ,    UpperCAmelCase=0.02       ,    UpperCAmelCase=0.1       ,    UpperCAmelCase=0.2       ,    UpperCAmelCase=0       ,    **UpperCAmelCase       ,    ):
							"""simple docstring"""







							_UpperCAmelCase									=			vocab_size
							_UpperCAmelCase									=			max_position_embeddings
							_UpperCAmelCase									=			sinusoidal_pos_embds
							_UpperCAmelCase									=			n_layers
							_UpperCAmelCase									=			n_heads
							_UpperCAmelCase									=			dim
							_UpperCAmelCase									=			hidden_dim
							_UpperCAmelCase									=			dropout
							_UpperCAmelCase									=			attention_dropout
							_UpperCAmelCase									=			activation
							_UpperCAmelCase									=			initializer_range
							_UpperCAmelCase									=			qa_dropout
							_UpperCAmelCase									=			seq_classif_dropout
							super().__init__(**UpperCAmelCase       ,    pad_token_id=UpperCAmelCase    )





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							if self.task == "multiple-choice":
													_UpperCAmelCase									=			{0: 'batch', 1: 'choice', 2: 'sequence'}
							else:
													_UpperCAmelCase									=			{0: 'batch', 1: 'sequence'}
							return OrderedDict(
							    [
							        ('input_ids', dynamic_axis),
							        ('attention_mask', dynamic_axis),
							    ]    )
