



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







													'''simple docstring'''

													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()






						@property
						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= 3
													__SCREAMING_SNAKE_CASE									= (32, 32)

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






						@property
						def 	_A     (   self		):







													'''simple docstring'''

													torch.manual_seed(0		)
													__SCREAMING_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=_A				,							only_cross_attention=(True, True, False)				,							num_class_embeds=100				,							)
													return model






						@property
						def 	_A     (   self		):







													'''simple docstring'''

													torch.manual_seed(0		)
													__SCREAMING_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		):







													'''simple docstring'''

													torch.manual_seed(0		)
													__SCREAMING_SNAKE_CASE									= 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(_A		)






						def 	_A     (   self		):







													'''simple docstring'''

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

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

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

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

													__SCREAMING_SNAKE_CASE									= output.images

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

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

													__SCREAMING_SNAKE_CASE									= low_res_image.size[0] * 4
													assert image.shape == (1, expected_height_width, expected_height_width, 3)
													__SCREAMING_SNAKE_CASE									= np.array([0.3_1_1_3, 0.3_9_1_0, 0.4_2_7_2, 0.4_8_5_9, 0.5_0_6_1, 0.4_6_5_2, 0.5_3_6_2, 0.5_7_1_5, 0.5_6_6_1]		)

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







													'''simple docstring'''

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

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

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

													__SCREAMING_SNAKE_CASE									= 'A painting of a squirrel eating a burger'
													__SCREAMING_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'				,							)
													__SCREAMING_SNAKE_CASE									= output.images
													assert image.shape[0] == 2

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






						@unittest.skipIf(torch_device != 'cuda'				,							'This test requires a GPU'		)
						def 	_A     (   self		):







													'''simple docstring'''

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

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

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

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

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

													__SCREAMING_SNAKE_CASE									= 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 	_A     (   self		):







													'''simple docstring'''

													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()






						def 	_A     (   self		):







													'''simple docstring'''

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

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

													__SCREAMING_SNAKE_CASE									= 'a cat sitting on a park bench'

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

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






						def 	_A     (   self		):







													'''simple docstring'''

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

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

													__SCREAMING_SNAKE_CASE									= 'a cat sitting on a park bench'

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

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






						def 	_A     (   self		):







													'''simple docstring'''

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

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

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

													__SCREAMING_SNAKE_CASE									= 'a cat sitting on a park bench'

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

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











from dataclasses import dataclass, field
from typing import Optional

from transformers import AutoConfig, AutoImageProcessor, AutoTokenizer, FlaxVisionEncoderDecoderModel, HfArgumentParser

@dataclass
class        UpperCAmelCase_  :

						'''simple docstring'''




						UpperCamelCase__      :			str      =  field(
						    metadata={'''help''': '''The output directory where the model will be written.'''}   ,   )
						UpperCamelCase__      :			str      =  field(
						    metadata={
						        '''help''': (
						            '''The encoder model checkpoint for weights initialization.'''
						            '''Don\'t set if you want to train an encoder model from scratch.'''
						        )
						    }   ,   )
						UpperCamelCase__      :			str      =  field(
						    metadata={
						        '''help''': (
						            '''The decoder model checkpoint for weights initialization.'''
						            '''Don\'t set if you want to train a decoder model from scratch.'''
						        )
						    }   ,   )
						UpperCamelCase__      :			Optional[str]      =  field(
						    default=UpperCamelCase_   ,   metadata={'''help''': '''Pretrained encoder config name or path if not the same as encoder_model_name'''}	)
						UpperCamelCase__      :			Optional[str]      =  field(
						    default=UpperCamelCase_   ,   metadata={'''help''': '''Pretrained decoder config name or path if not the same as decoder_model_name'''}	)







def      __lowercase       (			)    ->			Union[str, Any]:
							__SCREAMING_SNAKE_CASE									= HfArgumentParser((ModelArguments,)       )
							((__SCREAMING_SNAKE_CASE)							,	)									= parser.parse_args_into_dataclasses()

							# Load pretrained model and tokenizer

							# Use explicit specified encoder config
							if model_args.encoder_config_name:
														__SCREAMING_SNAKE_CASE									= AutoConfig.from_pretrained(model_args.encoder_config_name       )
							# Use pretrained encoder model's config
							else:
														__SCREAMING_SNAKE_CASE									= AutoConfig.from_pretrained(model_args.encoder_model_name_or_path       )

							# Use explicit specified decoder config
							if model_args.decoder_config_name:
														__SCREAMING_SNAKE_CASE									= AutoConfig.from_pretrained(model_args.decoder_config_name       )
							# Use pretrained decoder model's config
							else:
														__SCREAMING_SNAKE_CASE									= AutoConfig.from_pretrained(model_args.decoder_model_name_or_path       )

							# necessary for `from_encoder_decoder_pretrained` when `decoder_config` is passed
							__SCREAMING_SNAKE_CASE									= True
							__SCREAMING_SNAKE_CASE									= True

							__SCREAMING_SNAKE_CASE									= FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
							    encoder_pretrained_model_name_or_path=model_args.encoder_model_name_or_path							,       decoder_pretrained_model_name_or_path=model_args.decoder_model_name_or_path							,       encoder_config=a__							,       decoder_config=a__							,       )

							# GPT2 only has bos/eos tokens but not decoder_start/pad tokens
							__SCREAMING_SNAKE_CASE									= decoder_config.decoder_start_token_id
							__SCREAMING_SNAKE_CASE									= decoder_config.pad_token_id
							if decoder_start_token_id is None:
														__SCREAMING_SNAKE_CASE									= decoder_config.bos_token_id
							if pad_token_id is None:
														__SCREAMING_SNAKE_CASE									= decoder_config.eos_token_id

							# This is necessary to make Flax's generate() work
							__SCREAMING_SNAKE_CASE									= decoder_config.eos_token_id
							__SCREAMING_SNAKE_CASE									= decoder_start_token_id
							__SCREAMING_SNAKE_CASE									= pad_token_id

							__SCREAMING_SNAKE_CASE									= AutoImageProcessor.from_pretrained(model_args.encoder_model_name_or_path       )

							__SCREAMING_SNAKE_CASE									= AutoTokenizer.from_pretrained(model_args.decoder_model_name_or_path       )
							__SCREAMING_SNAKE_CASE									= tokenizer.convert_ids_to_tokens(model.config.pad_token_id       )

							model.save_pretrained(model_args.output_dir       )
							image_processor.save_pretrained(model_args.output_dir       )
							tokenizer.save_pretrained(model_args.output_dir       )


if __name__ == "__main__":
			main()






