

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 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		Dict       ):





    """simple docstring"""


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



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

  __lowerCAmelCase     :	Tuple						 =  True

  @property
  def 						lowerCAmelCase__     (					self):
      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(_A)
      return image

  def 						lowerCAmelCase__     (					self):
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_									=						UNetaDConditionModel(
          act_fn='gelu'						,		attention_head_dim=8						,		norm_num_groups=_A						,		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=_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_									=						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-05						,		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(_A)
      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 						lowerCAmelCase__     (					self						,		_A						,		_A=0):
      if str(_A).startswith('mps'):
          SCREAMING_SNAKE_CASE_									=						torch.manual_seed(_A)
      else:
          SCREAMING_SNAKE_CASE_									=						torch.Generator(device=_A).manual_seed(_A)
      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 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						'cpu'

      SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
      SCREAMING_SNAKE_CASE_									=						self.pipeline_class(**_A)
      pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)

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

      self.assertEqual(image.shape						,		(1, 256, 256, 3))
      SCREAMING_SNAKE_CASE_									=						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_									=						np.abs(image_slice.flatten() - expected_slice).max()
      self.assertLessEqual(_A						,		1E-3)

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

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

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

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

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

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

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


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						[
          'DDIMScheduler',
          'DDPMScheduler',
          'PNDMScheduler',
          'HeunDiscreteScheduler',
          'EulerAncestralDiscreteScheduler',
          'KDPM2DiscreteScheduler',
          'KDPM2AncestralDiscreteScheduler',
          'DPMSolverSDEScheduler',
      ]
      SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
      SCREAMING_SNAKE_CASE_									=						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_									=						self.get_dummy_inputs(_A)
      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(_A						,		scheduler_enum.name)
          SCREAMING_SNAKE_CASE_									=						scheduler_cls.from_config(pipe.scheduler.config)
          SCREAMING_SNAKE_CASE_									=						pipe(**_A)[0]
          outputs.append(_A)

      assert check_same_shape(_A)






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

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

  def 						lowerCAmelCase__     (					self):
      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(_A						,		generator=_A						,		output_type='latent').images

      SCREAMING_SNAKE_CASE_									=						upscaler(
          prompt=_A						,		image=_A						,		num_inference_steps=20						,		guidance_scale=0						,		generator=_A						,		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 						lowerCAmelCase__     (					self):
      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=_A						,		image=_A						,		num_inference_steps=20						,		guidance_scale=0						,		generator=_A						,		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 Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput



class       __snake_case  (    lowerCAmelCase__       ):

  def __init__(					self						,		_A						,		_A						,		_A						,		_A = None						,		):
      super().__init__()
      self.register_modules(transformer=_A						,		vae=_A						,		scheduler=_A)

      # create a imagenet -> id dictionary for easier use
      SCREAMING_SNAKE_CASE_									=						{}
      if idalabel is not None:
          for key, value in idalabel.items():
              for label in value.split(','):
                  SCREAMING_SNAKE_CASE_									=						int(_A)
          SCREAMING_SNAKE_CASE_									=						dict(sorted(self.labels.items()))

  def 						lowerCAmelCase__     (					self						,		_A):

      if not isinstance(_A						,		_A):
          SCREAMING_SNAKE_CASE_									=						list(_A)

      for l in label:
          if l not in self.labels:
              raise ValueError(
                  f"""{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}.""")

      return [self.labels[l] for l in label]


  @torch.no_grad()
  def __call__(					self						,		_A						,		_A = 4.0						,		_A = None						,		_A = 50						,		_A = "pil"						,		_A = True						,		):
      SCREAMING_SNAKE_CASE_									=						len(_A)
      SCREAMING_SNAKE_CASE_									=						self.transformer.config.sample_size
      SCREAMING_SNAKE_CASE_									=						self.transformer.config.in_channels

      SCREAMING_SNAKE_CASE_									=						randn_tensor(
          shape=(batch_size, latent_channels, latent_size, latent_size)						,		generator=_A						,		device=self.device						,		dtype=self.transformer.dtype						,		)
      SCREAMING_SNAKE_CASE_									=						torch.cat([latents] * 2) if guidance_scale > 1 else latents

      SCREAMING_SNAKE_CASE_									=						torch.tensor(_A						,		device=self.device).reshape(-1)
      SCREAMING_SNAKE_CASE_									=						torch.tensor([1000] * batch_size						,		device=self.device)
      SCREAMING_SNAKE_CASE_									=						torch.cat([class_labels, class_null]						,		0) if guidance_scale > 1 else class_labels

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

      for t in self.progress_bar(self.scheduler.timesteps):
          if guidance_scale > 1:
              SCREAMING_SNAKE_CASE_									=						latent_model_input[: len(_A) // 2]
              SCREAMING_SNAKE_CASE_									=						torch.cat([half, half]						,		dim=0)
          SCREAMING_SNAKE_CASE_									=						self.scheduler.scale_model_input(_A						,		_A)

          SCREAMING_SNAKE_CASE_									=						t
          if not torch.is_tensor(_A):
              # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
              # This would be a good case for the `match` statement (Python 3.10+)
              SCREAMING_SNAKE_CASE_									=						latent_model_input.device.type == 'mps'
              if isinstance(_A						,		_A):
                  SCREAMING_SNAKE_CASE_									=						torch.floataa if is_mps else torch.floataa
              else:
                  SCREAMING_SNAKE_CASE_									=						torch.intaa if is_mps else torch.intaa
              SCREAMING_SNAKE_CASE_									=						torch.tensor([timesteps]						,		dtype=_A						,		device=latent_model_input.device)
          elif len(timesteps.shape) == 0:
              SCREAMING_SNAKE_CASE_									=						timesteps[None].to(latent_model_input.device)
          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          SCREAMING_SNAKE_CASE_									=						timesteps.expand(latent_model_input.shape[0])
          # predict noise model_output
          SCREAMING_SNAKE_CASE_									=						self.transformer(
              _A						,		timestep=_A						,		class_labels=_A).sample

          # perform guidance
          if guidance_scale > 1:
              SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
              SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						torch.split(_A						,		len(_A) // 2						,		dim=0)

              SCREAMING_SNAKE_CASE_									=						uncond_eps + guidance_scale * (cond_eps - uncond_eps)
              SCREAMING_SNAKE_CASE_									=						torch.cat([half_eps, half_eps]						,		dim=0)

              SCREAMING_SNAKE_CASE_									=						torch.cat([eps, rest]						,		dim=1)

          # learned sigma
          if self.transformer.config.out_channels // 2 == latent_channels:
              SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						torch.split(_A						,		_A						,		dim=1)
          else:
              SCREAMING_SNAKE_CASE_									=						noise_pred

          # compute previous image: x_t -> x_t-1
          SCREAMING_SNAKE_CASE_									=						self.scheduler.step(_A						,		_A						,		_A).prev_sample

      if guidance_scale > 1:
          SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						latent_model_input.chunk(2						,		dim=0)
      else:
          SCREAMING_SNAKE_CASE_									=						latent_model_input

      SCREAMING_SNAKE_CASE_									=						1 / self.vae.config.scaling_factor * latents
      SCREAMING_SNAKE_CASE_									=						self.vae.decode(_A).sample

      SCREAMING_SNAKE_CASE_									=						(samples / 2 + 0.5).clamp(0						,		1)

      # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
      SCREAMING_SNAKE_CASE_									=						samples.cpu().permute(0						,		2						,		3						,		1).float().numpy()

      if output_type == "pil":
          SCREAMING_SNAKE_CASE_									=						self.numpy_to_pil(_A)

      if not return_dict:
          return (samples,)

      return ImagePipelineOutput(images=_A)


