

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 import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
     from transformers.models.esm.modeling_esm import (
         ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
         EsmEmbeddings,
         create_position_ids_from_input_ids,
     )


class  A						:







    """simple docstring"""





    def __init__(      self :       int,lowercase_ :       List[str],lowercase_ :       List[str]=1_3,lowercase_ :       Tuple=7,lowercase_ :       Any=False,lowercase_ :       Optional[int]=True,lowercase_ :       str=False,lowercase_ :       str=True,lowercase_ :       Tuple=3_3,lowercase_ :       Any=3_2,lowercase_ :       int=5,lowercase_ :       Optional[Any]=4,lowercase_ :       int=3_7,lowercase_ :       Tuple="gelu",lowercase_ :       List[Any]=0.1,lowercase_ :       Tuple=0.1,lowercase_ :       Any=5_1_2,lowercase_ :       Any=1_6,lowercase_ :       List[str]=2,lowercase_ :       Any=0.02,lowercase_ :       Union[str, Any]=3,lowercase_ :       str=4,lowercase_ :       int=None,)->    Optional[int]:





        '''simple docstring'''







        A__          =   parent
        A__          =   batch_size
        A__          =   seq_length
        A__          =   is_training
        A__          =   use_input_mask
        A__          =   use_token_type_ids
        A__          =   use_labels
        A__          =   vocab_size
        A__          =   hidden_size
        A__          =   num_hidden_layers
        A__          =   num_attention_heads
        A__          =   intermediate_size
        A__          =   hidden_act
        A__          =   hidden_dropout_prob
        A__          =   attention_probs_dropout_prob
        A__          =   max_position_embeddings
        A__          =   type_vocab_size
        A__          =   type_sequence_label_size
        A__          =   initializer_range
        A__          =   num_labels
        A__          =   num_choices
        A__          =   scope


    def 						snake_case__     (      self :       Tuple		)->    Tuple:





        '''simple docstring'''







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

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

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

        A__          =   self.get_config()

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


    def 						snake_case__     (      self :       Dict		)->    Any:





        '''simple docstring'''







        return EsmConfig(
            vocab_size=self.vocab_size,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,)


    def 						snake_case__     (      self :       Any,lowercase_ :       Tuple,lowercase_ :       List[str],lowercase_ :       Union[str, Any],lowercase_ :       List[str],lowercase_ :       Tuple,lowercase_ :       str		)->    str:





        '''simple docstring'''







        A__          =   EsmModel(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__          =   model(lowercase_,attention_mask=lowercase_		)
        A__          =   model(lowercase_		)
        A__          =   model(lowercase_		)

        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)		)
        self.parent.assertEqual(result.pooler_output.shape,(self.batch_size, self.hidden_size)		)


    def 						snake_case__     (      self :       Tuple,lowercase_ :       str,lowercase_ :       Dict,lowercase_ :       Dict,lowercase_ :       str,lowercase_ :       Union[str, Any],lowercase_ :       Any		)->    List[Any]:





        '''simple docstring'''







        A__          =   EsmForMaskedLM(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__          =   model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size)		)


    def 						snake_case__     (      self :       Tuple,lowercase_ :       str,lowercase_ :       Dict,lowercase_ :       List[str],lowercase_ :       Any,lowercase_ :       str,lowercase_ :       str		)->    List[Any]:





        '''simple docstring'''







        A__          =   self.num_labels
        A__          =   EsmForTokenClassification(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__          =   model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.num_labels)		)


    def 						snake_case__     (      self :       int		)->    List[Any]:





        '''simple docstring'''







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




@require_torch
class  A						(							_UpperCAmelCase    ,							_UpperCAmelCase    ,							unittest.TestCase ):







    """simple docstring"""





    lowerCamelCase			=							False

    lowerCamelCase			=							(
        (
            EsmForMaskedLM,
            EsmModel,
            EsmForSequenceClassification,
            EsmForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowerCamelCase			=							()
    lowerCamelCase			=							(
        {
            'feature-extraction': EsmModel,
            'fill-mask': EsmForMaskedLM,
            'text-classification': EsmForSequenceClassification,
            'token-classification': EsmForTokenClassification,
            'zero-shot': EsmForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowerCamelCase			=							True


    def 						snake_case__     (      self :       str		)->    str:





        '''simple docstring'''







        A__          =   EsmModelTester(self		)
        A__          =   ConfigTester(self,config_class=lowercase_,hidden_size=3_7		)


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





        '''simple docstring'''







        self.config_tester.run_common_tests()


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





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase_		)


    def 						snake_case__     (      self :       Tuple		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
            A__          =   type
            self.model_tester.create_and_check_model(*lowercase_		)


    def 						snake_case__     (      self :       Union[str, Any]		)->    Tuple:





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowercase_		)


    def 						snake_case__     (      self :       List[str]		)->    Tuple:





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowercase_		)


    @slow
    def 						snake_case__     (      self :       Tuple		)->    List[str]:





        '''simple docstring'''







        for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            A__          =   EsmModel.from_pretrained(lowercase_		)
            self.assertIsNotNone(lowercase_		)


    def 						snake_case__     (      self :       Tuple		)->    List[str]:





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()[0]
        A__          =   EsmEmbeddings(config=lowercase_		)

        A__          =   torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]		)
        A__          =   torch.as_tensor(
            [
                [
                    0 + model.padding_idx + 1,
                    1 + model.padding_idx + 1,
                    2 + model.padding_idx + 1,
                    model.padding_idx,
                ]
            ]		)
        A__          =   create_position_ids_from_input_ids(lowercase_,model.padding_idx		)
        self.assertEqual(position_ids.shape,expected_positions.shape		)
        self.assertTrue(torch.all(torch.eq(lowercase_,lowercase_		)		)		)


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





        '''simple docstring'''







        A__          =   self.model_tester.prepare_config_and_inputs()[0]
        A__          =   EsmEmbeddings(config=lowercase_		)

        A__          =   torch.empty(2,4,3_0		)
        A__          =   [
            0 + embeddings.padding_idx + 1,
            1 + embeddings.padding_idx + 1,
            2 + embeddings.padding_idx + 1,
            3 + embeddings.padding_idx + 1,
        ]
        A__          =   torch.as_tensor([expected_single_positions, expected_single_positions]		)
        A__          =   embeddings.create_position_ids_from_inputs_embeds(lowercase_		)
        self.assertEqual(position_ids.shape,expected_positions.shape		)
        self.assertTrue(torch.all(torch.eq(lowercase_,lowercase_		)		)		)


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





        '''simple docstring'''







        pass


    @unittest.skip('Esm does not support embedding resizing'		)
    def 						snake_case__     (      self :       Optional[int]		)->    int:





        '''simple docstring'''







        pass


    @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'		)
    def 						snake_case__     (      self :       Optional[Any]		)->    Any:





        '''simple docstring'''







        pass




@require_torch
class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    @slow
    def 						snake_case__     (      self :       str		)->    int:





        '''simple docstring'''







        with torch.no_grad():
            A__          =   EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'		)
            model.eval()
            A__          =   torch.tensor([[0, 1, 2, 3, 4, 5]]		)
            A__          =   model(lowercase_		)[0]

            A__          =   3_3

            A__          =   torch.Size((1, 6, vocab_size)		)
            self.assertEqual(output.shape,lowercase_		)

            A__          =   torch.tensor(
                [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]		)
            self.assertTrue(torch.allclose(output[:, :3, :3],lowercase_,atol=1E-4		)		)


    @slow
    def 						snake_case__     (      self :       int		)->    Any:





        '''simple docstring'''







        with torch.no_grad():
            A__          =   EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'		)
            model.eval()

            A__          =   torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]		)
            A__          =   model(lowercase_		)[0]
            # compare the actual values for a slice.
            A__          =   torch.tensor(
                [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]		)
            self.assertTrue(torch.allclose(output[:, :3, :3],lowercase_,atol=1E-4		)		)


"""simple docstring"""






import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


__A			   = False


class     UpperCAmelCase						(unittest.TestCase  ):



    """simple docstring"""

    pass



@nightly
@require_torch_gpu
class     UpperCAmelCase						(unittest.TestCase  ):



    """simple docstring"""

    def _snake_case			(   self							):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






    def _snake_case			(   self							):
       lowercase__:   Dict					=      VersatileDiffusionTextToImagePipeline.from_pretrained('''shi-labs/versatile-diffusion'''							)
       # remove text_unet
       pipe.remove_unused_weights()
       pipe.to(_UpperCAmelCase							)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase							)

       lowercase__:   List[str]					=      '''A painting of a squirrel eating a burger '''
       lowercase__:   str					=      torch.manual_seed(0							)
       lowercase__:   Union[str, Any]					=      pipe(
           prompt=_UpperCAmelCase    ,			generator=_UpperCAmelCase    ,			guidance_scale=7.5    ,			num_inference_steps=2    ,			output_type='''numpy'''							).images

       with tempfile.TemporaryDirectory() as tmpdirname:
          pipe.save_pretrained(_UpperCAmelCase							)
          lowercase__:   Optional[Any]					=      VersatileDiffusionTextToImagePipeline.from_pretrained(_UpperCAmelCase							)
       pipe.to(_UpperCAmelCase							)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase							)

       lowercase__:   Optional[int]					=      generator.manual_seed(0							)
       lowercase__:   List[str]					=      pipe(
           prompt=_UpperCAmelCase    ,			generator=_UpperCAmelCase    ,			guidance_scale=7.5    ,			num_inference_steps=2    ,			output_type='''numpy'''							).images

       assert np.abs(image - new_image							).sum() < 1e-5, "Models don't have the same forward pass"






    def _snake_case			(   self							):
       lowercase__:   Dict					=      VersatileDiffusionTextToImagePipeline.from_pretrained(
           '''shi-labs/versatile-diffusion'''    ,			torch_dtype=torch.floataa							)
       pipe.to(_UpperCAmelCase							)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase							)

       lowercase__:   Tuple					=      '''A painting of a squirrel eating a burger '''
       lowercase__:   Optional[Any]					=      torch.manual_seed(0							)
       lowercase__:   Tuple					=      pipe(
           prompt=_UpperCAmelCase    ,			generator=_UpperCAmelCase    ,			guidance_scale=7.5    ,			num_inference_steps=50    ,			output_type='''numpy'''							).images

       lowercase__:   Union[str, Any]					=      image[0, 253:256, 253:256, -1]

       assert image.shape == (1, 512, 512, 3)
       lowercase__:   Any					=      np.array([0.3_367, 0.3_169, 0.2_656, 0.3_870, 0.4_790, 0.3_796, 0.4_009, 0.4_878, 0.4_778]							)

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