






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 				__UpperCAmelCase		(unittest.TestCase   ):







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									super().tearDown()
									gc.collect()







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/sd2-inpaint/init_image.png"""					)
									_SCREAMING_SNAKE_CASE     				=  load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""					)

									_SCREAMING_SNAKE_CASE     				=  """xvjiarui/stable-diffusion-2-inpainting"""
									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  FlaxStableDiffusionInpaintPipeline.from_pretrained(UpperCAmelCase_				,     safety_checker=UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """Face of a yellow cat, high resolution, sitting on a park bench"""

									_SCREAMING_SNAKE_CASE     				=  jax.random.PRNGKey(0					)
									_SCREAMING_SNAKE_CASE     				=  50

									_SCREAMING_SNAKE_CASE     				=  jax.device_count()
									_SCREAMING_SNAKE_CASE     				=  num_samples * [prompt]
									_SCREAMING_SNAKE_CASE     				=  num_samples * [init_image]
									_SCREAMING_SNAKE_CASE     				=  num_samples * [mask_image]
									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  pipeline.prepare_inputs(UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_					)

									# shard inputs and rng
									_SCREAMING_SNAKE_CASE     				=  replicate(UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  jax.random.split(UpperCAmelCase_				,     jax.device_count()					)
									_SCREAMING_SNAKE_CASE     				=  shard(UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  shard(UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  shard(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  pipeline(
									    UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     jit=UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  output.images.reshape(UpperCAmelCase_				,     512				,     512				,     3					)

									_SCREAMING_SNAKE_CASE     				=  images[0, 253:256, 253:256, -1]

									_SCREAMING_SNAKE_CASE     				=  jnp.asarray(jax.device_get(image_slice.flatten()					)					)
									_SCREAMING_SNAKE_CASE     				=  jnp.array(
									    [0.3_61_13_07, 0.37_64_97_36, 0.3_75_74_08, 0.38_21_39_53, 0.39_29_51_67, 0.3_84_16_31, 0.41_55_49_78, 0.4_13_74_75, 0.4_21_70_84]					)
									print(F'output_slice: {output_slice}'					)
									assert jnp.abs(output_slice - expected_slice					).max() < 1E-2









from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers.models.esm.modeling_tf_esm import (
		    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFEsmForMaskedLM,
		    TFEsmForSequenceClassification,
		    TFEsmForTokenClassification,
		    TFEsmModel,
		)



class 				__UpperCAmelCase		:







			def __init__(     self:  Tuple				,     UpperCAmelCase_:  Tuple				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  13
									_SCREAMING_SNAKE_CASE     				=  7
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  99
									_SCREAMING_SNAKE_CASE     				=  32
									_SCREAMING_SNAKE_CASE     				=  2
									_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







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




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

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

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

									_SCREAMING_SNAKE_CASE     				=  EsmConfig(
									    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     pad_token_id=1				,     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				,     )

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







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  self.prepare_config_and_inputs()

									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
									_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     vocab_size=2					)

									return (
									    config,
									    input_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    encoder_hidden_states,
									    encoder_attention_mask,
									)







			def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModel(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  [input_ids, input_mask]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

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







			def 							UpperCamelCase				(     self:  Optional[Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  List[str]				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  True

									_SCREAMING_SNAKE_CASE     				=  TFEsmModel(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """encoder_hidden_states""": encoder_hidden_states,
									    """encoder_attention_mask""": encoder_attention_mask,
									}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  [input_ids, input_mask]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     encoder_hidden_states=UpperCAmelCase_					)

									# Also check the case where encoder outputs are not passed
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_					)

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







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmForMaskedLM(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model([input_ids, input_mask]					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size)					)







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.num_labels
									_SCREAMING_SNAKE_CASE     				=  TFEsmForTokenClassification(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels)					)







			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




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



@require_tf
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				List[Any]									=							(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			__snake_case    :				Tuple									=							(
			    {
			        "feature-extraction": TFEsmModel,
			        "fill-mask": TFEsmForMaskedLM,
			        "text-classification": TFEsmForSequenceClassification,
			        "token-classification": TFEsmForTokenClassification,
			        "zero-shot": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			__snake_case    :				List[str]									=							False
			__snake_case    :				Union[str, Any]									=							False







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_decoder()
									self.model_tester.create_and_check_model_as_decoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFEsmModel.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)







			@unittest.skip("""Protein models do not support embedding resizing."""					)
			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									pass







			@unittest.skip("""Protein models do not support embedding resizing."""					)
			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									pass







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															assert isinstance(model.get_input_embeddings()				,     tf.keras.layers.Layer					)
															if model_class is TFEsmForMaskedLM:
																					# Output embedding test differs from the main test because they're a matrix, not a layer
																					_SCREAMING_SNAKE_CASE     				=  model.get_bias()
																					assert isinstance(UpperCAmelCase_				,     UpperCAmelCase_					)
																					for k, v in name.items():
																											assert isinstance(UpperCAmelCase_				,     tf.Variable					)
															else:
																					_SCREAMING_SNAKE_CASE     				=  model.get_output_embeddings()
																					assert x is None
																					_SCREAMING_SNAKE_CASE     				=  model.get_bias()
																					assert name is None



@require_tf
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant([[0, 1, 2, 3, 4, 5]]					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]
									_SCREAMING_SNAKE_CASE     				=  [1, 6, 33]
									self.assertEqual(list(output.numpy().shape					)				,     UpperCAmelCase_					)
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
									            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
									            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,     expected_slice.numpy()				,     atol=1E-2					)					)







			@slow
			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
									            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
									            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,     expected_slice.numpy()				,     atol=1E-4					)					)

