



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 		UpperCAmelCase						(						unittest.TestCase				):





   def 	__lowerCAmelCase							(	self			):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





   @property
   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  1
         _lowerCAmelCase 			=  3
         _lowerCAmelCase 			=  (32, 32)

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





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  UNetaDConditionModel(
             block_out_channels=(32, 32, 64)       ,     layers_per_block=2       ,     sample_size=32       ,     in_channels=7       ,     out_channels=4       ,     down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''', '''CrossAttnDownBlock2D''')       ,     up_block_types=('''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,     cross_attention_dim=32       ,     attention_head_dim=8       ,     use_linear_projection=_lowerCAmelCase       ,     only_cross_attention=(True, True, False)       ,     num_class_embeds=100       ,     )
         return model





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  AutoencoderKL(
             block_out_channels=[32, 32, 64]       ,     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 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  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=1_000       ,     hidden_act='''gelu'''       ,     projection_dim=512       ,     )
         return CLIPTextModel(_lowerCAmelCase			)





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

         _lowerCAmelCase 			=  self.dummy_image.cpu().permute(0       ,     2       ,     3       ,     1			)[0]
         _lowerCAmelCase 			=  Image.fromarray(np.uinta(_lowerCAmelCase			)			).convert('''RGB'''			).resize((64, 64)			)

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

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

         _lowerCAmelCase 			=  output.images

         _lowerCAmelCase 			=  torch.Generator(device=_lowerCAmelCase			).manual_seed(0			)
         _lowerCAmelCase 			=  sd_pipe(
             [prompt]       ,     image=_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     guidance_scale=6.0       ,     noise_level=20       ,     num_inference_steps=2       ,     output_type='''np'''       ,     return_dict=_lowerCAmelCase       ,     )[0]

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

         _lowerCAmelCase 			=  low_res_image.size[0] * 4
         assert image.shape == (1, expected_height_width, expected_height_width, 3)
         _lowerCAmelCase 			=  np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661]			)

         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 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         _lowerCAmelCase 			=  self.dummy_cond_unet_upscale
         _lowerCAmelCase 			=  DDPMScheduler()
         _lowerCAmelCase 			=  DDIMScheduler(prediction_type='''v_prediction'''			)
         _lowerCAmelCase 			=  self.dummy_vae
         _lowerCAmelCase 			=  self.dummy_text_encoder
         _lowerCAmelCase 			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

         _lowerCAmelCase 			=  self.dummy_image.cpu().permute(0       ,     2       ,     3       ,     1			)[0]
         _lowerCAmelCase 			=  Image.fromarray(np.uinta(_lowerCAmelCase			)			).convert('''RGB'''			).resize((64, 64)			)

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

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

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





   @unittest.skipIf(torch_device != '''cuda'''       ,     '''This test requires a GPU'''			)
   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.dummy_cond_unet_upscale
         _lowerCAmelCase 			=  DDPMScheduler()
         _lowerCAmelCase 			=  DDIMScheduler(prediction_type='''v_prediction'''			)
         _lowerCAmelCase 			=  self.dummy_vae
         _lowerCAmelCase 			=  self.dummy_text_encoder
         _lowerCAmelCase 			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

         _lowerCAmelCase 			=  self.dummy_image.cpu().permute(0       ,     2       ,     3       ,     1			)[0]
         _lowerCAmelCase 			=  Image.fromarray(np.uinta(_lowerCAmelCase			)			).convert('''RGB'''			).resize((64, 64)			)

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

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

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

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






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





   def 	__lowerCAmelCase							(	self			):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/sd2-upscale/low_res_cat.png'''			)
         _lowerCAmelCase 			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
             '''/upsampled_cat.npy'''			)

         _lowerCAmelCase 			=  '''stabilityai/stable-diffusion-x4-upscaler'''
         _lowerCAmelCase 			=  StableDiffusionUpscalePipeline.from_pretrained(_lowerCAmelCase			)
         pipe.to(_lowerCAmelCase			)
         pipe.set_progress_bar_config(disable=_lowerCAmelCase			)
         pipe.enable_attention_slicing()

         _lowerCAmelCase 			=  '''a cat sitting on a park bench'''

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(
             prompt=_lowerCAmelCase       ,     image=_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     output_type='''np'''       ,     )
         _lowerCAmelCase 			=  output.images[0]

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





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/sd2-upscale/low_res_cat.png'''			)
         _lowerCAmelCase 			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
             '''/upsampled_cat_fp16.npy'''			)

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

         _lowerCAmelCase 			=  '''a cat sitting on a park bench'''

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(
             prompt=_lowerCAmelCase       ,     image=_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     output_type='''np'''       ,     )
         _lowerCAmelCase 			=  output.images[0]

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





   def 	__lowerCAmelCase							(	self			):
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

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

         _lowerCAmelCase 			=  '''stabilityai/stable-diffusion-x4-upscaler'''
         _lowerCAmelCase 			=  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()

         _lowerCAmelCase 			=  '''a cat sitting on a park bench'''

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(
             prompt=_lowerCAmelCase       ,     image=_lowerCAmelCase       ,     generator=_lowerCAmelCase       ,     num_inference_steps=5       ,     output_type='''np'''       ,     )

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







from __future__ import annotations
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	list				)->list:
      if len(_SCREAMING_SNAKE_CASE				) == 0:
            return []
      _lowerCAmelCase			,   _lowerCAmelCase 			=  min(_SCREAMING_SNAKE_CASE				), max(_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  int(max_value - min_value				) + 1
      _lowerCAmelCase 			=  [[] for _ in range(_SCREAMING_SNAKE_CASE				)]

      for i in my_list:
            buckets[int(i - min_value				)].append(_SCREAMING_SNAKE_CASE				)

      return [v for bucket in buckets for v in sorted(_SCREAMING_SNAKE_CASE				)]


if __name__ == "__main__":
    from doctest import testmod

    testmod()
    assert bucket_sort([4, 5, 3, 2, 1]) == [1, 2, 3, 4, 5]
    assert bucket_sort([0, 1, -1_0, 1_5, 2, -2]) == [-1_0, -2, 0, 1, 2, 1_5]