'''simple docstring'''




# flake8: noqa
# Lint as: python3
_a   :	Dict 							=   [
    "VerificationMode",
    "Version",
    "disable_progress_bar",
    "enable_progress_bar",
    "is_progress_bar_enabled",
    "experimental",
]

from .info_utils import VerificationMode
from .logging import disable_progress_bar, enable_progress_bar, is_progress_bar_enabled
from .version import Version
from .experimental import experimental







'''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 _lowercase	(      lowerCamelCase__						)					->       Union[str, Any]:





       """simple docstring"""


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






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

       snake_case       :Optional[int]				    =					True






       @property
       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Optional[int]       				=		1
              __UpperCAmelCase  :				Dict       				=		4
              __UpperCAmelCase  :				List[str]       				=		(16, 16)

              __UpperCAmelCase  :				Dict       				=		floats_tensor((batch_size, num_channels) + sizes				,     rng=random.Random(0      )      ).to(UpperCamelCase_      )
              return image






       def 			_snake_case						(      self      ):
              torch.manual_seed(0      )
              __UpperCAmelCase  :				List[str]       				=		UNetaDConditionModel(
                  act_fn="gelu"				,     attention_head_dim=8				,     norm_num_groups=UpperCamelCase_				,     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=UpperCamelCase_				,     only_cross_attention=UpperCamelCase_				,     out_channels=5				,     resnet_time_scale_shift="scale_shift"				,     time_embedding_type="fourier"				,     timestep_post_act="gelu"				,     up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D")				,     )
              __UpperCAmelCase  :				int       				=		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				,     )
              __UpperCAmelCase  :				Optional[int]       				=		EulerDiscreteScheduler(prediction_type="sample"      )
              __UpperCAmelCase  :				int       				=		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				,     )
              __UpperCAmelCase  :				List[str]       				=		CLIPTextModel(UpperCamelCase_      )
              __UpperCAmelCase  :				Tuple       				=		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

              __UpperCAmelCase  :				Union[str, Any]       				=		{
                  "unet": model.eval(),
                  "vae": vae.eval(),
                  "scheduler": scheduler,
                  "text_encoder": text_encoder,
                  "tokenizer": tokenizer,
              }

              return components






       def 			_snake_case						(      self				,     UpperCamelCase_				,     UpperCamelCase_=0      ):
              if str(UpperCamelCase_      ).startswith("mps"      ):
                     __UpperCAmelCase  :				str       				=		torch.manual_seed(UpperCamelCase_      )
              else:
                     __UpperCAmelCase  :				Optional[int]       				=		torch.Generator(device=UpperCamelCase_      ).manual_seed(UpperCamelCase_      )
              __UpperCAmelCase  :				Any       				=		{
                  "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 			_snake_case						(      self      ):
              __UpperCAmelCase  :				List[str]       				=		"cpu"

              __UpperCAmelCase  :				List[str]       				=		self.get_dummy_components()
              __UpperCAmelCase  :				Tuple       				=		self.pipeline_class(**UpperCamelCase_      )
              pipe.to(UpperCamelCase_      )
              pipe.set_progress_bar_config(disable=UpperCamelCase_      )

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

              self.assertEqual(image.shape				,     (1, 2_56, 2_56, 3)      )
              __UpperCAmelCase  :				Tuple       				=		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]      )
              __UpperCAmelCase  :				List[str]       				=		np.abs(image_slice.flatten() - expected_slice      ).max()
              self.assertLessEqual(UpperCamelCase_				,     1E-3      )






       def 			_snake_case						(      self      ):
              super().test_attention_slicing_forward_pass(expected_max_diff=7E-3      )






       def 			_snake_case						(      self      ):
              super().test_cpu_offload_forward_pass(expected_max_diff=3E-3      )






       def 			_snake_case						(      self      ):
              super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3      )






       def 			_snake_case						(      self      ):
              super().test_inference_batch_single_identical(expected_max_diff=7E-3      )






       def 			_snake_case						(      self      ):
              super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3      )






       def 			_snake_case						(      self      ):
              super().test_save_load_local(expected_max_difference=3E-3      )






       def 			_snake_case						(      self      ):
              super().test_save_load_optional_components(expected_max_difference=3E-3      )







       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Dict       				=		[
                  "DDIMScheduler",
                  "DDPMScheduler",
                  "PNDMScheduler",
                  "HeunDiscreteScheduler",
                  "EulerAncestralDiscreteScheduler",
                  "KDPM2DiscreteScheduler",
                  "KDPM2AncestralDiscreteScheduler",
                  "DPMSolverSDEScheduler",
              ]
              __UpperCAmelCase  :				Tuple       				=		self.get_dummy_components()
              __UpperCAmelCase  :				Union[str, Any]       				=		self.pipeline_class(**UpperCamelCase_      )

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

              pipe.to(UpperCamelCase_      )
              pipe.set_progress_bar_config(disable=UpperCamelCase_      )
              __UpperCAmelCase  :				Tuple       				=		self.get_dummy_inputs(UpperCamelCase_      )
              __UpperCAmelCase  :				List[str]       				=		2

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

                     __UpperCAmelCase  :				Optional[int]       				=		getattr(UpperCamelCase_				,     scheduler_enum.name      )
                     __UpperCAmelCase  :				List[str]       				=		scheduler_cls.from_config(pipe.scheduler.config      )
                     __UpperCAmelCase  :				Optional[int]       				=		pipe(**UpperCamelCase_      )[0]
                     outputs.append(UpperCamelCase_      )

              assert check_same_shape(UpperCamelCase_      )






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






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






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Optional[int]       				=		torch.manual_seed(33      )

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

              __UpperCAmelCase  :				Union[str, Any]       				=		StableDiffusionLatentUpscalePipeline.from_pretrained(
                  "stabilityai/sd-x2-latent-upscaler"				,     torch_dtype=torch.floataa      )
              upscaler.to("cuda"      )

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

              __UpperCAmelCase  :				Any       				=		pipe(UpperCamelCase_				,     generator=UpperCamelCase_				,     output_type="latent"      ).images

              __UpperCAmelCase  :				int       				=		upscaler(
                  prompt=UpperCamelCase_				,     image=UpperCamelCase_				,     num_inference_steps=20				,     guidance_scale=0				,     generator=UpperCamelCase_				,     output_type="np"				,     ).images[0]

              __UpperCAmelCase  :				Optional[Any]       				=		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 			_snake_case						(      self      ):
              __UpperCAmelCase  :				List[Any]       				=		torch.manual_seed(33      )

              __UpperCAmelCase  :				Union[str, Any]       				=		StableDiffusionLatentUpscalePipeline.from_pretrained(
                  "stabilityai/sd-x2-latent-upscaler"				,     torch_dtype=torch.floataa      )
              upscaler.to("cuda"      )

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

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

              __UpperCAmelCase  :				Dict       				=		upscaler(
                  prompt=UpperCamelCase_				,     image=UpperCamelCase_				,     num_inference_steps=20				,     guidance_scale=0				,     generator=UpperCamelCase_				,     output_type="np"				,     ).images[0]

              __UpperCAmelCase  :				Tuple       				=		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






