




'''simple docstring'''



import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers.models.esm.modeling_esmfold import EsmForProteinFolding





class      lowerCAmelCase :


       def __init__( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=13     ,	__SCREAMING_SNAKE_CASE=7     ,	__SCREAMING_SNAKE_CASE=False     ,	__SCREAMING_SNAKE_CASE=True     ,	__SCREAMING_SNAKE_CASE=False     ,	__SCREAMING_SNAKE_CASE=False     ,	__SCREAMING_SNAKE_CASE=19     ,	__SCREAMING_SNAKE_CASE=32     ,	__SCREAMING_SNAKE_CASE=5     ,	__SCREAMING_SNAKE_CASE=4     ,	__SCREAMING_SNAKE_CASE=37     ,	__SCREAMING_SNAKE_CASE="gelu"     ,	__SCREAMING_SNAKE_CASE=0.1     ,	__SCREAMING_SNAKE_CASE=0.1     ,	__SCREAMING_SNAKE_CASE=512     ,	__SCREAMING_SNAKE_CASE=16     ,	__SCREAMING_SNAKE_CASE=2     ,	__SCREAMING_SNAKE_CASE=0.02     ,	__SCREAMING_SNAKE_CASE=3     ,	__SCREAMING_SNAKE_CASE=4     ,	__SCREAMING_SNAKE_CASE=None     ,	)   ->	List[Any]:



            '''simple docstring'''




            __snake_case            =	parent
            __snake_case            =	batch_size
            __snake_case            =	seq_length
            __snake_case            =	is_training
            __snake_case            =	use_input_mask
            __snake_case            =	use_token_type_ids
            __snake_case            =	use_labels
            __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_act
            __snake_case            =	hidden_dropout_prob
            __snake_case            =	attention_probs_dropout_prob
            __snake_case            =	max_position_embeddings
            __snake_case            =	type_vocab_size
            __snake_case            =	type_sequence_label_size
            __snake_case            =	initializer_range
            __snake_case            =	num_labels
            __snake_case            =	num_choices
            __snake_case            =	scope


       def    lowerCAmelCase   ( self	)   ->	Dict:



            '''simple docstring'''




            __snake_case            =	ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size	)

            __snake_case            =	None
            if self.use_input_mask:
                 __snake_case            =	random_attention_mask([self.batch_size, self.seq_length]	)

            __snake_case            =	None
            __snake_case            =	None
            __snake_case            =	None
            if self.use_labels:
                 __snake_case            =	ids_tensor([self.batch_size]     ,	self.type_sequence_label_size	)
                 __snake_case            =	ids_tensor([self.batch_size, self.seq_length]     ,	self.num_labels	)
                 __snake_case            =	ids_tensor([self.batch_size]     ,	self.num_choices	)

            __snake_case            =	self.get_config()

            return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


       def    lowerCAmelCase   ( self	)   ->	Any:



            '''simple docstring'''




            __snake_case            =	EsmConfig(
                vocab_size=33     ,	hidden_size=self.hidden_size     ,	pad_token_id=1     ,	num_hidden_layers=self.num_hidden_layers     ,	num_attention_heads=self.num_attention_heads     ,	intermediate_size=self.intermediate_size     ,	hidden_act=self.hidden_act     ,	hidden_dropout_prob=self.hidden_dropout_prob     ,	attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,	max_position_embeddings=self.max_position_embeddings     ,	type_vocab_size=self.type_vocab_size     ,	initializer_range=self.initializer_range     ,	is_folding_model=__SCREAMING_SNAKE_CASE     ,	esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}     ,	)
            return config


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Optional[Any]:



            '''simple docstring'''




            __snake_case            =	EsmForProteinFolding(config=__SCREAMING_SNAKE_CASE	).float()
            model.to(__SCREAMING_SNAKE_CASE	)
            model.eval()
            __snake_case            =	model(__SCREAMING_SNAKE_CASE     ,	attention_mask=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	model(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	model(__SCREAMING_SNAKE_CASE	)

            self.parent.assertEqual(result.positions.shape     ,	(8, self.batch_size, self.seq_length, 14, 3)	)
            self.parent.assertEqual(result.angles.shape     ,	(8, self.batch_size, self.seq_length, 7, 2)	)




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



            '''simple docstring'''




            __snake_case            =	self.prepare_config_and_inputs()
            (
                (
                __snake_case
            )       ,     (
                __snake_case
            )       ,     (
                __snake_case
            )       ,     (
                __snake_case
            )       ,     (
                __snake_case
            )       ,     (
                __snake_case
            )       ,     
            )            =	config_and_inputs
            __snake_case            =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
            return config, inputs_dict





@require_torch
class      lowerCAmelCase (			__lowerCAmelCase					, __lowerCAmelCase					, unittest.TestCase):
       __lowercase	:   List[str]										=      False

       __lowercase	:   Dict										=      (EsmForProteinFolding,) if is_torch_available() else ()
       __lowercase	:   Union[str, Any]										=      ()
       __lowercase	:   List[str]										=      {} if is_torch_available() else {}
       __lowercase	:   List[Any]										=      False


       def    lowerCAmelCase   ( self	)   ->	Any:



            '''simple docstring'''




            __snake_case            =	EsmFoldModelTester(self	)
            __snake_case            =	ConfigTester(self     ,	config_class=__SCREAMING_SNAKE_CASE     ,	hidden_size=37	)


       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            self.config_tester.run_common_tests()


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



            '''simple docstring'''




            __snake_case            =	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE	)


       @unittest.skip('''Does not support attention outputs'''	)
       def    lowerCAmelCase   ( self	)   ->	Dict:



            '''simple docstring'''




            pass


       @unittest.skip
       def    lowerCAmelCase   ( self	)   ->	Union[str, Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''Esm does not support embedding resizing'''	)
       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''Esm does not support embedding resizing'''	)
       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support passing input embeds!'''	)
       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support head pruning.'''	)
       def    lowerCAmelCase   ( self	)   ->	Optional[int]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support head pruning.'''	)
       def    lowerCAmelCase   ( self	)   ->	List[str]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support head pruning.'''	)
       def    lowerCAmelCase   ( self	)   ->	int:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support head pruning.'''	)
       def    lowerCAmelCase   ( self	)   ->	Optional[Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support head pruning.'''	)
       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not output hidden states in the normal way.'''	)
       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMfold does not output hidden states in the normal way.'''	)
       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold only has one output format.'''	)
       def    lowerCAmelCase   ( self	)   ->	Optional[Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''	)
       def    lowerCAmelCase   ( self	)   ->	Dict:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold does not support input chunking.'''	)
       def    lowerCAmelCase   ( self	)   ->	Optional[Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''	)
       def    lowerCAmelCase   ( self	)   ->	Optional[int]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''	)
       def    lowerCAmelCase   ( self	)   ->	Any:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''	)
       def    lowerCAmelCase   ( self	)   ->	int:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''	)
       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            pass


       @unittest.skip('''ESMFold doesn\'t support data parallel.'''	)
       def    lowerCAmelCase   ( self	)   ->	Any:



            '''simple docstring'''




            pass




       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''	)
       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




            pass





@require_torch
class      lowerCAmelCase (			__lowerCAmelCase):


       @slow
       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            __snake_case            =	EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''	).float()
            model.eval()
            __snake_case            =	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
            __snake_case            =	model(__SCREAMING_SNAKE_CASE	)['''positions''']
            __snake_case            =	torch.tensor([2.5_828, 0.7_993, -10.9_334]     ,	dtype=torch.floataa	)
            self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]     ,	__SCREAMING_SNAKE_CASE     ,	atol=1E-4	)	)














"""simple docstring"""



import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
 import jax
 import jax.numpy as jnp
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard




@slow
@require_flax
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()




   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/sd2-inpaint/init_image.png"							)
        UpperCAmelCase__						:    Any							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"							)

        UpperCAmelCase__						:    str							=							"xvjiarui/stable-diffusion-2-inpainting"
        UpperCAmelCase__							,				UpperCAmelCase__						:    List[Any]							=							FlaxStableDiffusionInpaintPipeline.from_pretrained(snake_case__							,			safety_checker=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							"Face of a yellow cat, high resolution, sitting on a park bench"

        UpperCAmelCase__						:    List[str]							=							jax.random.PRNGKey(0							)
        UpperCAmelCase__						:    Optional[int]							=							5_0

        UpperCAmelCase__						:    List[Any]							=							jax.device_count()
        UpperCAmelCase__						:    int							=							num_samples * [prompt]
        UpperCAmelCase__						:    str							=							num_samples * [init_image]
        UpperCAmelCase__						:    Optional[Any]							=							num_samples * [mask_image]
        UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__						:    Optional[int]							=							pipeline.prepare_inputs(snake_case__							,			snake_case__							,			snake_case__							)

        # shard inputs and rng
        UpperCAmelCase__						:    Dict							=							replicate(snake_case__							)
        UpperCAmelCase__						:    Dict							=							jax.random.split(snake_case__							,			jax.device_count()							)
        UpperCAmelCase__						:    List[str]							=							shard(snake_case__							)
        UpperCAmelCase__						:    str							=							shard(snake_case__							)
        UpperCAmelCase__						:    Optional[int]							=							shard(snake_case__							)

        UpperCAmelCase__						:    Dict							=							pipeline(
            snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			jit=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							output.images.reshape(snake_case__							,			5_1_2							,			5_1_2							,			3							)

        UpperCAmelCase__						:    Tuple							=							images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

        UpperCAmelCase__						:    str							=							jnp.asarray(jax.device_get(image_slice.flatten()							)							)
        UpperCAmelCase__						:    int							=							jnp.array(
            [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084]							)
        print(f'output_slice: {output_slice}'							)
        assert jnp.abs(output_slice - expected_slice							).max() < 1e-2


