


"""simple docstring"""

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    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       __lowercase									=	1
       __lowercase									=	3
       __lowercase									=	(3_2, 3_2)

       __lowercase									=	floats_tensor((batch_size, num_channels) + sizes      ,       rng=random.Random(0							)							).to(A_							)
       return image
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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=A_      ,       only_cross_attention=(True, True, False)      ,       num_class_embeds=1_0_0      ,       )
       return model
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	CLIPTextConfig(
           bos_token_id=0      ,       eos_token_id=2      ,       hidden_size=3_2      ,       intermediate_size=3_7      ,       layer_norm_eps=1e-0_5      ,       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(A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


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

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

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

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

       __lowercase									=	output.images

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

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

       __lowercase									=	low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)
       __lowercase									=	np.array([0.31_13, 0.39_10, 0.42_72, 0.48_59, 0.50_61, 0.46_52, 0.53_62, 0.57_15, 0.56_61]							)

       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 SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


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

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

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

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

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






      @unittest.skipIf(torch_device != """cuda"""      ,       """This test requires a GPU"""							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


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

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

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

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

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

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






@slow
@require_torch_gpu
class    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


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

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

       __lowercase									=	"""a cat sitting on a park bench"""

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

       assert image.shape == (5_1_2, 5_1_2, 3)
       assert np.abs(expected_image - image							).max() < 1e-3
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


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

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

       __lowercase									=	"""a cat sitting on a park bench"""

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

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






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


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

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

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

       __lowercase									=	"""a cat sitting on a park bench"""

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

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

# Algorithm for the pigeonhole sorting




def lowerCAmelCase_      (				UpperCamelCase__	: Dict ):






 """simple docstring"""







 __lowercase									=	min(UpperCamelCase__ )  # min() finds the minimum value
 __lowercase									=	max(UpperCamelCase__ )  # max() finds the maximum value

 __lowercase									=	max_val - min_val + 1  # size is difference of max and min values plus one

 # list of pigeonholes of size equal to the variable size
 __lowercase									=	[0] * size

 # Populate the pigeonholes.
 for x in a:
  assert isinstance(UpperCamelCase__				,   UpperCamelCase__ ), "integers only please"
  holes[x - min_val] += 1

 # Putting the elements back into the array in an order.
 __lowercase									=	0
 for count in range(UpperCamelCase__ ):
  while holes[count] > 0:
   holes[count] -= 1
   __lowercase									=	count + min_val
   i += 1


def lowerCAmelCase_      (				):






 """simple docstring"""







 __lowercase									=	[8, 3, 2, 7, 4, 6, 8]
 pigeonhole_sort(UpperCamelCase__ )
 print("""Sorted order is:"""				,   """ """.join(UpperCamelCase__ ) )


if __name__ == "__main__":
       main()
