






"""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 						_SCREAMING_SNAKE_CASE (    __snake_case			:  Optional[int]    ):





  '''simple docstring'''





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

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

    UpperCAmelCase_ :     List[Any]	     =True







    @property
    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       1
      lowercase							    =       4
      lowercase							    =       (1_6, 1_6)

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







    def 			UpperCamelCase_      (			self       ):
      torch.manual_seed(0       )
      lowercase							    =       UNetaDConditionModel(
          act_fn='gelu'							,		attention_head_dim=8							,		norm_num_groups=_lowerCamelCase							,		block_out_channels=[3_2, 3_2, 6_4, 6_4]							,		time_cond_proj_dim=1_6_0							,		conv_in_kernel=1							,		conv_out_kernel=1							,		cross_attention_dim=3_2							,		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')							,		)
      lowercase							    =       AutoencoderKL(
          block_out_channels=[3_2, 3_2, 6_4, 6_4]							,		in_channels=3							,		out_channels=3							,		down_block_types=[
              'DownEncoderBlock2D',
              'DownEncoderBlock2D',
              'DownEncoderBlock2D',
              'DownEncoderBlock2D',
          ]							,		up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']							,		latent_channels=4							,		)
      lowercase							    =       EulerDiscreteScheduler(prediction_type='sample'       )
      lowercase							    =       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='quick_gelu'							,		projection_dim=5_1_2							,		)
      lowercase							    =       CLIPTextModel(_lowerCamelCase       )
      lowercase							    =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

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

      return components







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase=0       ):
      if str(_lowerCamelCase       ).startswith('mps'       ):
        lowercase							    =       torch.manual_seed(_lowerCamelCase       )
      else:
        lowercase							    =       torch.Generator(device=_lowerCamelCase       ).manual_seed(_lowerCamelCase       )
      lowercase							    =       {
          '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       ):
      lowercase							    =       'cpu'

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

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

      self.assertEqual(image.shape							,		(1, 2_5_6, 2_5_6, 3)       )
      lowercase							    =       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]       )
      lowercase							    =       np.abs(image_slice.flatten() - expected_slice       ).max()
      self.assertLessEqual(_lowerCamelCase							,		1e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_attention_slicing_forward_pass(expected_max_diff=7e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_cpu_offload_forward_pass(expected_max_diff=3e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_inference_batch_single_identical(expected_max_diff=7e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_save_load_local(expected_max_difference=3e-3       )







    def 			UpperCamelCase_      (			self       ):
      super().test_save_load_optional_components(expected_max_difference=3e-3       )







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

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

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

      assert check_same_shape(_lowerCamelCase       )


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







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







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       torch.manual_seed(3_3       )

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

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

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

      lowercase							    =       pipe(_lowerCamelCase							,		generator=_lowerCamelCase							,		output_type='latent'       ).images

      lowercase							    =       upscaler(
          prompt=_lowerCamelCase							,		image=_lowerCamelCase							,		num_inference_steps=2_0							,		guidance_scale=0							,		generator=_lowerCamelCase							,		output_type='np'							,		).images[0]

      lowercase							    =       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       ):
      lowercase							    =       torch.manual_seed(3_3       )

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

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

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

      lowercase							    =       upscaler(
          prompt=_lowerCamelCase							,		image=_lowerCamelCase							,		num_inference_steps=2_0							,		guidance_scale=0							,		generator=_lowerCamelCase							,		output_type='np'							,		).images[0]

      lowercase							    =       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"""



# limitations under the License.


from typing import Optional, Tuple, Union

import torch

from diffusers import DiffusionPipeline, ImagePipelineOutput

class        a     (   a_	):







    def __init__(			self							,		_lowerCamelCase							,		_lowerCamelCase       ):
      super().__init__()
      self.register_modules(unet=_lowerCamelCase							,		scheduler=_lowerCamelCase       )







    @torch.no_grad()
    def __call__(			self							,		_lowerCamelCase = 1							,		_lowerCamelCase = None							,		_lowerCamelCase = 5_0							,		_lowerCamelCase = "pil"							,		_lowerCamelCase = True							,		**_lowerCamelCase							,		):
      lowercase							    =       torch.randn(
          (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)							,		generator=_lowerCamelCase							,		)
      lowercase							    =       image.to(self.device       )

      # set step values
      self.scheduler.set_timesteps(_lowerCamelCase       )

      for t in self.progress_bar(self.scheduler.timesteps       ):
        # 1. predict noise model_output
        lowercase							    =       self.unet(_lowerCamelCase							,		_lowerCamelCase       ).sample

        # 2. predict previous mean of image x_t-1 and add variance depending on eta
        # eta corresponds to η in paper and should be between [0, 1]
        # do x_t -> x_t-1
        lowercase							    =       self.scheduler.step(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       ).prev_sample

      lowercase							    =       (image / 2 + 0.5).clamp(0							,		1       )
      lowercase							    =       image.cpu().permute(0							,		2							,		3							,		1       ).numpy()
      if output_type == "pil":
        lowercase							    =       self.numpy_to_pil(_lowerCamelCase       )

      if not return_dict:
        return (image,), "This is a local test"

      return ImagePipelineOutput(images=_lowerCamelCase       ), "This is a local test"


