



"""simple docstring"""
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 					__lowerCamelCase       (							a_					:		Union[str, Any]							)      ->				Tuple:
      __SCREAMING_SNAKE_CASE       :List[Any]		    =  [tensor.shape for tensor in tensor_list]
      return all(shape == shapes[0] for shape in shapes[1:]							)






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

       SCREAMING_SNAKE_CASE_       :				Tuple     						=  True




       @property
       def    _UpperCamelCase					(						self							)      ->				List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  1
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  4
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  (16, 16)

             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  floats_tensor((batch_size, num_channels) + sizes		,rng=random.Random(0							)							).to(_a							)
             return image




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  UNetaDConditionModel(
                 act_fn='''gelu'''		,attention_head_dim=8		,norm_num_groups=_a		,block_out_channels=[32, 32, 64, 64]		,time_cond_proj_dim=1_60		,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=_a		,only_cross_attention=_a		,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       :Any		    =  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       :List[str]		    =  EulerDiscreteScheduler(prediction_type='''sample'''							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  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='''quick_gelu'''		,projection_dim=5_12		,)
             __SCREAMING_SNAKE_CASE       :int		    =  CLIPTextModel(_a							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  {
                 """unet""": model.eval(),
                 """vae""": vae.eval(),
                 """scheduler""": scheduler,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
             }

             return components




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=0							)      ->				List[Any]:





             """simple docstring"""






             if str(_a							).startswith('''mps'''							):
                   __SCREAMING_SNAKE_CASE       :int		    =  torch.manual_seed(_a							)
             else:
                   __SCREAMING_SNAKE_CASE       :Tuple		    =  torch.Generator(device=_a							).manual_seed(_a							)
             __SCREAMING_SNAKE_CASE       :int		    =  {
                 """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    _UpperCamelCase					(						self							)      ->				List[str]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  """cpu"""

             __SCREAMING_SNAKE_CASE       :Dict		    =  self.get_dummy_components()
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.pipeline_class(**_a							)
             pipe.to(_a							)
             pipe.set_progress_bar_config(disable=_a							)

             __SCREAMING_SNAKE_CASE       :Any		    =  self.get_dummy_inputs(_a							)
             __SCREAMING_SNAKE_CASE       :int		    =  pipe(**_a							).images
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  image[0, -3:, -3:, -1]

             self.assertEqual(image.shape		,(1, 2_56, 2_56, 3)							)
             __SCREAMING_SNAKE_CASE       :List[str]		    =  np.array(
                 [0.4_7_2_2_2_4_1_2, 0.4_1_9_2_1_6_3_3, 0.4_4_7_1_7_4_3_4, 0.4_6_8_7_4_1_9_2, 0.4_2_5_8_8_2_5_8, 0.4_6_1_5_0_7_2_6, 0.4_6_7_7_5_3_4, 0.4_5_5_8_3_8_3_2, 0.4_8_5_7_9_0_5_5]							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  np.abs(image_slice.flatten() - expected_slice							).max()
             self.assertLessEqual(_a		,1E-3							)




       def    _UpperCamelCase					(						self							)      ->				Tuple:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Optional[Any]:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Dict:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Optional[Any]:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				List[Any]:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  [
                 """DDIMScheduler""",
                 """DDPMScheduler""",
                 """PNDMScheduler""",
                 """HeunDiscreteScheduler""",
                 """EulerAncestralDiscreteScheduler""",
                 """KDPM2DiscreteScheduler""",
                 """KDPM2AncestralDiscreteScheduler""",
                 """DPMSolverSDEScheduler""",
             ]
             __SCREAMING_SNAKE_CASE       :int		    =  self.get_dummy_components()
             __SCREAMING_SNAKE_CASE       :Dict		    =  self.pipeline_class(**_a							)

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

             pipe.to(_a							)
             pipe.set_progress_bar_config(disable=_a							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  self.get_dummy_inputs(_a							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  2

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

                   __SCREAMING_SNAKE_CASE       :str		    =  getattr(_a		,scheduler_enum.name							)
                   __SCREAMING_SNAKE_CASE       :Dict		    =  scheduler_cls.from_config(pipe.scheduler.config							)
                   __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  pipe(**_a							)[0]
                   outputs.append(_a							)

             assert check_same_shape(_a							)






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




       def    _UpperCamelCase					(						self							)      ->				List[Any]:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Any		    =  torch.manual_seed(33							)

             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  StableDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''		,torch_dtype=torch.floataa							)
             pipe.to('''cuda'''							)

             __SCREAMING_SNAKE_CASE       :str		    =  StableDiffusionLatentUpscalePipeline.from_pretrained(
                 '''stabilityai/sd-x2-latent-upscaler'''		,torch_dtype=torch.floataa							)
             upscaler.to('''cuda'''							)

             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  """a photo of an astronaut high resolution, unreal engine, ultra realistic"""

             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  pipe(_a		,generator=_a		,output_type='''latent'''							).images

             __SCREAMING_SNAKE_CASE       :str		    =  upscaler(
                 prompt=_a		,image=_a		,num_inference_steps=20		,guidance_scale=0		,generator=_a		,output_type='''np'''		,).images[0]

             __SCREAMING_SNAKE_CASE       :Tuple		    =  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    _UpperCamelCase					(						self							)      ->				Optional[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  torch.manual_seed(33							)

             __SCREAMING_SNAKE_CASE       :int		    =  StableDiffusionLatentUpscalePipeline.from_pretrained(
                 '''stabilityai/sd-x2-latent-upscaler'''		,torch_dtype=torch.floataa							)
             upscaler.to('''cuda'''							)

             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  """the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"""

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

             __SCREAMING_SNAKE_CASE       :List[str]		    =  upscaler(
                 prompt=_a		,image=_a		,num_inference_steps=20		,guidance_scale=0		,generator=_a		,output_type='''np'''		,).images[0]

             __SCREAMING_SNAKE_CASE       :str		    =  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






"""simple docstring"""
from collections import OrderedDict
from typing import Mapping

from packaging import version

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


lowerCamelCase_       =		logging.get_logger(__name__)

lowerCamelCase_       =		{
    "facebook/levit-128S": "https://huggingface.co/facebook/levit-128S/resolve/main/config.json",
    # See all LeViT models at https://huggingface.co/models?filter=levit
}






class  _SCREAMING_SNAKE_CASE(       A  ):
       SCREAMING_SNAKE_CASE_       :				List[Any]     						=  '''levit'''




       def __init__(						self		,SCREAMING_SNAKE_CASE__=2_24		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=1		,SCREAMING_SNAKE_CASE__=16		,SCREAMING_SNAKE_CASE__=[1_28, 2_56, 3_84]		,SCREAMING_SNAKE_CASE__=[4, 8, 12]		,SCREAMING_SNAKE_CASE__=[4, 4, 4]		,SCREAMING_SNAKE_CASE__=[16, 16, 16]		,SCREAMING_SNAKE_CASE__=0		,SCREAMING_SNAKE_CASE__=[2, 2, 2]		,SCREAMING_SNAKE_CASE__=[2, 2, 2]		,SCREAMING_SNAKE_CASE__=0.0_2		,**SCREAMING_SNAKE_CASE__		,)      ->				Tuple:





             """simple docstring"""






             super().__init__(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  image_size
             __SCREAMING_SNAKE_CASE       :Dict		    =  num_channels
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  kernel_size
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  stride
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  padding
             __SCREAMING_SNAKE_CASE       :Tuple		    =  hidden_sizes
             __SCREAMING_SNAKE_CASE       :List[str]		    =  num_attention_heads
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  depths
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  key_dim
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  drop_path_rate
             __SCREAMING_SNAKE_CASE       :Tuple		    =  patch_size
             __SCREAMING_SNAKE_CASE       :int		    =  attention_ratio
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  mlp_ratio
             __SCREAMING_SNAKE_CASE       :str		    =  initializer_range
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  [
                 ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
                 ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
             ]






class  _SCREAMING_SNAKE_CASE(       A  ):
       SCREAMING_SNAKE_CASE_       :				int     						=  version.parse('''1.11'''  )




       @property
       def    _UpperCamelCase					(						self							)      ->				Mapping[str, Mapping[int, str]]:





             """simple docstring"""






             return OrderedDict(
                 [
                     ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
                 ]							)




       @property
       def    _UpperCamelCase					(						self							)      ->				float:





             """simple docstring"""






             return 1E-4