


from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


A							:					int     			=     logging.get_logger(__name__)

# TODO Update this
A							:					int     			=     {
    'facebook/esm-1b': 'https://huggingface.co/facebook/esm-1b/resolve/main/config.json',
    # See all ESM models at https://huggingface.co/models?filter=esm
}






class     UpperCamelCase(  _a   ):
					snake_case_		:				str										=	"""esm"""


					def __init__(						self	:		Dict     ,    SCREAMING_SNAKE_CASE	:		List[Any]=None     ,    SCREAMING_SNAKE_CASE	:		str=None     ,    SCREAMING_SNAKE_CASE	:		List[Any]=None     ,    SCREAMING_SNAKE_CASE	:		str=7_6_8     ,    SCREAMING_SNAKE_CASE	:		int=1_2     ,    SCREAMING_SNAKE_CASE	:		List[str]=1_2     ,    SCREAMING_SNAKE_CASE	:		int=3_0_7_2     ,    SCREAMING_SNAKE_CASE	:		Tuple=0.1     ,    SCREAMING_SNAKE_CASE	:		List[Any]=0.1     ,    SCREAMING_SNAKE_CASE	:		List[Any]=1_0_2_6     ,    SCREAMING_SNAKE_CASE	:		str=0.02     ,    SCREAMING_SNAKE_CASE	:		List[Any]=1e-1_2     ,    SCREAMING_SNAKE_CASE	:		int="absolute"     ,    SCREAMING_SNAKE_CASE	:		str=True     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=None     ,    SCREAMING_SNAKE_CASE	:		int=False     ,    SCREAMING_SNAKE_CASE	:		Tuple=False     ,    SCREAMING_SNAKE_CASE	:		List[Any]=None     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=None     ,    **SCREAMING_SNAKE_CASE	:		str     ,    )			->	int:




											'''simple docstring'''

											super().__init__(pad_token_id=SCREAMING_SNAKE_CASE     ,    mask_token_id=SCREAMING_SNAKE_CASE     ,    **SCREAMING_SNAKE_CASE						)

											__snake_case				=		vocab_size
											__snake_case				=		hidden_size
											__snake_case				=		num_hidden_layers
											__snake_case				=		num_attention_heads
											__snake_case				=		intermediate_size
											__snake_case				=		hidden_dropout_prob
											__snake_case				=		attention_probs_dropout_prob
											__snake_case				=		max_position_embeddings
											__snake_case				=		initializer_range
											__snake_case				=		layer_norm_eps
											__snake_case				=		position_embedding_type
											__snake_case				=		use_cache
											__snake_case				=		emb_layer_norm_before
											__snake_case				=		token_dropout
											__snake_case				=		is_folding_model
											if is_folding_model:
																	if esmfold_config is None:
																							logger.info("No esmfold_config supplied for folding model, using default values."						)
																							__snake_case				=		EsmFoldConfig()
																	elif isinstance(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE						):
																							__snake_case				=		EsmFoldConfig(**SCREAMING_SNAKE_CASE						)
																	__snake_case				=		esmfold_config
																	if vocab_list is None:
																							logger.warning("No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"						)
																							__snake_case				=		get_default_vocab_list()
																	else:
																							__snake_case				=		vocab_list
											else:
																	__snake_case				=		None
																	__snake_case				=		None
											if self.esmfold_config is not None and getattr(self.esmfold_config     ,    "use_esm_attn_map"     ,    SCREAMING_SNAKE_CASE						):
																	raise ValueError("The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"						)






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		super().to_dict()
											if isinstance(self.esmfold_config     ,    SCREAMING_SNAKE_CASE						):
																	__snake_case				=		self.esmfold_config.to_dict()
											return output






@dataclass
class     UpperCamelCase:
					snake_case_		:				str										=	None
					snake_case_		:				bool										=	True
					snake_case_		:				bool										=	False
					snake_case_		:				bool										=	False
					snake_case_		:				bool										=	False
					snake_case_		:				float										=	0

					snake_case_		:				bool										=	True
					snake_case_		:				bool										=	False

					snake_case_		:				int										=	128
					snake_case_		:				"TrunkConfig"										=	None


					def 							SCREAMING_SNAKE_CASE_  (						self	:		str						)			->	List[Any]:




											'''simple docstring'''

											if self.trunk is None:
																	__snake_case				=		TrunkConfig()
											elif isinstance(self.trunk     ,    SCREAMING_SNAKE_CASE						):
																	__snake_case				=		TrunkConfig(**self.trunk						)






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any						)			->	Optional[Any]:




											'''simple docstring'''

											__snake_case				=		asdict(self						)
											__snake_case				=		self.trunk.to_dict()
											return output






@dataclass
class     UpperCamelCase:
					snake_case_		:				int										=	48
					snake_case_		:				int										=	1024
					snake_case_		:				int										=	128
					snake_case_		:				int										=	32
					snake_case_		:				int										=	32
					snake_case_		:				int										=	32
					snake_case_		:				float										=	0
					snake_case_		:				float										=	0
					snake_case_		:				bool										=	False
					snake_case_		:				int										=	4
					snake_case_		:				Optional[int]										=	128
					snake_case_		:				"StructureModuleConfig"										=	None


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Dict:




											'''simple docstring'''

											if self.structure_module is None:
																	__snake_case				=		StructureModuleConfig()
											elif isinstance(self.structure_module     ,    SCREAMING_SNAKE_CASE						):
																	__snake_case				=		StructureModuleConfig(**self.structure_module						)

											if self.max_recycles <= 0:
																	raise ValueError(f'''`max_recycles` should be positive, got {self.max_recycles}.'''						)
											if self.sequence_state_dim % self.sequence_state_dim != 0:
																	raise ValueError(
																	    "`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"
																	    f''' {self.sequence_state_dim} and {self.sequence_state_dim}.'''						)
											if self.pairwise_state_dim % self.pairwise_state_dim != 0:
																	raise ValueError(
																	    "`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"
																	    f''' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'''						)

											__snake_case				=		self.sequence_state_dim // self.sequence_head_width
											__snake_case				=		self.pairwise_state_dim // self.pairwise_head_width

											if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
																	raise ValueError(
																	    "`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"
																	    f''' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'''						)
											if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
																	raise ValueError(
																	    "`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"
																	    f''' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'''						)
											if self.pairwise_state_dim % 2 != 0:
																	raise ValueError(f'''`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'''						)

											if self.dropout >= 0.4:
																	raise ValueError(f'''`dropout` should not be greater than 0.4, got {self.dropout}.'''						)






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		asdict(self						)
											__snake_case				=		self.structure_module.to_dict()
											return output






@dataclass
class     UpperCamelCase:
					snake_case_		:				int										=	384
					snake_case_		:				int										=	128
					snake_case_		:				int										=	16
					snake_case_		:				int										=	128
					snake_case_		:				int										=	12
					snake_case_		:				int										=	4
					snake_case_		:				int										=	8
					snake_case_		:				float										=	0.1
					snake_case_		:				int										=	8
					snake_case_		:				int										=	1
					snake_case_		:				int										=	2
					snake_case_		:				int										=	7
					snake_case_		:				int										=	10
					snake_case_		:				float										=	1E-8
					snake_case_		:				float										=	1E5


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[Any]						)			->	Any:




											'''simple docstring'''

											return asdict(self						)





def 	_lowerCAmelCase       (    )						->	Tuple:


						'''simple docstring'''






						return (
						    "<cls>",
						    "<pad>",
						    "<eos>",
						    "<unk>",
						    "L",
						    "A",
						    "G",
						    "V",
						    "S",
						    "E",
						    "R",
						    "T",
						    "I",
						    "D",
						    "P",
						    "K",
						    "Q",
						    "N",
						    "F",
						    "Y",
						    "M",
						    "H",
						    "W",
						    "C",
						    "X",
						    "B",
						    "U",
						    "Z",
						    "O",
						    ".",
						    "-",
						    "<null_1>",
						    "<mask>",
						)









import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
	import onnxruntime as ort






class     UpperCamelCase(  _a     ,							unittest.TestCase   ):
					# TODO: is there an appropriate internal test set?
					snake_case_		:				Dict										=	"""ssube/stable-diffusion-x4-upscaler-onnx"""


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict     ,    SCREAMING_SNAKE_CASE	:		List[str]=0						)			->	Union[str, Any]:




											'''simple docstring'''

											__snake_case				=		floats_tensor((1, 3, 1_2_8, 1_2_8)     ,    rng=random.Random(SCREAMING_SNAKE_CASE						)						)
											__snake_case				=		torch.manual_seed(SCREAMING_SNAKE_CASE						)
											__snake_case				=		{
											    "prompt": "A painting of a squirrel eating a burger",
											    "image": image,
											    "generator": generator,
											    "num_inference_steps": 3,
											    "guidance_scale": 7.5,
											    "output_type": "numpy",
											}
											return inputs


					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	List[str]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1].flatten()

											# started as 128, should now be 512
											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
											assert np.abs(image_slice - expected_slice						).max() < 1e-1


					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	Optional[Any]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		PNDMScheduler.from_config(pipe.scheduler.config     ,    skip_prk_steps=SCREAMING_SNAKE_CASE						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]						)
											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	Dict:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1






					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[Any]						)			->	List[str]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1






@nightly
@require_onnxruntime
@require_torch_gpu
class     UpperCamelCase(  unittest.TestCase   ):


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		str						)			->	List[Any]:




											'''simple docstring'''

											return (
											    "CUDAExecutionProvider",
											    {
											        "gpu_mem_limit": "15000000000",  # 15GB
											        "arena_extend_strategy": "kSameAsRequested",
											    },
											)


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Tuple:




											'''simple docstring'''

											__snake_case				=		ort.SessionOptions()
											__snake_case				=		False
											return options


					def 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"						)
											__snake_case				=		init_image.resize((1_2_8, 1_2_8)						)
											# using the PNDM scheduler by default
											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    provider=self.gpu_provider     ,    sess_options=self.gpu_options     ,    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		"A fantasy landscape, trending on artstation"

											__snake_case				=		torch.manual_seed(0						)
											__snake_case				=		pipe(
											    prompt=SCREAMING_SNAKE_CASE     ,    image=SCREAMING_SNAKE_CASE     ,    guidance_scale=7.5     ,    num_inference_steps=1_0     ,    generator=SCREAMING_SNAKE_CASE     ,    output_type="np"     ,    )
											__snake_case				=		output.images
											__snake_case				=		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert images.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]						)
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"						)
											__snake_case				=		init_image.resize((1_2_8, 1_2_8)						)
											__snake_case				=		LMSDiscreteScheduler.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    subfolder="scheduler"						)
											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    scheduler=SCREAMING_SNAKE_CASE     ,    provider=self.gpu_provider     ,    sess_options=self.gpu_options     ,    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		"A fantasy landscape, trending on artstation"

											__snake_case				=		torch.manual_seed(0						)
											__snake_case				=		pipe(
											    prompt=SCREAMING_SNAKE_CASE     ,    image=SCREAMING_SNAKE_CASE     ,    guidance_scale=7.5     ,    num_inference_steps=2_0     ,    generator=SCREAMING_SNAKE_CASE     ,    output_type="np"     ,    )
											__snake_case				=		output.images
											__snake_case				=		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert images.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]						)
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2





