




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



					'''simple docstring'''





					def     __snake_case						(	self     :						Optional[int])    ->							Union[str, Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					@property
					def     __snake_case						(	self     :						List[Any])    ->							Dict:
								A_									= 1
								A_									= 3
								A_									= (32, 32)

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



					@property
					def     __snake_case						(	self     :						Optional[int])    ->							Tuple:
								torch.manual_seed(0)
								A_									= 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=100      ,       )
								return model



					@property
					def     __snake_case						(	self     :						Union[str, Any])    ->							Any:
								torch.manual_seed(0)
								A_									= 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     :						str)    ->							str:
								torch.manual_seed(0)
								A_									= 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(_lowercase)



					def     __snake_case						(	self     :						Union[str, Any])    ->							Dict:
								A_									= 'cpu'  # ensure determinism for the device-dependent torch.Generator
								A_									= self.dummy_cond_unet_upscale
								A_									= DDPMScheduler()
								A_									= DDIMScheduler(prediction_type='v_prediction')
								A_									= self.dummy_vae
								A_									= self.dummy_text_encoder
								A_									= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

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

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

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

								A_									= output.images

								A_									= torch.Generator(device=_lowercase).manual_seed(0)
								A_									= 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]

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

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

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

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

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

								A_									= torch.Generator(device=_lowercase).manual_seed(0)
								A_									= 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'      ,       )
								A_									= output.images
								assert image.shape[0] == 2



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

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

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

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

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

								A_									= 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     :						List[Any])    ->							Optional[Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def     __snake_case						(	self     :						Union[str, Any])    ->							Tuple:
								A_									= load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/sd2-upscale/low_res_cat.png')
								A_									= load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
								    '/upsampled_cat.npy')

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

								A_									= 'a cat sitting on a park bench'

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

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



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

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

								A_									= 'a cat sitting on a park bench'

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

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



					def     __snake_case						(	self     :						List[Any])    ->							Union[str, Any]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

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

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

								A_									= 'a cat sitting on a park bench'

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

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










'''simple docstring'''
import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()





def 			lowerCamelCase(							SCREAMING_SNAKE_CASE_       ,SCREAMING_SNAKE_CASE_       ,SCREAMING_SNAKE_CASE_     )		->		str:
			# Construct model
			if openai_config_file == "":
						A_									= OpenAIGPTConfig()
			else:
						A_									= OpenAIGPTConfig.from_json_file(SCREAMING_SNAKE_CASE_     )
			A_									= OpenAIGPTModel(SCREAMING_SNAKE_CASE_     )

			# Load weights from numpy
			load_tf_weights_in_openai_gpt(SCREAMING_SNAKE_CASE_       ,SCREAMING_SNAKE_CASE_       ,SCREAMING_SNAKE_CASE_     )

			# Save pytorch-model
			A_									= pytorch_dump_folder_path + '/' + WEIGHTS_NAME
			A_									= pytorch_dump_folder_path + '/' + CONFIG_NAME
			print(F'Save PyTorch model to {pytorch_weights_dump_path}'     )
			torch.save(model.state_dict()       ,SCREAMING_SNAKE_CASE_     )
			print(F'Save configuration file to {pytorch_config_dump_path}'     )
			with open(SCREAMING_SNAKE_CASE_       ,'w'       ,encoding='utf-8'     ) as f:
						f.write(config.to_json_string()     )


if __name__ == "__main__":
							_SCREAMING_SNAKE_CASE			=						argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    "--openai_checkpoint_folder_path",
							    default=None,
							    type=str,
							    required=True,
							    help="Path to the TensorFlow checkpoint path.",
							)
							parser.add_argument(
							    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
							)
							parser.add_argument(
							    "--openai_config_file",
							    default="",
							    type=str,
							    help=(
							        "An optional config json file corresponding to the pre-trained OpenAI model. \n"
							        "This specifies the model architecture."
							    ),
							)
							_SCREAMING_SNAKE_CASE			=						parser.parse_args()
							convert_openai_checkpoint_to_pytorch(
							    args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
							)




