

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




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




				@property
				def 					a							(      self ):
											snake_case_           =					1
											snake_case_           =					3
											snake_case_           =					(32, 32)

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




				@property
				def 					a							(      self ):
											torch.manual_seed(0 )
											snake_case_           =					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=snake_case   ,     only_cross_attention=(True, True, False)   ,     num_class_embeds=100   ,     )
											return model




				@property
				def 					a							(      self ):
											torch.manual_seed(0 )
											snake_case_           =					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 					a							(      self ):
											torch.manual_seed(0 )
											snake_case_           =					CLIPTextConfig(
											    bos_token_id=0   ,     eos_token_id=2   ,     hidden_size=32   ,     intermediate_size=37   ,     layer_norm_eps=1e-0_5   ,     num_attention_heads=4   ,     num_hidden_layers=5   ,     pad_token_id=1   ,     vocab_size=1000   ,     hidden_act='gelu'   ,     projection_dim=512   ,     )
											return CLIPTextModel(snake_case )




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

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

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

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

											snake_case_           =					output.images

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

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

											snake_case_           =					low_res_image.size[0] * 4
											assert image.shape == (1, expected_height_width, expected_height_width, 3)
											snake_case_           =					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 					a							(      self ):
											snake_case_           =					'cpu'  # ensure determinism for the device-dependent torch.Generator
											snake_case_           =					self.dummy_cond_unet_upscale
											snake_case_           =					DDPMScheduler()
											snake_case_           =					DDIMScheduler(prediction_type='v_prediction' )
											snake_case_           =					self.dummy_vae
											snake_case_           =					self.dummy_text_encoder
											snake_case_           =					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

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

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

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

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




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

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

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

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

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

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







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




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




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

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

											snake_case_           =					'a cat sitting on a park bench'

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

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




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

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

											snake_case_           =					'a cat sitting on a park bench'

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

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




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

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

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

											snake_case_           =					'a cat sitting on a park bench'

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

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







import os

try:
					from .build_directory_md import good_file_paths
except ImportError:
					from build_directory_md import good_file_paths  # type: ignore

_UpperCAmelCase	:     Optional[int]		       = list(good_file_paths())
assert filepaths, "good_file_paths() failed!"

_UpperCAmelCase	:     Dict		       = [file for file in filepaths if file != file.lower()]
if upper_files:
					print(F'''{len(upper_files)} files contain uppercase characters:''')
					print("""\n""".join(upper_files) + """\n""")

_UpperCAmelCase	:     Tuple		       = [file for file in filepaths if """ """ in file]
if space_files:
					print(F'''{len(space_files)} files contain space characters:''')
					print("""\n""".join(space_files) + """\n""")

_UpperCAmelCase	:     Union[str, Any]		       = [file for file in filepaths if """-""" in file]
if hyphen_files:
					print(F'''{len(hyphen_files)} files contain hyphen characters:''')
					print("""\n""".join(hyphen_files) + """\n""")

_UpperCAmelCase	:     List[Any]		       = [file for file in filepaths if os.sep not in file]
if nodir_files:
					print(F'''{len(nodir_files)} files are not in a directory:''')
					print("""\n""".join(nodir_files) + """\n""")

_UpperCAmelCase	:     List[str]		       = len(upper_files + space_files + hyphen_files + nodir_files)
if bad_files:
					import sys

					sys.exit(bad_files)




