





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



def        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       Optional[int]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						[tensor.shape for tensor in tensor_list]
      return all(shape == shapes[0] for shape in shapes[1:]     )

class 		UpperCamelCase__   (				a__							,				a__							,				a__							,				unittest.TestCase       ):



     '''simple docstring'''

     __snake_case   :     int		     =      StableDiffusionLatentUpscalePipeline
     __snake_case   :     Any		     =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
         "height",
         "width",
         "cross_attention_kwargs",
         "negative_prompt_embeds",
         "prompt_embeds",
     }
     __snake_case   :     Optional[Any]		     =      PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
     __snake_case   :     int		     =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     __snake_case   :     List[str]		     =      frozenset(
         []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     __snake_case   :     Tuple		     =      frozenset([]       )

     __snake_case   :     Tuple		     =      True






     @property
     def     SCREAMING_SNAKE_CASE__  (							self						:   str	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						1
           SCREAMING_SNAKE_CASE		=						4
           SCREAMING_SNAKE_CASE		=						(16, 16)

           SCREAMING_SNAKE_CASE		=						floats_tensor((batch_size, num_channels) + sizes						,rng=random.Random(0	)	).to(_lowerCamelCase	)
           return image






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		Dict:




           '''simple docstring'''


           torch.manual_seed(0	)
           SCREAMING_SNAKE_CASE		=						UNetaDConditionModel(
               act_fn="""gelu"""						,attention_head_dim=8						,norm_num_groups=_lowerCamelCase						,block_out_channels=[32, 32, 64, 64]						,time_cond_proj_dim=160						,conv_in_kernel=1						,conv_out_kernel=1						,cross_attention_dim=32						,down_block_types=(
                   """KDownBlock2D""",
                   """KCrossAttnDownBlock2D""",
                   """KCrossAttnDownBlock2D""",
                   """KCrossAttnDownBlock2D""",
               )						,in_channels=8						,mid_block_type=_lowerCamelCase						,only_cross_attention=_lowerCamelCase						,out_channels=5						,resnet_time_scale_shift="""scale_shift"""						,time_embedding_type="""fourier"""						,timestep_post_act="""gelu"""						,up_block_types=("""KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KUpBlock2D""")						,)
           SCREAMING_SNAKE_CASE		=						AutoencoderKL(
               block_out_channels=[32, 32, 64, 64]						,in_channels=3						,out_channels=3						,down_block_types=[
                   """DownEncoderBlock2D""",
                   """DownEncoderBlock2D""",
                   """DownEncoderBlock2D""",
                   """DownEncoderBlock2D""",
               ]						,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,latent_channels=4						,)
           SCREAMING_SNAKE_CASE		=						EulerDiscreteScheduler(prediction_type="""sample"""	)
           SCREAMING_SNAKE_CASE		=						CLIPTextConfig(
               bos_token_id=0						,eos_token_id=2						,hidden_size=32						,intermediate_size=37						,layer_norm_eps=1e-0_5						,num_attention_heads=4						,num_hidden_layers=5						,pad_token_id=1						,vocab_size=1000						,hidden_act="""quick_gelu"""						,projection_dim=512						,)
           SCREAMING_SNAKE_CASE		=						CLIPTextModel(_lowerCamelCase	)
           SCREAMING_SNAKE_CASE		=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

           SCREAMING_SNAKE_CASE		=						{
               '''unet''': model.eval(),
               '''vae''': vae.eval(),
               '''scheduler''': scheduler,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
           }

           return components






     def     SCREAMING_SNAKE_CASE__  (							self						:   int						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   Dict=0	)					->		List[Any]:




           '''simple docstring'''


           if str(_lowerCamelCase	).startswith("""mps"""	):
                 SCREAMING_SNAKE_CASE		=						torch.manual_seed(_lowerCamelCase	)
           else:
                 SCREAMING_SNAKE_CASE		=						torch.Generator(device=_lowerCamelCase	).manual_seed(_lowerCamelCase	)
           SCREAMING_SNAKE_CASE		=						{
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''image''': self.dummy_image.cpu(),
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''output_type''': '''numpy''',
           }
           return inputs






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						'''cpu'''

           SCREAMING_SNAKE_CASE		=						self.get_dummy_components()
           SCREAMING_SNAKE_CASE		=						self.pipeline_class(**_lowerCamelCase	)
           pipe.to(_lowerCamelCase	)
           pipe.set_progress_bar_config(disable=_lowerCamelCase	)

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(_lowerCamelCase	)
           SCREAMING_SNAKE_CASE		=						pipe(**_lowerCamelCase	).images
           SCREAMING_SNAKE_CASE		=						image[0, -3:, -3:, -1]

           self.assertEqual(image.shape						,(1, 256, 256, 3)	)
           SCREAMING_SNAKE_CASE		=						np.array(
               [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055]	)
           SCREAMING_SNAKE_CASE		=						np.abs(image_slice.flatten() - expected_slice	).max()
           self.assertLessEqual(_lowerCamelCase						,1e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		List[str]:




           '''simple docstring'''


           super().test_attention_slicing_forward_pass(expected_max_diff=7e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		int:




           '''simple docstring'''


           super().test_cpu_offload_forward_pass(expected_max_diff=3e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		Dict:




           '''simple docstring'''


           super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		Union[str, Any]:




           '''simple docstring'''


           super().test_inference_batch_single_identical(expected_max_diff=7e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		str:




           '''simple docstring'''


           super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		Union[str, Any]:




           '''simple docstring'''


           super().test_save_load_local(expected_max_difference=3e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		Union[str, Any]:




           '''simple docstring'''


           super().test_save_load_optional_components(expected_max_difference=3e-3	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						[
               '''DDIMScheduler''',
               '''DDPMScheduler''',
               '''PNDMScheduler''',
               '''HeunDiscreteScheduler''',
               '''EulerAncestralDiscreteScheduler''',
               '''KDPM2DiscreteScheduler''',
               '''KDPM2AncestralDiscreteScheduler''',
               '''DPMSolverSDEScheduler''',
           ]
           SCREAMING_SNAKE_CASE		=						self.get_dummy_components()
           SCREAMING_SNAKE_CASE		=						self.pipeline_class(**_lowerCamelCase	)

           # make sure that PNDM does not need warm-up
           pipe.scheduler.register_to_config(skip_prk_steps=_lowerCamelCase	)

           pipe.to(_lowerCamelCase	)
           pipe.set_progress_bar_config(disable=_lowerCamelCase	)
           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(_lowerCamelCase	)
           SCREAMING_SNAKE_CASE		=						2

           SCREAMING_SNAKE_CASE		=						[]
           for scheduler_enum in KarrasDiffusionSchedulers:
                 if scheduler_enum.name in skip_schedulers:
                       # no sigma schedulers are not supported
                       # no schedulers
                       continue

                 SCREAMING_SNAKE_CASE		=						getattr(_lowerCamelCase						,scheduler_enum.name	)
                 SCREAMING_SNAKE_CASE		=						scheduler_cls.from_config(pipe.scheduler.config	)
                 SCREAMING_SNAKE_CASE		=						pipe(**_lowerCamelCase	)[0]
                 outputs.append(_lowerCamelCase	)

           assert check_same_shape(_lowerCamelCase	)

@require_torch_gpu
@slow
class 		UpperCamelCase__   (				unittest.TestCase       ):



     '''simple docstring'''

     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[int]	)					->		List[Any]:




           '''simple docstring'''


           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()






     def     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						torch.manual_seed(33	)

           SCREAMING_SNAKE_CASE		=						StableDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""						,torch_dtype=torch.floataa	)
           pipe.to("""cuda"""	)

           SCREAMING_SNAKE_CASE		=						StableDiffusionLatentUpscalePipeline.from_pretrained(
               """stabilityai/sd-x2-latent-upscaler"""						,torch_dtype=torch.floataa	)
           upscaler.to("""cuda"""	)

           SCREAMING_SNAKE_CASE		=						'''a photo of an astronaut high resolution, unreal engine, ultra realistic'''

           SCREAMING_SNAKE_CASE		=						pipe(_lowerCamelCase						,generator=_lowerCamelCase						,output_type="""latent"""	).images

           SCREAMING_SNAKE_CASE		=						upscaler(
               prompt=_lowerCamelCase						,image=_lowerCamelCase						,num_inference_steps=20						,guidance_scale=0						,generator=_lowerCamelCase						,output_type="""np"""						,).images[0]

           SCREAMING_SNAKE_CASE		=						load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy"""	)
           assert np.abs((expected_image - image).mean()	) < 5e-2






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						torch.manual_seed(33	)

           SCREAMING_SNAKE_CASE		=						StableDiffusionLatentUpscalePipeline.from_pretrained(
               """stabilityai/sd-x2-latent-upscaler"""						,torch_dtype=torch.floataa	)
           upscaler.to("""cuda"""	)

           SCREAMING_SNAKE_CASE		=						'''the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'''

           SCREAMING_SNAKE_CASE		=						load_image(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png"""	)

           SCREAMING_SNAKE_CASE		=						upscaler(
               prompt=_lowerCamelCase						,image=_lowerCamelCase						,num_inference_steps=20						,guidance_scale=0						,generator=_lowerCamelCase						,output_type="""np"""						,).images[0]

           SCREAMING_SNAKE_CASE		=						load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy"""	)
           assert np.abs((expected_image - image).max()	) < 5e-2









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE_						       =	{
    """configuration_mvp""": ["""MVP_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MvpConfig""", """MvpOnnxConfig"""],
    """tokenization_mvp""": ["""MvpTokenizer"""],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	["""MvpTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	[
          """MVP_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MvpForCausalLM""",
          """MvpForConditionalGeneration""",
          """MvpForQuestionAnswering""",
          """MvpForSequenceClassification""",
          """MvpModel""",
          """MvpPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
      from .tokenization_mvp import MvpTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_mvp_fast import MvpTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_mvp import (
                MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
                MvpForCausalLM,
                MvpForConditionalGeneration,
                MvpForQuestionAnswering,
                MvpForSequenceClassification,
                MvpModel,
                MvpPreTrainedModel,
            )

else:
      import sys

      SCREAMING_SNAKE_CASE_						       =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


