






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    (      _A							,     _A							,     unittest.TestCase			):
	_A 							= StableDiffusionXLImgaImgPipeline
	_A 							= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
	_A 							= PipelineTesterMixin.required_optional_params - {'latents'}
	_A 							= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	_A 							= IMAGE_TO_IMAGE_IMAGE_PARAMS
	_A 							= IMAGE_TO_IMAGE_IMAGE_PARAMS






	def  lowerCAmelCase_				(		self							)    ->     List[str]:
					torch.manual_seed(0							)
					snake_case__   :Union[str, 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=a_       ,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__   :Optional[Any]       = EulerDiscreteScheduler(
					    beta_start=0.00085       ,beta_end=0.012       ,steps_offset=1       ,beta_schedule="scaled_linear"       ,timestep_spacing="leading"       ,)
					torch.manual_seed(0							)
					snake_case__   :Dict       = 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__   :Dict       = 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(a_							)
					snake_case__   :Optional[Any]       = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       ,local_files_only=a_							)

					snake_case__   :Tuple       = CLIPTextModelWithProjection(a_							)
					snake_case__   :int       = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       ,local_files_only=a_							)

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

					snake_case__   :List[Any]       = self.get_dummy_inputs(a_							)
					snake_case__   :Any       = sd_pipe(**a_							).images
					snake_case__   :Tuple       = image[0, -3:, -3:, -1]

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

					snake_case__   :Optional[int]       = np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]							)

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






	def  lowerCAmelCase_				(		self							)    ->     List[Any]:
					super().test_attention_slicing_forward_pass(expected_max_diff=3E-3							)






	def  lowerCAmelCase_				(		self							)    ->     List[str]:
					super().test_inference_batch_single_identical(expected_max_diff=3E-3							)






	def  lowerCAmelCase_				(		self							)    ->     str:
					pass






	def  lowerCAmelCase_				(		self							)    ->     Tuple:
					snake_case__   :List[str]       = self.get_dummy_components()
					snake_case__   :Tuple       = StableDiffusionXLImgaImgPipeline(**a_							)
					snake_case__   :int       = sd_pipe.to(a_							)
					snake_case__   :Optional[int]       = sd_pipe.to(a_							)
					sd_pipe.set_progress_bar_config(disable=a_							)

					# forward without prompt embeds
					snake_case__   :Optional[Any]       = self.get_dummy_inputs(a_							)
					snake_case__   :List[Any]       = 3 * ["this is a negative prompt"]
					snake_case__   :Tuple       = negative_prompt
					snake_case__   :Optional[Any]       = 3 * [inputs["prompt"]]

					snake_case__   :Optional[Any]       = sd_pipe(**a_							)
					snake_case__   :Any       = output.images[0, -3:, -3:, -1]

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

					(
					    snake_case__		
					)   :Dict       = sd_pipe.encode_prompt(a_       ,negative_prompt=a_							)

					snake_case__   :List[str]       = sd_pipe(
					    **a_       ,prompt_embeds=a_       ,negative_prompt_embeds=a_       ,pooled_prompt_embeds=a_       ,negative_pooled_prompt_embeds=a_       ,)
					snake_case__   :List[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							)    ->     Tuple:
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






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

					snake_case__   :Any       = self.get_inputs(a_							)
					snake_case__   :Any       = pipe(**a_							).images
					snake_case__   :Any       = image[0, -3:, -3:, -1].flatten()

					assert image.shape == (1, 512, 512, 3)
					snake_case__   :Union[str, Any]       = np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]							)
					assert np.abs(image_slice - expected_slice							).max() < 7E-3






'''simple docstring'''






# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCAmelCase									=   {"""configuration_mra""": ["""MRA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MraConfig"""]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCAmelCase									=   [
					    """MRA_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """MraForMaskedLM""",
					    """MraForMultipleChoice""",
					    """MraForQuestionAnswering""",
					    """MraForSequenceClassification""",
					    """MraForTokenClassification""",
					    """MraLayer""",
					    """MraModel""",
					    """MraPreTrainedModel""",
					]


if TYPE_CHECKING:
					from .configuration_mra import MRA_PRETRAINED_CONFIG_ARCHIVE_MAP, MraConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mra import (
										    MRA_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MraForMaskedLM,
										    MraForMultipleChoice,
										    MraForQuestionAnswering,
										    MraForSequenceClassification,
										    MraForTokenClassification,
										    MraLayer,
										    MraModel,
										    MraPreTrainedModel,
										)
else:
					import sys

					lowerCAmelCase									=   _LazyModule(__name__, globals()["""__file__"""], _import_structure)



