





"""simple docstring"""

import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__UpperCAmelCase        =logging.get_logger(__name__)

enable_full_determinism()


class       lowerCAmelCase__				(   lowercase__	,					lowercase__	,					unittest.TestCase  ):
 lowercase__     :  Tuple				=   UNetaDModel
 lowercase__     :  Union[str, Any]				=   """sample"""







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					4
    A__											=					3
    A__											=					(32, 32)

    A__											=					floats_tensor((batch_size, num_channels) + sizes					).to(UpperCAmelCase__					)
    A__											=					torch.tensor([10]					).to(UpperCAmelCase__					)

    return {"sample": noise, "timestep": time_step}







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (3, 32, 32)







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (3, 32, 32)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					{
        '''block_out_channels''': (32, 64),
        '''down_block_types''': ('''DownBlock2D''', '''AttnDownBlock2D'''),
        '''up_block_types''': ('''AttnUpBlock2D''', '''UpBlock2D'''),
        '''attention_head_dim''': 3,
        '''out_channels''': 3,
        '''in_channels''': 3,
        '''layers_per_block''': 2,
        '''sample_size''': 32,
    }
    A__											=					self.dummy_input
    return init_dict, inputs_dict





class       lowerCAmelCase__				(   lowercase__	,					lowercase__	,					unittest.TestCase  ):
 lowercase__     :  Union[str, Any]				=   UNetaDModel
 lowercase__     :  Optional[int]				=   """sample"""







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					4
    A__											=					4
    A__											=					(32, 32)

    A__											=					floats_tensor((batch_size, num_channels) + sizes					).to(UpperCAmelCase__					)
    A__											=					torch.tensor([10]					).to(UpperCAmelCase__					)

    return {"sample": noise, "timestep": time_step}







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (4, 32, 32)







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (4, 32, 32)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					{
        '''sample_size''': 32,
        '''in_channels''': 4,
        '''out_channels''': 4,
        '''layers_per_block''': 2,
        '''block_out_channels''': (32, 64),
        '''attention_head_dim''': 32,
        '''down_block_types''': ('''DownBlock2D''', '''DownBlock2D'''),
        '''up_block_types''': ('''UpBlock2D''', '''UpBlock2D'''),
    }
    A__											=					self.dummy_input
    return init_dict, inputs_dict







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"		,	output_loading_info=UpperCAmelCase__					)

    self.assertIsNotNone(UpperCAmelCase__					)
    self.assertEqual(len(loading_info["missing_keys"]					)		,	0					)

    model.to(UpperCAmelCase__					)
    A__											=					model(**self.dummy_input					).sample

    assert image is not None, "Make sure output is not None"







 @unittest.skipIf(torch_device != "cuda"		,	"This test is supposed to run on GPU"					)
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"		,	output_loading_info=UpperCAmelCase__					)
    model.to(UpperCAmelCase__					)
    A__											=					model(**self.dummy_input					).sample

    assert image is not None, "Make sure output is not None"







 @unittest.skipIf(torch_device != "cuda"		,	"This test is supposed to run on GPU"					)
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"		,	output_loading_info=UpperCAmelCase__					)
    model_accelerate.to(UpperCAmelCase__					)
    model_accelerate.eval()

    A__											=					torch.randn(
        1		,	model_accelerate.config.in_channels		,	model_accelerate.config.sample_size		,	model_accelerate.config.sample_size		,	generator=torch.manual_seed(0					)		,	)
    A__											=					noise.to(UpperCAmelCase__					)
    A__											=					torch.tensor([10] * noise.shape[0]					).to(UpperCAmelCase__					)

    A__											=					model_accelerate(UpperCAmelCase__		,	UpperCAmelCase__					)['''sample''']

    # two models don't need to stay in the device at the same time
    del model_accelerate
    torch.cuda.empty_cache()
    gc.collect()

    A__											=					UNetaDModel.from_pretrained(
        "fusing/unet-ldm-dummy-update"		,	output_loading_info=UpperCAmelCase__		,	low_cpu_mem_usage=UpperCAmelCase__					)
    model_normal_load.to(UpperCAmelCase__					)
    model_normal_load.eval()
    A__											=					model_normal_load(UpperCAmelCase__		,	UpperCAmelCase__					)['''sample''']

    assert torch_all_close(UpperCAmelCase__		,	UpperCAmelCase__		,	rtol=1e-3					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"					)
    model.eval()
    model.to(UpperCAmelCase__					)

    A__											=					torch.randn(
        1		,	model.config.in_channels		,	model.config.sample_size		,	model.config.sample_size		,	generator=torch.manual_seed(0					)		,	)
    A__											=					noise.to(UpperCAmelCase__					)
    A__											=					torch.tensor([10] * noise.shape[0]					).to(UpperCAmelCase__					)

    with torch.no_grad():
       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					).sample

    A__											=					output[0, -1, -3:, -3:].flatten().cpu()
    # fmt: off
    A__											=					torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]					)
    # fmt: on

    self.assertTrue(torch_all_close(UpperCAmelCase__		,	UpperCAmelCase__		,	rtol=1e-3					)					)





class       lowerCAmelCase__				(   lowercase__	,					lowercase__	,					unittest.TestCase  ):
 lowercase__     :  Optional[int]				=   UNetaDModel
 lowercase__     :  str				=   """sample"""







 @property
 def  lowercase_       (   self		,	UpperCamelCase__=(32, 32)					):

    '''simple docstring'''


    A__											=					4
    A__											=					3

    A__											=					floats_tensor((batch_size, num_channels) + sizes					).to(UpperCAmelCase__					)
    A__											=					torch.tensor(batch_size * [10]					).to(dtype=torch.intaa		,	device=UpperCAmelCase__					)

    return {"sample": noise, "timestep": time_step}







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (3, 32, 32)







 @property
 def  lowercase_       (   self					):

    '''simple docstring'''


    return (3, 32, 32)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					{
        '''block_out_channels''': [32, 64, 64, 64],
        '''in_channels''': 3,
        '''layers_per_block''': 1,
        '''out_channels''': 3,
        '''time_embedding_type''': '''fourier''',
        '''norm_eps''': 1e-6,
        '''mid_block_scale_factor''': math.sqrt(2.0					),
        '''norm_num_groups''': None,
        '''down_block_types''': [
            '''SkipDownBlock2D''',
            '''AttnSkipDownBlock2D''',
            '''SkipDownBlock2D''',
            '''SkipDownBlock2D''',
        ],
        '''up_block_types''': [
            '''SkipUpBlock2D''',
            '''SkipUpBlock2D''',
            '''AttnSkipUpBlock2D''',
            '''SkipUpBlock2D''',
        ],
    }
    A__											=					self.dummy_input
    return init_dict, inputs_dict







 @slow
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"		,	output_loading_info=UpperCAmelCase__					)
    self.assertIsNotNone(UpperCAmelCase__					)
    self.assertEqual(len(loading_info["missing_keys"]					)		,	0					)

    model.to(UpperCAmelCase__					)
    A__											=					self.dummy_input
    A__											=					floats_tensor((4, 3) + (2_56, 2_56)					).to(UpperCAmelCase__					)
    A__											=					noise
    A__											=					model(**UpperCAmelCase__					)

    assert image is not None, "Make sure output is not None"







 @slow
 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"					)
    model.to(UpperCAmelCase__					)

    A__											=					4
    A__											=					3
    A__											=					(2_56, 2_56)

    A__											=					torch.ones((batch_size, num_channels) + sizes					).to(UpperCAmelCase__					)
    A__											=					torch.tensor(batch_size * [1e-4]					).to(UpperCAmelCase__					)

    with torch.no_grad():
       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					).sample

    A__											=					output[0, -3:, -3:, -1].flatten().cpu()
    # fmt: off
    A__											=					torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]					)
    # fmt: on

    self.assertTrue(torch_all_close(UpperCAmelCase__		,	UpperCAmelCase__		,	rtol=1e-2					)					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update"					)
    model.to(UpperCAmelCase__					)

    A__											=					4
    A__											=					3
    A__											=					(32, 32)

    A__											=					torch.ones((batch_size, num_channels) + sizes					).to(UpperCAmelCase__					)
    A__											=					torch.tensor(batch_size * [1e-4]					).to(UpperCAmelCase__					)

    with torch.no_grad():
       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					).sample

    A__											=					output[0, -3:, -3:, -1].flatten().cpu()
    # fmt: off
    A__											=					torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]					)
    # fmt: on

    self.assertTrue(torch_all_close(UpperCAmelCase__		,	UpperCAmelCase__		,	rtol=1e-2					)					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    pass


'''simple docstring'''





from dataclasses import asdict, dataclass
from typing import Optional

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


UpperCamelCase_    						=       logging.get_logger(__name__)

# TODO Update this
UpperCamelCase_    						=       {
    """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 		__SCREAMING_SNAKE_CASE							(			lowercase__      ):
				lowerCamelCase_		       =    'esm'







				def __init__(       self     :		Optional[Any]		,					UpperCAmelCase__     :		str=None		,					UpperCAmelCase__     :		Dict=None		,					UpperCAmelCase__     :		Any=None		,					UpperCAmelCase__     :		int=768		,					UpperCAmelCase__     :		Union[str, Any]=12		,					UpperCAmelCase__     :		List[Any]=12		,					UpperCAmelCase__     :		Optional[Any]=3072		,					UpperCAmelCase__     :		Dict=0.1		,					UpperCAmelCase__     :		int=0.1		,					UpperCAmelCase__     :		int=1026		,					UpperCAmelCase__     :		Tuple=0.02		,					UpperCAmelCase__     :		Optional[int]=1E-12		,					UpperCAmelCase__     :		List[Any]="absolute"		,					UpperCAmelCase__     :		int=True		,					UpperCAmelCase__     :		Dict=None		,					UpperCAmelCase__     :		str=False		,					UpperCAmelCase__     :		int=False		,					UpperCAmelCase__     :		Optional[Any]=None		,					UpperCAmelCase__     :		List[str]=None		,					**UpperCAmelCase__     :		int		,					):
										'''simple docstring'''
										super().__init__(pad_token_id=UpperCAmelCase__		,					mask_token_id=UpperCAmelCase__		,					**UpperCAmelCase__   )

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







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										lowercase :							Union[str, Any]    		=super().to_dict()
										if isinstance(self.esmfold_config		,					UpperCAmelCase__   ):
																lowercase :							Optional[Any]    		=self.esmfold_config.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    None
				lowerCamelCase_		       =    True
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    0

				lowerCamelCase_		       =    True
				lowerCamelCase_		       =    False

				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    None







				def  lowerCamelCase_	(       self     :		Tuple   ):
										'''simple docstring'''
										if self.trunk is None:
																lowercase :							str    		=TrunkConfig()
										elif isinstance(self.trunk		,					UpperCAmelCase__   ):
																lowercase :							int    		=TrunkConfig(**self.trunk   )







				def  lowerCamelCase_	(       self     :		List[Any]   ):
										'''simple docstring'''
										lowercase :							str    		=asdict(self   )
										lowercase :							Union[str, Any]    		=self.trunk.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    48
				lowerCamelCase_		       =    10_24
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    0
				lowerCamelCase_		       =    0
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    4
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    None







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										if self.structure_module is None:
																lowercase :							Any    		=StructureModuleConfig()
										elif isinstance(self.structure_module		,					UpperCAmelCase__   ):
																lowercase :							Union[str, Any]    		=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}.'''   )

										lowercase :							str    		=self.sequence_state_dim // self.sequence_head_width
										lowercase :							int    		=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  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										lowercase :							List[Any]    		=asdict(self   )
										lowercase :							Any    		=self.structure_module.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    3_84
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    16
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    12
				lowerCamelCase_		       =    4
				lowerCamelCase_		       =    8
				lowerCamelCase_		       =    0.1
				lowerCamelCase_		       =    8
				lowerCamelCase_		       =    1
				lowerCamelCase_		       =    2
				lowerCamelCase_		       =    7
				lowerCamelCase_		       =    10
				lowerCamelCase_		       =    1E-8
				lowerCamelCase_		       =    1E5







				def  lowerCamelCase_	(       self     :		Dict   ):
										'''simple docstring'''
										return asdict(self   )






def 				_lowerCAmelCase			(   )				->  Optional[int]:
						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>",
						)





