




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  ):





				'''simple docstring'''



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



				@property
				def       __snake_case		( self						):
											lowerCAmelCase   =		1
											lowerCAmelCase   =		3
											lowerCAmelCase   =		(32, 32)

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



				@property
				def       __snake_case		( 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=lowercase_       , only_cross_attention=(True, True, False)       , num_class_embeds=1_00       , )
											return model



				@property
				def       __snake_case		( 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       __snake_case		( 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-0_5       , num_attention_heads=4       , num_hidden_layers=5       , pad_token_id=1       , vocab_size=10_00       , hidden_act='''gelu'''       , projection_dim=5_12       , )
											return CLIPTextModel(lowercase_						)



				def       __snake_case		( 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(lowercase_						)						).convert('''RGB'''						).resize((64, 64)						)

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

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

											lowerCAmelCase   =		output.images

											lowerCAmelCase   =		torch.Generator(device=lowercase_						).manual_seed(0						)
											lowerCAmelCase   =		sd_pipe(
											    [prompt]       , image=lowercase_       , generator=lowercase_       , guidance_scale=6.0       , noise_level=20       , num_inference_steps=2       , output_type='''np'''       , return_dict=lowercase_       , )[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.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       __snake_case		( 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(lowercase_						)						).convert('''RGB'''						).resize((64, 64)						)

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

											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=lowercase_						).manual_seed(0						)
											lowerCAmelCase   =		sd_pipe(
											    [prompt]       , image=lowercase_       , generator=lowercase_       , 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       __snake_case		( 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(lowercase_						)						).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=lowercase_       , low_res_scheduler=lowercase_       , scheduler=lowercase_       , vae=lowercase_       , text_encoder=lowercase_       , tokenizer=lowercase_       , max_noise_level=3_50       , )
											lowerCAmelCase   =		sd_pipe.to(lowercase_						)
											sd_pipe.set_progress_bar_config(disable=lowercase_						)

											lowerCAmelCase   =		'''A painting of a squirrel eating a burger'''
											lowerCAmelCase   =		torch.manual_seed(0						)
											lowerCAmelCase   =		sd_pipe(
											    [prompt]       , image=lowercase_       , generator=lowercase_       , 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  ):





				'''simple docstring'''



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



				def       __snake_case		( 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(lowercase_						)
											pipe.to(lowercase_						)
											pipe.set_progress_bar_config(disable=lowercase_						)
											pipe.enable_attention_slicing()

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

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

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



				def       __snake_case		( 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(
											    lowercase_       , torch_dtype=torch.floataa       , )
											pipe.to(lowercase_						)
											pipe.set_progress_bar_config(disable=lowercase_						)
											pipe.enable_attention_slicing()

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

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

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






				def       __snake_case		( 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(
											    lowercase_       , torch_dtype=torch.floataa       , )
											pipe.to(lowercase_						)
											pipe.set_progress_bar_config(disable=lowercase_						)
											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=lowercase_       , image=lowercase_       , generator=lowercase_       , 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









import json
import os
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from requests.exceptions import HTTPError

from transformers.utils import (
    CONFIG_NAME,
    FLAX_WEIGHTS_NAME,
    TF2_WEIGHTS_NAME,
    TRANSFORMERS_CACHE,
    WEIGHTS_NAME,
    cached_file,
    get_file_from_repo,
    has_file,
)


UpperCAmelCase_						="""hf-internal-testing/tiny-random-bert"""
UpperCAmelCase_						=os.path.join(TRANSFORMERS_CACHE, """models--hf-internal-testing--tiny-random-bert""")
UpperCAmelCase_						="""9b8c223d42b2188cb49d29af482996f9d0f3e5a6"""


class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				def       __snake_case		( self						):
											lowerCAmelCase   =		cached_file(UpperCAmelCase_       , UpperCAmelCase_						)
											# Should have downloaded the file in here
											self.assertTrue(os.path.isdir(UpperCAmelCase_						)						)
											# Cache should contain at least those three subfolders:
											for subfolder in ["blobs", "refs", "snapshots"]:
																		self.assertTrue(os.path.isdir(os.path.join(UpperCAmelCase_       , UpperCAmelCase_						)						)						)
											with open(os.path.join(UpperCAmelCase_       , '''refs'''       , '''main'''						)						) as f:
																		lowerCAmelCase   =		f.read()
											self.assertEqual(UpperCAmelCase_       , os.path.join(UpperCAmelCase_       , '''snapshots'''       , UpperCAmelCase_       , UpperCAmelCase_						)						)
											self.assertTrue(os.path.isfile(UpperCAmelCase_						)						)

											# File is cached at the same place the second time.
											lowerCAmelCase   =		cached_file(UpperCAmelCase_       , UpperCAmelCase_						)
											self.assertEqual(UpperCAmelCase_       , UpperCAmelCase_						)

											# Using a specific revision to test the full commit hash.
											lowerCAmelCase   =		cached_file(UpperCAmelCase_       , UpperCAmelCase_       , revision='''9b8c223'''						)
											self.assertEqual(UpperCAmelCase_       , os.path.join(UpperCAmelCase_       , '''snapshots'''       , UpperCAmelCase_       , UpperCAmelCase_						)						)



				def       __snake_case		( self						):
											with self.assertRaisesRegex(UpperCAmelCase_       , '''is not a valid model identifier'''						):
																		lowerCAmelCase   =		cached_file('''tiny-random-bert'''       , UpperCAmelCase_						)

											with self.assertRaisesRegex(UpperCAmelCase_       , '''is not a valid git identifier'''						):
																		lowerCAmelCase   =		cached_file(UpperCAmelCase_       , UpperCAmelCase_       , revision='''aaaa'''						)

											with self.assertRaisesRegex(UpperCAmelCase_       , '''does not appear to have a file named'''						):
																		lowerCAmelCase   =		cached_file(UpperCAmelCase_       , '''conf'''						)



				def       __snake_case		( self						):
											with self.assertRaisesRegex(UpperCAmelCase_       , '''does not appear to have a file named'''						):
																		lowerCAmelCase   =		cached_file(UpperCAmelCase_       , '''conf'''						)

											with open(os.path.join(UpperCAmelCase_       , '''refs'''       , '''main'''						)						) as f:
																		lowerCAmelCase   =		f.read()
											self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase_       , '''.no_exist'''       , UpperCAmelCase_       , '''conf'''						)						)						)

											lowerCAmelCase   =		cached_file(UpperCAmelCase_       , '''conf'''       , _raise_exceptions_for_missing_entries=UpperCAmelCase_						)
											self.assertIsNone(UpperCAmelCase_						)

											lowerCAmelCase   =		cached_file(UpperCAmelCase_       , '''conf'''       , local_files_only=UpperCAmelCase_       , _raise_exceptions_for_missing_entries=UpperCAmelCase_						)
											self.assertIsNone(UpperCAmelCase_						)

											lowerCAmelCase   =		mock.Mock()
											lowerCAmelCase   =		5_00
											lowerCAmelCase   =		{}
											lowerCAmelCase   =		HTTPError
											lowerCAmelCase   =		{}

											# Under the mock environment we get a 500 error when trying to reach the tokenizer.
											with mock.patch('''requests.Session.request'''       , return_value=UpperCAmelCase_						) as mock_head:
																		lowerCAmelCase   =		cached_file(UpperCAmelCase_       , '''conf'''       , _raise_exceptions_for_connection_errors=UpperCAmelCase_						)
																		self.assertIsNone(UpperCAmelCase_						)
																		# This check we did call the fake head request
																		mock_head.assert_called()



				def       __snake_case		( self						):
											self.assertTrue(has_file('''hf-internal-testing/tiny-bert-pt-only'''       , UpperCAmelCase_						)						)
											self.assertFalse(has_file('''hf-internal-testing/tiny-bert-pt-only'''       , UpperCAmelCase_						)						)
											self.assertFalse(has_file('''hf-internal-testing/tiny-bert-pt-only'''       , UpperCAmelCase_						)						)



				def       __snake_case		( self						):
											# `get_file_from_repo` returns None if the file does not exist
											self.assertIsNone(get_file_from_repo('''bert-base-cased'''       , '''ahah.txt'''						)						)

											# The function raises if the repository does not exist.
											with self.assertRaisesRegex(UpperCAmelCase_       , '''is not a valid model identifier'''						):
																		get_file_from_repo('''bert-base-case'''       , UpperCAmelCase_						)

											# The function raises if the revision does not exist.
											with self.assertRaisesRegex(UpperCAmelCase_       , '''is not a valid git identifier'''						):
																		get_file_from_repo('''bert-base-cased'''       , UpperCAmelCase_       , revision='''ahaha'''						)

											lowerCAmelCase   =		get_file_from_repo('''bert-base-cased'''       , UpperCAmelCase_						)
											# The name is the cached name which is not very easy to test, so instead we load the content.
											lowerCAmelCase   =		json.loads(open(UpperCAmelCase_       , '''r'''						).read()						)
											self.assertEqual(config['''hidden_size''']       , 7_68						)






				def       __snake_case		( self						):
											with tempfile.TemporaryDirectory() as tmp_dir:
																		lowerCAmelCase   =		Path(UpperCAmelCase_						) / '''a.txt'''
																		filename.touch()
																		self.assertEqual(get_file_from_repo(UpperCAmelCase_       , '''a.txt'''						)       , str(UpperCAmelCase_						)						)

																		self.assertIsNone(get_file_from_repo(UpperCAmelCase_       , '''b.txt'''						)						)



