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

      def __init__(						self		,       __lowerCAmelCase		,       __lowerCAmelCase=13		,       __lowerCAmelCase=7		,       __lowerCAmelCase=False		,       __lowerCAmelCase=True		,       __lowerCAmelCase=False		,       __lowerCAmelCase=False		,       __lowerCAmelCase=19		,       __lowerCAmelCase=32		,       __lowerCAmelCase=5		,       __lowerCAmelCase=4		,       __lowerCAmelCase=37		,       __lowerCAmelCase="gelu"		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=512		,       __lowerCAmelCase=16		,       __lowerCAmelCase=2		,       __lowerCAmelCase=0.0_2		,       __lowerCAmelCase=3		,       __lowerCAmelCase=4		,       __lowerCAmelCase=None		,       ):







            """simple docstring"""



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

      def 				A__       (						self						):







            """simple docstring"""



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

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

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

            lowercase														=	self.get_config()

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

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	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=__lowerCAmelCase		,       esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}		,       )
            return config

      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            lowercase														=	EsmForProteinFolding(config=__lowerCAmelCase						).float()
            model.to(__lowerCAmelCase						)
            model.eval()
            lowercase														=	model(__lowerCAmelCase		,       attention_mask=__lowerCAmelCase						)
            lowercase														=	model(__lowerCAmelCase						)
            lowercase														=	model(__lowerCAmelCase						)

            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 				A__       (						self						):







            """simple docstring"""



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






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

      snake_case__			:   int							  =	(EsmForProteinFolding,) if is_torch_available() else ()
      snake_case__			:   List[Any]							  =	()
      snake_case__			:   str							  =	{} if is_torch_available() else {}
      snake_case__			:   List[Any]							  =	False

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	EsmFoldModelTester(self						)
            lowercase														=	ConfigTester(self		,       config_class=__lowerCAmelCase		,       hidden_size=37						)

      def 				A__       (						self						):







            """simple docstring"""



            self.config_tester.run_common_tests()

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__lowerCAmelCase						)

      @unittest.skip("""Does not support attention outputs"""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""Esm does not support embedding resizing"""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""Esm does not support embedding resizing"""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support passing input embeds!"""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support head pruning."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support head pruning."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support head pruning."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support head pruning."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support head pruning."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not output hidden states in the normal way."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMfold does not output hidden states in the normal way."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold only has one output format."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold does not support input chunking."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold doesn't support torchscript compilation."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold doesn't support torchscript compilation."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold doesn't support torchscript compilation."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass

      @unittest.skip("""ESMFold doesn't support data parallel."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass







      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""						)
      def 				A__       (						self						):







            """simple docstring"""



            pass






@require_torch
class  _A						(       lowerCAmelCase	):

      @slow
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""						).float()
            model.eval()
            lowercase														=	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
            lowercase														=	model(__lowerCAmelCase						)["""positions"""]
            lowercase														=	torch.tensor([2.5_8_2_8, 0.7_9_9_3, -1_0.9_3_3_4]		,       dtype=torch.floataa						)
            self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]		,       __lowerCAmelCase		,       atol=1E-4						)						)




"""simple docstring"""






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np





class  _A						:

      def 				A__       (						self						):







            """simple docstring"""



            torch.manual_seed(0						)
            lowercase														=	TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""						)

            torch.manual_seed(0						)
            lowercase														=	AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""						)

            torch.manual_seed(0						)
            lowercase														=	UNetaDConditionModel(
                sample_size=32		,       layers_per_block=1		,       block_out_channels=[32, 64]		,       down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]		,       mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""		,       up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]		,       in_channels=3		,       out_channels=6		,       cross_attention_dim=32		,       encoder_hid_dim=32		,       attention_head_dim=8		,       addition_embed_type="""text"""		,       addition_embed_type_num_heads=2		,       cross_attention_norm="""group_norm"""		,       resnet_time_scale_shift="""scale_shift"""		,       act_fn="""gelu"""		,       )
            unet.set_attn_processor(AttnAddedKVProcessor()						)  # For reproducibility tests

            torch.manual_seed(0						)
            lowercase														=	DDPMScheduler(
                num_train_timesteps=1000		,       beta_schedule="""squaredcos_cap_v2"""		,       beta_start=0.0_0_0_1		,       beta_end=0.0_2		,       thresholding=__lowerCAmelCase		,       dynamic_thresholding_ratio=0.9_5		,       sample_max_value=1.0		,       prediction_type="""epsilon"""		,       variance_type="""learned_range"""		,       )

            torch.manual_seed(0						)
            lowercase														=	IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }

      def 				A__       (						self						):







            """simple docstring"""



            torch.manual_seed(0						)
            lowercase														=	TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""						)

            torch.manual_seed(0						)
            lowercase														=	AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5"""						)

            torch.manual_seed(0						)
            lowercase														=	UNetaDConditionModel(
                sample_size=32		,       layers_per_block=[1, 2]		,       block_out_channels=[32, 64]		,       down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]		,       mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""		,       up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]		,       in_channels=6		,       out_channels=6		,       cross_attention_dim=32		,       encoder_hid_dim=32		,       attention_head_dim=8		,       addition_embed_type="""text"""		,       addition_embed_type_num_heads=2		,       cross_attention_norm="""group_norm"""		,       resnet_time_scale_shift="""scale_shift"""		,       act_fn="""gelu"""		,       class_embed_type="""timestep"""		,       mid_block_scale_factor=1.4_1_4		,       time_embedding_act_fn="""gelu"""		,       time_embedding_dim=32		,       )
            unet.set_attn_processor(AttnAddedKVProcessor()						)  # For reproducibility tests

            torch.manual_seed(0						)
            lowercase														=	DDPMScheduler(
                num_train_timesteps=1000		,       beta_schedule="""squaredcos_cap_v2"""		,       beta_start=0.0_0_0_1		,       beta_end=0.0_2		,       thresholding=__lowerCAmelCase		,       dynamic_thresholding_ratio=0.9_5		,       sample_max_value=1.0		,       prediction_type="""epsilon"""		,       variance_type="""learned_range"""		,       )

            torch.manual_seed(0						)
            lowercase														=	DDPMScheduler(
                num_train_timesteps=1000		,       beta_schedule="""squaredcos_cap_v2"""		,       beta_start=0.0_0_0_1		,       beta_end=0.0_2		,       )

            torch.manual_seed(0						)
            lowercase														=	IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "image_noising_scheduler": image_noising_scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	self.get_dummy_components()
            lowercase														=	self.pipeline_class(**__lowerCAmelCase						)
            pipe.to(__lowerCAmelCase						)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase						)

            lowercase														=	self.get_dummy_inputs(__lowerCAmelCase						)

            lowercase														=	inputs["""prompt"""]
            lowercase														=	inputs["""generator"""]
            lowercase														=	inputs["""num_inference_steps"""]
            lowercase														=	inputs["""output_type"""]

            if "image" in inputs:
                  lowercase														=	inputs["""image"""]
            else:
                  lowercase														=	None

            if "mask_image" in inputs:
                  lowercase														=	inputs["""mask_image"""]
            else:
                  lowercase														=	None

            if "original_image" in inputs:
                  lowercase														=	inputs["""original_image"""]
            else:
                  lowercase														=	None

            lowercase							,							lowercase														=	pipe.encode_prompt(__lowerCAmelCase						)

            # inputs with prompt converted to embeddings
            lowercase														=	{
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowercase														=	image

            if mask_image is not None:
                  lowercase														=	mask_image

            if original_image is not None:
                  lowercase														=	original_image

            # set all optional components to None
            for optional_component in pipe._optional_components:
                  setattr(__lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						)

            lowercase														=	pipe(**__lowerCAmelCase						)[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(__lowerCAmelCase						)
                  lowercase														=	self.pipeline_class.from_pretrained(__lowerCAmelCase						)
                  pipe_loaded.to(__lowerCAmelCase						)
                  pipe_loaded.set_progress_bar_config(disable=__lowerCAmelCase						)

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()						)  # For reproducibility tests

            for optional_component in pipe._optional_components:
                  self.assertTrue(
                      getattr(__lowerCAmelCase		,       __lowerCAmelCase						) is None		,       f'`{optional_component}` did not stay set to None after loading.'		,       )

            lowercase														=	self.get_dummy_inputs(__lowerCAmelCase						)

            lowercase														=	inputs["""generator"""]
            lowercase														=	inputs["""num_inference_steps"""]
            lowercase														=	inputs["""output_type"""]

            # inputs with prompt converted to embeddings
            lowercase														=	{
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowercase														=	image

            if mask_image is not None:
                  lowercase														=	mask_image

            if original_image is not None:
                  lowercase														=	original_image

            lowercase														=	pipe_loaded(**__lowerCAmelCase						)[0]

            lowercase														=	np.abs(to_np(__lowerCAmelCase						) - to_np(__lowerCAmelCase						)						).max()
            self.assertLess(__lowerCAmelCase		,       1E-4						)







      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	self.get_dummy_components()
            lowercase														=	self.pipeline_class(**__lowerCAmelCase						)
            pipe.to(__lowerCAmelCase						)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase						)

            lowercase														=	self.get_dummy_inputs(__lowerCAmelCase						)
            lowercase														=	pipe(**__lowerCAmelCase						)[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(__lowerCAmelCase						)
                  lowercase														=	self.pipeline_class.from_pretrained(__lowerCAmelCase						)
                  pipe_loaded.to(__lowerCAmelCase						)
                  pipe_loaded.set_progress_bar_config(disable=__lowerCAmelCase						)

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()						)  # For reproducibility tests

            lowercase														=	self.get_dummy_inputs(__lowerCAmelCase						)
            lowercase														=	pipe_loaded(**__lowerCAmelCase						)[0]

            lowercase														=	np.abs(to_np(__lowerCAmelCase						) - to_np(__lowerCAmelCase						)						).max()
            self.assertLess(__lowerCAmelCase		,       1E-4						)



