

"""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        _snake_case   (			UpperCAmelCase_     :    Optional[int]	):
     A__								=     [tensor.shape for tensor in tensor_list]
     return all(shape == shapes[0] for shape in shapes[1:]	)



class     a  (				_lowerCamelCase,		_lowerCamelCase,		_lowerCamelCase,		unittest.TestCase ):



    """simple docstring"""



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

    UpperCAmelCase         =    True

    @property
    def  UpperCamelCase		(   self:       Union[str, Any]				):






         """simple docstring"""


         A__								=     1
         A__								=     4
         A__								=     (16, 16)

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

    def  UpperCamelCase		(   self:       int				):






         """simple docstring"""


         torch.manual_seed(0				)
         A__								=     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""")			,  )
         A__								=     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			,  )
         A__								=     EulerDiscreteScheduler(prediction_type="""sample"""				)
         A__								=     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=10_00			,  hidden_act="""quick_gelu"""			,  projection_dim=5_12			,  )
         A__								=     CLIPTextModel(UpperCamelCase				)
         A__								=     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

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

         return components

    def  UpperCamelCase		(   self:       Optional[Any]			,  UpperCamelCase:       Optional[int]			,  UpperCamelCase:       Optional[int]=0				):






         """simple docstring"""


         if str(UpperCamelCase				).startswith("""mps"""				):
              A__								=     torch.manual_seed(UpperCamelCase				)
         else:
              A__								=     torch.Generator(device=UpperCamelCase				).manual_seed(UpperCamelCase				)
         A__								=     {
             '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:       Tuple				):






         """simple docstring"""


         A__								=     'cpu'

         A__								=     self.get_dummy_components()
         A__								=     self.pipeline_class(**UpperCamelCase				)
         pipe.to(UpperCamelCase				)
         pipe.set_progress_bar_config(disable=UpperCamelCase				)

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

         self.assertEqual(image.shape			,  (1, 2_56, 2_56, 3)				)
         A__								=     np.array(
             [0.47_222_412, 0.41_921_633, 0.44_717_434, 0.46_874_192, 0.42_588_258, 0.46_150_726, 0.4_677_534, 0.45_583_832, 0.48_579_055]				)
         A__								=     np.abs(image_slice.flatten() - expected_slice				).max()
         self.assertLessEqual(UpperCamelCase			,  1e-3				)

    def  UpperCamelCase		(   self:       Union[str, Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       int				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Union[str, Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Optional[Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Any				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Union[str, Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Union[str, Any]				):






         """simple docstring"""


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





    def  UpperCamelCase		(   self:       int				):






         """simple docstring"""


         A__								=     [
             'DDIMScheduler',
             'DDPMScheduler',
             'PNDMScheduler',
             'HeunDiscreteScheduler',
             'EulerAncestralDiscreteScheduler',
             'KDPM2DiscreteScheduler',
             'KDPM2AncestralDiscreteScheduler',
             'DPMSolverSDEScheduler',
         ]
         A__								=     self.get_dummy_components()
         A__								=     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				)
         A__								=     self.get_dummy_inputs(UpperCamelCase				)
         A__								=     2

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

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

         assert check_same_shape(UpperCamelCase				)



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



    """simple docstring"""



    def  UpperCamelCase		(   self:       List[Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Dict				):






         """simple docstring"""


         A__								=     torch.manual_seed(33				)

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

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

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

         A__								=     pipe(UpperCamelCase			,  generator=UpperCamelCase			,  output_type="""latent"""				).images

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

         A__								=     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:       List[Any]				):






         """simple docstring"""


         A__								=     torch.manual_seed(33				)

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

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

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

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

         A__								=     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 List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class      A__							(     __magic_name__	):
				lowercase			    =  42
				lowercase			    =  42






				def __init__(						self					:      Any    ,     a					:      UNetaDModel    ,     a					:      ScoreSdeVeScheduler   ):






							'''simple docstring'''







							super().__init__()
							self.register_modules(unet=a    ,     scheduler=a   )






				@torch.no_grad()
				def __call__(						self					:      List[str]    ,     a					:      int = 1    ,     a					:      int = 2_000    ,     a					:      Optional[Union[torch.Generator, List[torch.Generator]]] = None    ,     a					:      Optional[str] = "pil"    ,     a					:      bool = True    ,     **a					:      int    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[str]         =		self.unet.config.sample_size
							lowerCAmelCase__      :      Union[str, Any]         =		(batch_size, 3, img_size, img_size)

							lowerCAmelCase__      :      Tuple         =		self.unet

							lowerCAmelCase__      :      Optional[Any]         =		randn_tensor(a    ,     generator=a   ) * self.scheduler.init_noise_sigma
							lowerCAmelCase__      :      int         =		sample.to(self.device   )

							self.scheduler.set_timesteps(a   )
							self.scheduler.set_sigmas(a   )

							for i, t in enumerate(self.progress_bar(self.scheduler.timesteps   )   ):
										lowerCAmelCase__      :      Dict         =		self.scheduler.sigmas[i] * torch.ones(shape[0]    ,     device=self.device   )

										# correction step
										for _ in range(self.scheduler.config.correct_steps   ):
													lowerCAmelCase__      :      Optional[Any]         =		self.unet(a    ,     a   ).sample
													lowerCAmelCase__      :      Dict         =		self.scheduler.step_correct(a    ,     a    ,     generator=a   ).prev_sample

										# prediction step
										lowerCAmelCase__      :      Optional[int]         =		model(a    ,     a   ).sample
										lowerCAmelCase__      :      Optional[Any]         =		self.scheduler.step_pred(a    ,     a    ,     a    ,     generator=a   )

										lowerCAmelCase__      ,				lowerCAmelCase__      :      str         =		output.prev_sample, output.prev_sample_mean

							lowerCAmelCase__      :      Any         =		sample_mean.clamp(0    ,     1   )
							lowerCAmelCase__      :      List[str]         =		sample.cpu().permute(0    ,     2    ,     3    ,     1   ).numpy()
							if output_type == "pil":
										lowerCAmelCase__      :      int         =		self.numpy_to_pil(a   )

							if not return_dict:
										return (sample,)

							return ImagePipelineOutput(images=a   )