





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 				_UpperCAmelCase    (    _UpperCamelCase		,    _UpperCamelCase		,    unittest.TestCase    ):

  """simple docstring"""


  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 				lowercase		(				self :		str     )							->						Optional[int]:
     torch.manual_seed(0     )
     __lowerCAmelCase								=	UNetaDConditionModel(
         block_out_channels=(3_2, 6_4)					,  layers_per_block=2					,  sample_size=3_2					,  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=8_0					,  cross_attention_dim=6_4					,  )
     __lowerCAmelCase								=	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     )
     __lowerCAmelCase								=	AutoencoderKL(
         block_out_channels=[3_2, 6_4]					,  in_channels=3					,  out_channels=3					,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']					,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']					,  latent_channels=4					,  sample_size=1_2_8					,  )
     torch.manual_seed(0     )
     __lowerCAmelCase								=	CLIPTextConfig(
         bos_token_id=0					,  eos_token_id=2					,  hidden_size=3_2					,  intermediate_size=3_7					,  layer_norm_eps=1e-05					,  num_attention_heads=4					,  num_hidden_layers=5					,  pad_token_id=1					,  vocab_size=1_0_0_0					,  hidden_act='gelu'					,  projection_dim=3_2					,  )
     __lowerCAmelCase								=	CLIPTextModel(lowerCAmelCase_     )
     __lowerCAmelCase								=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					,  local_files_only=lowerCAmelCase_     )

     __lowerCAmelCase								=	CLIPTextModelWithProjection(lowerCAmelCase_     )
     __lowerCAmelCase								=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					,  local_files_only=lowerCAmelCase_     )

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

     __lowerCAmelCase								=	self.get_dummy_inputs(lowerCAmelCase_     )
     __lowerCAmelCase								=	sd_pipe(**lowerCAmelCase_     ).images
     __lowerCAmelCase								=	image[0, -3:, -3:, -1]

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

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






  def 				lowercase		(				self :		Optional[int]     )							->						Optional[Any]:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3     )






  def 				lowercase		(				self :		int     )							->						Any:
     pass






  def 				lowercase		(				self :		Dict     )							->						Tuple:
     __lowerCAmelCase								=	self.get_dummy_components()
     __lowerCAmelCase								=	StableDiffusionXLImgaImgPipeline(**lowerCAmelCase_     )
     __lowerCAmelCase								=	sd_pipe.to(lowerCAmelCase_     )
     __lowerCAmelCase								=	sd_pipe.to(lowerCAmelCase_     )
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_     )

     # forward without prompt embeds
     __lowerCAmelCase								=	self.get_dummy_inputs(lowerCAmelCase_     )
     __lowerCAmelCase								=	3 * ['this is a negative prompt']
     __lowerCAmelCase								=	negative_prompt
     __lowerCAmelCase								=	3 * [inputs['prompt']]

     __lowerCAmelCase								=	sd_pipe(**lowerCAmelCase_     )
     __lowerCAmelCase								=	output.images[0, -3:, -3:, -1]

     # forward with prompt embeds
     __lowerCAmelCase								=	self.get_dummy_inputs(lowerCAmelCase_     )
     __lowerCAmelCase								=	3 * ['this is a negative prompt']
     __lowerCAmelCase								=	3 * [inputs.pop('prompt'     )]

     (
         (
         __lowerCAmelCase
     )	,	(
         __lowerCAmelCase
     )	,	(
         __lowerCAmelCase
     )	,	(
         __lowerCAmelCase
     )	,	
     )								=	sd_pipe.encode_prompt(lowerCAmelCase_					,  negative_prompt=lowerCAmelCase_     )

     __lowerCAmelCase								=	sd_pipe(
         **lowerCAmelCase_					,  prompt_embeds=lowerCAmelCase_					,  negative_prompt_embeds=lowerCAmelCase_					,  pooled_prompt_embeds=lowerCAmelCase_					,  negative_pooled_prompt_embeds=lowerCAmelCase_					,  )
     __lowerCAmelCase								=	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 				_UpperCAmelCase    (    unittest.TestCase    ):

  """simple docstring"""


  def 				lowercase		(				self :		int     )							->						str:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






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

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

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     __lowerCAmelCase								=	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








# 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


_snake_case  :    Dict						    =  {'configuration_mra': ['MRA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MraConfig']}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :    Optional[int]						    =  [
         '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

     _snake_case  :    Tuple						    =  _LazyModule(__name__, globals()['__file__'], _import_structure)

