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



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



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

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



				@property
				def  _lowercase					(self    :				List[Any])   ->     List[Any]:
					torch.manual_seed(0)
					__snake_case      :							Dict						=			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=1_00     ,  )
					return model



				@property
				def  _lowercase					(self    :				Dict)   ->     Tuple:
					torch.manual_seed(0)
					__snake_case      :							str						=			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  _lowercase					(self    :				Optional[Any])   ->     Dict:
					torch.manual_seed(0)
					__snake_case      :							int						=			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=10_00     ,  hidden_act='gelu'     ,  projection_dim=5_12     ,  )
					return CLIPTextModel(_A)



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

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

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

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

					__snake_case      :							List[Any]						=			output.images

					__snake_case      :							Optional[Any]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Optional[Any]						=			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]

					__snake_case      :							int						=			image[0, -3:, -3:, -1]
					__snake_case      :							Tuple						=			image_from_tuple[0, -3:, -3:, -1]

					__snake_case      :							Dict						=			low_res_image.size[0] * 4
					assert image.shape == (1, expected_height_width, expected_height_width, 3)
					__snake_case      :							Any						=			np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661])

					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  _lowercase					(self    :				Any)   ->     Any:
					__snake_case      :							int						=			'cpu'  # ensure determinism for the device-dependent torch.Generator
					__snake_case      :							Optional[Any]						=			self.dummy_cond_unet_upscale
					__snake_case      :							Tuple						=			DDPMScheduler()
					__snake_case      :							str						=			DDIMScheduler(prediction_type='v_prediction')
					__snake_case      :							Optional[int]						=			self.dummy_vae
					__snake_case      :							List[Any]						=			self.dummy_text_encoder
					__snake_case      :							Any						=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

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

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

					__snake_case      :							Union[str, Any]						=			'A painting of a squirrel eating a burger'
					__snake_case      :							Union[str, Any]						=			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      :							Dict						=			output.images
					assert image.shape[0] == 2

					__snake_case      :							List[Any]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							List[Any]						=			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'     ,  )
					__snake_case      :							int						=			output.images
					assert image.shape[0] == 2






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

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

					# put models in fp16, except vae as it overflows in fp16
					__snake_case      :							List[Any]						=			unet.half()
					__snake_case      :							List[str]						=			text_encoder.half()

					# make sure here that pndm scheduler skips prk
					__snake_case      :							Dict						=			StableDiffusionUpscalePipeline(
					    unet=_A     ,  low_res_scheduler=_A     ,  scheduler=_A     ,  vae=_A     ,  text_encoder=_A     ,  tokenizer=_A     ,  max_noise_level=3_50     ,  )
					__snake_case      :							Optional[Any]						=			sd_pipe.to(_A)
					sd_pipe.set_progress_bar_config(disable=_A)

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

					__snake_case      :							Optional[Any]						=			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						):



				def  _lowercase					(self    :				Tuple)   ->     Any:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()



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

					__snake_case      :							Optional[int]						=			'stabilityai/stable-diffusion-x4-upscaler'
					__snake_case      :							List[str]						=			StableDiffusionUpscalePipeline.from_pretrained(_A)
					pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)
					pipe.enable_attention_slicing()

					__snake_case      :							List[str]						=			'a cat sitting on a park bench'

					__snake_case      :							int						=			torch.manual_seed(0)
					__snake_case      :							List[str]						=			pipe(
					    prompt=_A     ,  image=_A     ,  generator=_A     ,  output_type='np'     ,  )
					__snake_case      :							Optional[int]						=			output.images[0]

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



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

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

					__snake_case      :							Union[str, Any]						=			'a cat sitting on a park bench'

					__snake_case      :							Dict						=			torch.manual_seed(0)
					__snake_case      :							Tuple						=			pipe(
					    prompt=_A     ,  image=_A     ,  generator=_A     ,  output_type='np'     ,  )
					__snake_case      :							int						=			output.images[0]

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






				def  _lowercase					(self    :				List[str])   ->     int:
					torch.cuda.empty_cache()
					torch.cuda.reset_max_memory_allocated()
					torch.cuda.reset_peak_memory_stats()

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

					__snake_case      :							List[Any]						=			'stabilityai/stable-diffusion-x4-upscaler'
					__snake_case      :							Tuple						=			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()

					__snake_case      :							Tuple						=			'a cat sitting on a park bench'

					__snake_case      :							Optional[Any]						=			torch.manual_seed(0)
					__snake_case      :							Optional[Any]						=			pipe(
					    prompt=_A     ,  image=_A     ,  generator=_A     ,  num_inference_steps=5     ,  output_type='np'     ,  )

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






"""simple docstring"""







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


_a					:					Tuple=			{
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}

class 		UpperCamelCase      (   lowercase						):
				UpperCAmelCase    :				str   =		"""albert"""



				def __init__(self    :				Union[str, Any]     ,  _A    :				int=3_00_00     ,  _A    :				Any=1_28     ,  _A    :				Tuple=40_96     ,  _A    :				int=12     ,  _A    :				Tuple=1     ,  _A    :				int=64     ,  _A    :				Optional[Any]=1_63_84     ,  _A    :				Optional[Any]=1     ,  _A    :				List[str]="gelu_new"     ,  _A    :				Any=0     ,  _A    :				Optional[Any]=0     ,  _A    :				List[Any]=5_12     ,  _A    :				List[Any]=2     ,  _A    :				Dict=0.02     ,  _A    :				Union[str, Any]=1E-12     ,  _A    :				Tuple=0.1     ,  _A    :				int="absolute"     ,  _A    :				List[Any]=0     ,  _A    :				str=2     ,  _A    :				int=3     ,  **_A    :				Tuple     ,  )   ->     Union[str, Any]:
					super().__init__(pad_token_id=_A     ,  bos_token_id=_A     ,  eos_token_id=_A     ,  **_A)

					__snake_case      :							str						=			vocab_size
					__snake_case      :							List[Any]						=			embedding_size
					__snake_case      :							Tuple						=			hidden_size
					__snake_case      :							Optional[Any]						=			num_hidden_layers
					__snake_case      :							List[str]						=			num_hidden_groups
					__snake_case      :							Dict						=			num_attention_heads
					__snake_case      :							Any						=			inner_group_num
					__snake_case      :							Union[str, Any]						=			hidden_act
					__snake_case      :							List[Any]						=			intermediate_size
					__snake_case      :							List[str]						=			hidden_dropout_prob
					__snake_case      :							int						=			attention_probs_dropout_prob
					__snake_case      :							Any						=			max_position_embeddings
					__snake_case      :							Optional[int]						=			type_vocab_size
					__snake_case      :							Optional[Any]						=			initializer_range
					__snake_case      :							List[Any]						=			layer_norm_eps
					__snake_case      :							Optional[Any]						=			classifier_dropout_prob
					__snake_case      :							List[str]						=			position_embedding_type







class 		UpperCamelCase      (   lowercase						):



				@property
				def  _lowercase					(self    :				Tuple)   ->     Mapping[str, Mapping[int, str]]:
					if self.task == "multiple-choice":
						__snake_case      :							int						=			{0: 'batch', 1: 'choice', 2: 'sequence'}
					else:
						__snake_case      :							Optional[Any]						=			{0: 'batch', 1: 'sequence'}
					return OrderedDict(
					    [
					        ('input_ids', dynamic_axis),
					        ('attention_mask', dynamic_axis),
					        ('token_type_ids', dynamic_axis),
					    ])





