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        lowercase_   (					__snake_case						,			__snake_case						,			unittest.TestCase		):
							_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 	UpperCamelCase	(						self    ):
													torch.manual_seed(0    )
													_snake_case						:					str    =   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=lowercase_				,    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				,    )
													_snake_case						:					Any    =   EulerDiscreteScheduler(
													    beta_start=0.00_085				,    beta_end=0.012				,    steps_offset=1				,    beta_schedule="scaled_linear"				,    timestep_spacing="leading"				,    )
													torch.manual_seed(0    )
													_snake_case						:					List[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    )
													_snake_case						:					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=1_000				,    hidden_act="gelu"				,    projection_dim=32				,    )
													_snake_case						:					Tuple    =   CLIPTextModel(lowercase_    )
													_snake_case						:					Optional[Any]    =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,    local_files_only=lowercase_    )

													_snake_case						:					Union[str, Any]    =   CLIPTextModelWithProjection(lowercase_    )
													_snake_case						:					Optional[int]    =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				,    local_files_only=lowercase_    )

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

													_snake_case						:					Dict    =   self.get_dummy_inputs(lowercase_    )
													_snake_case						:					str    =   sd_pipe(**lowercase_    ).images
													_snake_case						:					Tuple    =   image[0, -3:, -3:, -1]

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

													_snake_case						:					Optional[int]    =   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 	UpperCamelCase	(						self    ):
													super().test_attention_slicing_forward_pass(expected_max_diff=3e-3    )



							def 	UpperCamelCase	(						self    ):
													super().test_inference_batch_single_identical(expected_max_diff=3e-3    )



							def 	UpperCamelCase	(						self    ):
													pass



							def 	UpperCamelCase	(						self    ):
													_snake_case						:					Any    =   self.get_dummy_components()
													_snake_case						:					str    =   StableDiffusionXLImgaImgPipeline(**lowercase_    )
													_snake_case						:					Tuple    =   sd_pipe.to(lowercase_    )
													_snake_case						:					Union[str, Any]    =   sd_pipe.to(lowercase_    )
													sd_pipe.set_progress_bar_config(disable=lowercase_    )

													# forward without prompt embeds
													_snake_case						:					List[str]    =   self.get_dummy_inputs(lowercase_    )
													_snake_case						:					Dict    =   3 * ["this is a negative prompt"]
													_snake_case						:					Any    =   negative_prompt
													_snake_case						:					Tuple    =   3 * [inputs["prompt"]]

													_snake_case						:					List[str]    =   sd_pipe(**lowercase_    )
													_snake_case						:					Dict    =   output.images[0, -3:, -3:, -1]

													# forward with prompt embeds
													_snake_case						:					str    =   self.get_dummy_inputs(lowercase_    )
													_snake_case						:					Any    =   3 * ["this is a negative prompt"]
													_snake_case						:					Any    =   3 * [inputs.pop("prompt"    )]

													(
													    (
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,(
													    _snake_case
													)    ,
													)						:					Dict    =   sd_pipe.encode_prompt(lowercase_				,    negative_prompt=lowercase_    )

													_snake_case						:					Any    =   sd_pipe(
													    **lowercase_				,    prompt_embeds=lowercase_				,    negative_prompt_embeds=lowercase_				,    pooled_prompt_embeds=lowercase_				,    negative_pooled_prompt_embeds=lowercase_				,    )
													_snake_case						:					Dict    =   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        lowercase_   (					unittest.TestCase		):



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



							def 	UpperCamelCase	(						self				,    lowercase_				,    lowercase_="cpu"				,    lowercase_=torch.floataa				,    lowercase_=0    ):
													_snake_case						:					Tuple    =   torch.Generator(device=lowercase_    ).manual_seed(lowercase_    )
													_snake_case						:					str    =   np.random.RandomState(lowercase_    ).standard_normal((1, 4, 64, 64)    )
													_snake_case						:					int    =   torch.from_numpy(lowercase_    ).to(device=lowercase_				,    dtype=lowercase_    )
													_snake_case						:					Any    =   {
													    "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 	UpperCamelCase	(						self    ):
													_snake_case						:					str    =   DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"    )
													pipe.to(lowercase_    )
													pipe.set_progress_bar_config(disable=lowercase_    )

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

													assert image.shape == (1, 512, 512, 3)
													_snake_case						:					List[str]    =   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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__SCREAMING_SNAKE_CASE						:   Optional[int]         =			{
    'configuration_distilbert': [
        'DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'DistilBertConfig',
        'DistilBertOnnxConfig',
    ],
    'tokenization_distilbert': ['DistilBertTokenizer'],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__SCREAMING_SNAKE_CASE						:   Optional[int]         =			['DistilBertTokenizerFast']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__SCREAMING_SNAKE_CASE						:   int         =			[
				    'DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'DistilBertForMaskedLM',
				    'DistilBertForMultipleChoice',
				    'DistilBertForQuestionAnswering',
				    'DistilBertForSequenceClassification',
				    'DistilBertForTokenClassification',
				    'DistilBertModel',
				    'DistilBertPreTrainedModel',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__SCREAMING_SNAKE_CASE						:   Any         =			[
				    'TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'TFDistilBertForMaskedLM',
				    'TFDistilBertForMultipleChoice',
				    'TFDistilBertForQuestionAnswering',
				    'TFDistilBertForSequenceClassification',
				    'TFDistilBertForTokenClassification',
				    'TFDistilBertMainLayer',
				    'TFDistilBertModel',
				    'TFDistilBertPreTrainedModel',
				]

try:
				if not is_flax_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__SCREAMING_SNAKE_CASE						:   Optional[Any]         =			[
				    'FlaxDistilBertForMaskedLM',
				    'FlaxDistilBertForMultipleChoice',
				    'FlaxDistilBertForQuestionAnswering',
				    'FlaxDistilBertForSequenceClassification',
				    'FlaxDistilBertForTokenClassification',
				    'FlaxDistilBertModel',
				    'FlaxDistilBertPreTrainedModel',
				]


if TYPE_CHECKING:
				from .configuration_distilbert import (
				    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    DistilBertConfig,
				    DistilBertOnnxConfig,
				)
				from .tokenization_distilbert import DistilBertTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_distilbert_fast import DistilBertTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_distilbert import (
								    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    DistilBertForMaskedLM,
								    DistilBertForMultipleChoice,
								    DistilBertForQuestionAnswering,
								    DistilBertForSequenceClassification,
								    DistilBertForTokenClassification,
								    DistilBertModel,
								    DistilBertPreTrainedModel,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_distilbert import (
								    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFDistilBertForMaskedLM,
								    TFDistilBertForMultipleChoice,
								    TFDistilBertForQuestionAnswering,
								    TFDistilBertForSequenceClassification,
								    TFDistilBertForTokenClassification,
								    TFDistilBertMainLayer,
								    TFDistilBertModel,
								    TFDistilBertPreTrainedModel,
								)

				try:
								if not is_flax_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_flax_distilbert import (
								    FlaxDistilBertForMaskedLM,
								    FlaxDistilBertForMultipleChoice,
								    FlaxDistilBertForQuestionAnswering,
								    FlaxDistilBertForSequenceClassification,
								    FlaxDistilBertForTokenClassification,
								    FlaxDistilBertModel,
								    FlaxDistilBertPreTrainedModel,
								)

else:
				import sys

				__SCREAMING_SNAKE_CASE						:   str         =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)