import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
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


enable_full_determinism()





class      A							(				unittest.TestCase     ):


   def 			lowercase_      (self     :						Any		)						->				str:
         """simple docstring"""



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


   @property
   def 			lowercase_      (self     :						int		)						->				Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  1
         UpperCAmelCase__								=  3
         UpperCAmelCase__								=  (3_2, 3_2)

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


   @property
   def 			lowercase_      (self     :						Any		)						->				Union[str, Any]:
         """simple docstring"""



         torch.manual_seed(0		)
         UpperCAmelCase__								=  UNetaDConditionModel(
             block_out_channels=(3_2, 3_2, 6_4)       ,  layers_per_block=2       ,  sample_size=3_2       ,  in_channels=7       ,  out_channels=4       ,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")       ,  up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")       ,  cross_attention_dim=3_2       ,  attention_head_dim=8       ,  use_linear_projection=lowerCAmelCase__       ,  only_cross_attention=(True, True, False)       ,  num_class_embeds=1_0_0       ,  )
         return model


   @property
   def 			lowercase_      (self     :						Any		)						->				Dict:
         """simple docstring"""



         torch.manual_seed(0		)
         UpperCAmelCase__								=  AutoencoderKL(
             block_out_channels=[3_2, 3_2, 6_4]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]       ,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]       ,  latent_channels=4       ,  )
         return model


   @property
   def 			lowercase_      (self     :						Optional[Any]		)						->				int:
         """simple docstring"""



         torch.manual_seed(0		)
         UpperCAmelCase__								=  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="gelu"       ,  projection_dim=5_1_2       ,  )
         return CLIPTextModel(lowerCAmelCase__		)


   def 			lowercase_      (self     :						int		)						->				Any:
         """simple docstring"""



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.dummy_cond_unet_upscale
         UpperCAmelCase__								=  DDPMScheduler()
         UpperCAmelCase__								=  DDIMScheduler(prediction_type="v_prediction"		)
         UpperCAmelCase__								=  self.dummy_vae
         UpperCAmelCase__								=  self.dummy_text_encoder
         UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

         UpperCAmelCase__								=  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1		)[0]
         UpperCAmelCase__								=  Image.fromarray(np.uinta(lowerCAmelCase__		)		).convert("RGB"		).resize((6_4, 6_4)		)

         # make sure here that pndm scheduler skips prk
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline(
             unet=lowerCAmelCase__       ,  low_res_scheduler=lowerCAmelCase__       ,  scheduler=lowerCAmelCase__       ,  vae=lowerCAmelCase__       ,  text_encoder=lowerCAmelCase__       ,  tokenizer=lowerCAmelCase__       ,  max_noise_level=3_5_0       ,  )
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

         UpperCAmelCase__								=  "A painting of a squirrel eating a burger"
         UpperCAmelCase__								=  torch.Generator(device=lowerCAmelCase__		).manual_seed(0		)
         UpperCAmelCase__								=  sd_pipe(
             [prompt]       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  guidance_scale=6.0       ,  noise_level=2_0       ,  num_inference_steps=2       ,  output_type="np"       ,  )

         UpperCAmelCase__								=  output.images

         UpperCAmelCase__								=  torch.Generator(device=lowerCAmelCase__		).manual_seed(0		)
         UpperCAmelCase__								=  sd_pipe(
             [prompt]       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  guidance_scale=6.0       ,  noise_level=2_0       ,  num_inference_steps=2       ,  output_type="np"       ,  return_dict=lowerCAmelCase__       ,  )[0]

         UpperCAmelCase__								=  image[0, -3:, -3:, -1]
         UpperCAmelCase__								=  image_from_tuple[0, -3:, -3:, -1]

         UpperCAmelCase__								=  low_res_image.size[0] * 4
         assert image.shape == (1, expected_height_width, expected_height_width, 3)
         UpperCAmelCase__								=  np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]		)

         assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2


   def 			lowercase_      (self     :						int		)						->				Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.dummy_cond_unet_upscale
         UpperCAmelCase__								=  DDPMScheduler()
         UpperCAmelCase__								=  DDIMScheduler(prediction_type="v_prediction"		)
         UpperCAmelCase__								=  self.dummy_vae
         UpperCAmelCase__								=  self.dummy_text_encoder
         UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

         UpperCAmelCase__								=  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1		)[0]
         UpperCAmelCase__								=  Image.fromarray(np.uinta(lowerCAmelCase__		)		).convert("RGB"		).resize((6_4, 6_4)		)

         # make sure here that pndm scheduler skips prk
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline(
             unet=lowerCAmelCase__       ,  low_res_scheduler=lowerCAmelCase__       ,  scheduler=lowerCAmelCase__       ,  vae=lowerCAmelCase__       ,  text_encoder=lowerCAmelCase__       ,  tokenizer=lowerCAmelCase__       ,  max_noise_level=3_5_0       ,  )
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

         UpperCAmelCase__								=  "A painting of a squirrel eating a burger"
         UpperCAmelCase__								=  sd_pipe(
             2 * [prompt]       ,  image=2 * [low_res_image]       ,  guidance_scale=6.0       ,  noise_level=2_0       ,  num_inference_steps=2       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  output.images
         assert image.shape[0] == 2

         UpperCAmelCase__								=  torch.Generator(device=lowerCAmelCase__		).manual_seed(0		)
         UpperCAmelCase__								=  sd_pipe(
             [prompt]       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  num_images_per_prompt=2       ,  guidance_scale=6.0       ,  noise_level=2_0       ,  num_inference_steps=2       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  output.images
         assert image.shape[0] == 2





   @unittest.skipIf(torch_device != "cuda"       ,  "This test requires a GPU"		)
   def 			lowercase_      (self     :						str		)						->				Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  self.dummy_cond_unet_upscale
         UpperCAmelCase__								=  DDPMScheduler()
         UpperCAmelCase__								=  DDIMScheduler(prediction_type="v_prediction"		)
         UpperCAmelCase__								=  self.dummy_vae
         UpperCAmelCase__								=  self.dummy_text_encoder
         UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

         UpperCAmelCase__								=  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1		)[0]
         UpperCAmelCase__								=  Image.fromarray(np.uinta(lowerCAmelCase__		)		).convert("RGB"		).resize((6_4, 6_4)		)

         # put models in fp16, except vae as it overflows in fp16
         UpperCAmelCase__								=  unet.half()
         UpperCAmelCase__								=  text_encoder.half()

         # make sure here that pndm scheduler skips prk
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline(
             unet=lowerCAmelCase__       ,  low_res_scheduler=lowerCAmelCase__       ,  scheduler=lowerCAmelCase__       ,  vae=lowerCAmelCase__       ,  text_encoder=lowerCAmelCase__       ,  tokenizer=lowerCAmelCase__       ,  max_noise_level=3_5_0       ,  )
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

         UpperCAmelCase__								=  "A painting of a squirrel eating a burger"
         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  sd_pipe(
             [prompt]       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  num_inference_steps=2       ,  output_type="np"       ,  ).images

         UpperCAmelCase__								=  low_res_image.size[0] * 4
         assert image.shape == (1, expected_height_width, expected_height_width, 3)





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


   def 			lowercase_      (self     :						Optional[int]		)						->				Optional[int]:
         """simple docstring"""



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


   def 			lowercase_      (self     :						List[str]		)						->				str:
         """simple docstring"""



         UpperCAmelCase__								=  load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/sd2-upscale/low_res_cat.png"		)
         UpperCAmelCase__								=  load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
             "/upsampled_cat.npy"		)

         UpperCAmelCase__								=  "stabilityai/stable-diffusion-x4-upscaler"
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline.from_pretrained(lowerCAmelCase__		)
         pipe.to(lowerCAmelCase__		)
         pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
         pipe.enable_attention_slicing()

         UpperCAmelCase__								=  "a cat sitting on a park bench"

         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  pipe(
             prompt=lowerCAmelCase__       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  output.images[0]

         assert image.shape == (5_1_2, 5_1_2, 3)
         assert np.abs(expected_image - image		).max() < 1E-3


   def 			lowercase_      (self     :						Any		)						->				Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/sd2-upscale/low_res_cat.png"		)
         UpperCAmelCase__								=  load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
             "/upsampled_cat_fp16.npy"		)

         UpperCAmelCase__								=  "stabilityai/stable-diffusion-x4-upscaler"
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline.from_pretrained(
             lowerCAmelCase__       ,  torch_dtype=torch.floataa       ,  )
         pipe.to(lowerCAmelCase__		)
         pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
         pipe.enable_attention_slicing()

         UpperCAmelCase__								=  "a cat sitting on a park bench"

         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  pipe(
             prompt=lowerCAmelCase__       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  output_type="np"       ,  )
         UpperCAmelCase__								=  output.images[0]

         assert image.shape == (5_1_2, 5_1_2, 3)
         assert np.abs(expected_image - image		).max() < 5E-1





   def 			lowercase_      (self     :						Tuple		)						->				Optional[Any]:
         """simple docstring"""



         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         UpperCAmelCase__								=  load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/sd2-upscale/low_res_cat.png"		)

         UpperCAmelCase__								=  "stabilityai/stable-diffusion-x4-upscaler"
         UpperCAmelCase__								=  StableDiffusionUpscalePipeline.from_pretrained(
             lowerCAmelCase__       ,  torch_dtype=torch.floataa       ,  )
         pipe.to(lowerCAmelCase__		)
         pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
         pipe.enable_attention_slicing(1		)
         pipe.enable_sequential_cpu_offload()

         UpperCAmelCase__								=  "a cat sitting on a park bench"

         UpperCAmelCase__								=  torch.manual_seed(0		)
         UpperCAmelCase__								=  pipe(
             prompt=lowerCAmelCase__       ,  image=lowerCAmelCase__       ,  generator=lowerCAmelCase__       ,  num_inference_steps=5       ,  output_type="np"       ,  )

         UpperCAmelCase__								=  torch.cuda.max_memory_allocated()
         # make sure that less than 2.9 GB is allocated
         assert mem_bytes < 2.9 * 1_0**9










"""simple docstring"""




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class        __UpperCamelCase				(						a__		,       a__						):






							@register_to_config
							def __init__(   self      ,					*,
							    lowerCAmelCase__ = 4      ,					lowerCAmelCase__ = 768      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					)	-> Optional[Any]:
								super().__init__()

								a      :		Tuple	=				nn.Parameter(torch.zeros(lowerCAmelCase__      )      )

								# parameters for additional clip time embeddings
								a      :		str	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )
								a      :		Any	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )

								# parameters for encoder hidden states
								a      :		int	=				clip_extra_context_tokens
								a      :		int	=				nn.Linear(
								    lowerCAmelCase__      ,					self.clip_extra_context_tokens * cross_attention_dim      )
								a      :		Any	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )
								a      :		str	=				nn.LayerNorm(lowerCAmelCase__      )






							def     __a   (   self      ,					*, lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )	-> Union[str, Any]:
								if do_classifier_free_guidance:
									# Add the classifier free guidance embeddings to the image embeddings
									a      :		str	=				image_embeddings.shape[0]
									a      :		Optional[int]	=				self.learned_classifier_free_guidance_embeddings.unsqueeze(0      )
									a      :		Any	=				classifier_free_guidance_embeddings.expand(
									    lowerCAmelCase__      ,					-1      )
									a      :		Any	=				torch.cat([classifier_free_guidance_embeddings, image_embeddings]      ,					dim=0      )

								# The image embeddings batch size and the text embeddings batch size are equal
								assert image_embeddings.shape[0] == prompt_embeds.shape[0]

								a      :		List[str]	=				prompt_embeds.shape[0]

								# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
								# adding CLIP embeddings to the existing timestep embedding, ...
								a      :		Dict	=				self.embedding_proj(lowerCAmelCase__      )
								a      :		List[str]	=				self.clip_image_embeddings_project_to_time_embeddings(lowerCAmelCase__      )
								a      :		Dict	=				time_projected_image_embeddings + time_projected_prompt_embeds

								# ... and by projecting CLIP embeddings into four
								# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
								a      :		Union[str, Any]	=				self.clip_extra_context_tokens_proj(lowerCAmelCase__      )
								a      :		List[str]	=				clip_extra_context_tokens.reshape(lowerCAmelCase__      ,					-1      ,					self.clip_extra_context_tokens      )
								a      :		Optional[Any]	=				clip_extra_context_tokens.permute(0      ,					2      ,					1      )

								a      :		Optional[int]	=				self.encoder_hidden_states_proj(lowerCAmelCase__      )
								a      :		str	=				self.text_encoder_hidden_states_norm(lowerCAmelCase__      )
								a      :		List[str]	=				torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]      ,					dim=1      )

								return text_encoder_hidden_states, additive_clip_time_embeddings


