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 			__magic_name__	:





   """simple docstring"""




   def __init__(						self       :	Union[str, Any]		,	_lowercase       :	int		,	):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			parent
    _UpperCamelCase:  Optional[int]    						=			13
    _UpperCamelCase:  str    						=			7
    _UpperCamelCase:  Dict    						=			True
    _UpperCamelCase:  List[str]    						=			True
    _UpperCamelCase:  Optional[int]    						=			True
    _UpperCamelCase:  Dict    						=			99
    _UpperCamelCase:  List[str]    						=			32
    _UpperCamelCase:  Tuple    						=			2
    _UpperCamelCase:  Any    						=			4
    _UpperCamelCase:  Optional[int]    						=			37
    _UpperCamelCase:  int    						=			'''gelu'''
    _UpperCamelCase:  Union[str, Any]    						=			0.1
    _UpperCamelCase:  Optional[int]    						=			0.1
    _UpperCamelCase:  Union[str, Any]    						=			512
    _UpperCamelCase:  Any    						=			16
    _UpperCamelCase:  Union[str, Any]    						=			2
    _UpperCamelCase:  Optional[Any]    						=			0.02
    _UpperCamelCase:  Union[str, Any]    						=			3
    _UpperCamelCase:  Optional[Any]    						=			4
    _UpperCamelCase:  List[Any]    						=			None



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  List[str]    						=			None
    if self.use_input_mask:
     _UpperCamelCase:  Dict    						=			random_attention_mask([self.batch_size, self.seq_length]				)

    _UpperCamelCase:  int    						=			None
    _UpperCamelCase:  int    						=			None
    _UpperCamelCase:  List[str]    						=			None
    if self.use_labels:
     _UpperCamelCase:  int    						=			ids_tensor([self.batch_size]		,	self.type_sequence_label_size				)
     _UpperCamelCase:  Optional[Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.num_labels				)
     _UpperCamelCase:  Optional[int]    						=			ids_tensor([self.batch_size]		,	self.num_choices				)

    _UpperCamelCase:  Dict    						=			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 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    (
        (
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				
    ):  Optional[int]    						=			self.prepare_config_and_inputs()

    _UpperCamelCase:  List[Any]    						=			True
    _UpperCamelCase:  Optional[Any]    						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
    _UpperCamelCase:  Tuple    						=			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 	lowerCAmelCase    (						self       :	Optional[Any]		,	_lowercase       :	List[Any]		,	_lowercase       :	Union[str, Any]		,	_lowercase       :	List[Any]		,	_lowercase       :	Any		,	_lowercase       :	str		,	_lowercase       :	List[str]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			TFEsmModel(config=_lowercase				)
    _UpperCamelCase:  int    						=			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
    _UpperCamelCase:  Optional[int]    						=			model(_lowercase				)

    _UpperCamelCase:  Union[str, Any]    						=			[input_ids, input_mask]
    _UpperCamelCase:  Dict    						=			model(_lowercase				)

    _UpperCamelCase:  Tuple    						=			model(_lowercase				)

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



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	Dict		,	_lowercase       :	int		,	_lowercase       :	Any		,	_lowercase       :	Optional[int]		,	_lowercase       :	List[Any]		,	_lowercase       :	Tuple		,	_lowercase       :	List[Any]		,	_lowercase       :	Dict		,	):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			True

    _UpperCamelCase:  int    						=			TFEsmModel(config=_lowercase				)
    _UpperCamelCase:  Union[str, Any]    						=			{
        '''input_ids''': input_ids,
        '''attention_mask''': input_mask,
        '''encoder_hidden_states''': encoder_hidden_states,
        '''encoder_attention_mask''': encoder_attention_mask,
    }
    _UpperCamelCase:  int    						=			model(_lowercase				)

    _UpperCamelCase:  Optional[int]    						=			[input_ids, input_mask]
    _UpperCamelCase:  int    						=			model(_lowercase		,	encoder_hidden_states=_lowercase				)

    # Also check the case where encoder outputs are not passed
    _UpperCamelCase:  Tuple    						=			model(_lowercase		,	attention_mask=_lowercase				)

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



   def 	lowerCAmelCase    (						self       :	Any		,	_lowercase       :	str		,	_lowercase       :	str		,	_lowercase       :	Any		,	_lowercase       :	List[str]		,	_lowercase       :	Tuple		,	_lowercase       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			TFEsmForMaskedLM(config=_lowercase				)
    _UpperCamelCase:  str    						=			model([input_ids, input_mask]				)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.vocab_size)				)



   def 	lowerCAmelCase    (						self       :	Union[str, Any]		,	_lowercase       :	Optional[Any]		,	_lowercase       :	Dict		,	_lowercase       :	Union[str, Any]		,	_lowercase       :	List[str]		,	_lowercase       :	Optional[int]		,	_lowercase       :	List[str]				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			self.num_labels
    _UpperCamelCase:  Optional[int]    						=			TFEsmForTokenClassification(config=_lowercase				)
    _UpperCamelCase:  Tuple    						=			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
    _UpperCamelCase:  Optional[Any]    						=			model(_lowercase				)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.num_labels)				)



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			self.prepare_config_and_inputs()
    (
        (
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				(
        _UpperCamelCase
    )						,				
    ):  List[Any]    						=			config_and_inputs
    _UpperCamelCase:  List[Any]    						=			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict




@require_tf
class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     int	      = (
       (
           TFEsmModel,
           TFEsmForMaskedLM,
           TFEsmForSequenceClassification,
           TFEsmForTokenClassification,
       )
       if is_tf_available()
       else ()
   )
   lowerCAmelCase   :     str	      = (
       {
           '''feature-extraction''': TFEsmModel,
           '''fill-mask''': TFEsmForMaskedLM,
           '''text-classification''': TFEsmForSequenceClassification,
           '''token-classification''': TFEsmForTokenClassification,
           '''zero-shot''': TFEsmForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   lowerCAmelCase   :     int	      = False
   lowerCAmelCase   :     Dict	      = False



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			TFEsmModelTester(self				)
    _UpperCamelCase:  int    						=			ConfigTester(self		,	config_class=_lowercase		,	hidden_size=37				)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    self.config_tester.run_common_tests()



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



    """simple docstring"""



    _UpperCamelCase:  str    						=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*_lowercase				)



   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*_lowercase				)



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*_lowercase				)



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*_lowercase				)



   @slow
   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _UpperCamelCase:  int    						=			TFEsmModel.from_pretrained(_lowercase				)
     self.assertIsNotNone(_lowercase				)



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



    """simple docstring"""



    pass



   @unittest.skip('''Protein models do not support embedding resizing.'''				)
   def 	lowerCAmelCase    (						self       :	List[str]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    _UpperCamelCase						,				_UpperCamelCase:  List[str]    						=			self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
     _UpperCamelCase:  Optional[Any]    						=			model_class(_lowercase				)
     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
      _UpperCamelCase:  str    						=			model.get_bias()
      assert isinstance(_lowercase		,	_lowercase				)
      for k, v in name.items():
       assert isinstance(_lowercase		,	tf.Variable				)
     else:
      _UpperCamelCase:  List[Any]    						=			model.get_output_embeddings()
      assert x is None
      _UpperCamelCase:  int    						=			model.get_bias()
      assert name is None




@require_tf
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   @slow
   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''				)

    _UpperCamelCase:  List[str]    						=			tf.constant([[0, 1, 2, 3, 4, 5]]				)
    _UpperCamelCase:  Union[str, Any]    						=			model(_lowercase				)[0]
    _UpperCamelCase:  Optional[Any]    						=			[1, 6, 33]
    self.assertEqual(list(output.numpy().shape				)		,	_lowercase				)
    # compare the actual values for a slice.
    _UpperCamelCase:  Dict    						=			tf.constant(
        [
            [
                [8.921518, -10.589814, -6.4671307],
                [-6.3967156, -13.911377, -1.1211915],
                [-7.781247, -13.951557, -3.740592],
            ]
        ]				)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,	expected_slice.numpy()		,	atol=1E-2				)				)



   @slow
   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  List[Any]    						=			TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''				)

    _UpperCamelCase:  Union[str, Any]    						=			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]				)
    _UpperCamelCase:  Dict    						=			model(_lowercase				)[0]
    # compare the actual values for a slice.
    _UpperCamelCase:  Union[str, Any]    						=			tf.constant(
        [
            [
                [0.14443092, 0.54125327, 0.3247739],
                [0.30340484, 0.00526676, 0.31077722],
                [0.32278043, -0.24987096, 0.3414628],
            ]
        ]				)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,	expected_slice.numpy()		,	atol=1E-4				)				)
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     List[str]	      = StableDiffusionSAGPipeline
   lowerCAmelCase   :     List[Any]	      = TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase   :     List[str]	      = TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase   :     str	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = False



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			UNetaDConditionModel(
        block_out_channels=(32, 64)		,	layers_per_block=2		,	sample_size=32		,	in_channels=4		,	out_channels=4		,	down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')		,	up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')		,	cross_attention_dim=32		,	)
    _UpperCamelCase:  List[Any]    						=			DDIMScheduler(
        beta_start=0.00085		,	beta_end=0.012		,	beta_schedule='''scaled_linear'''		,	clip_sample=_lowercase		,	set_alpha_to_one=_lowercase		,	)
    torch.manual_seed(0				)
    _UpperCamelCase:  Optional[int]    						=			AutoencoderKL(
        block_out_channels=[32, 64]		,	in_channels=3		,	out_channels=3		,	down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,	up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,	latent_channels=4		,	)
    torch.manual_seed(0				)
    _UpperCamelCase:  int    						=			CLIPTextConfig(
        bos_token_id=0		,	eos_token_id=2		,	hidden_size=32		,	intermediate_size=37		,	layer_norm_eps=1E-05		,	num_attention_heads=4		,	num_hidden_layers=5		,	pad_token_id=1		,	vocab_size=1_000		,	)
    _UpperCamelCase:  Any    						=			CLIPTextModel(_lowercase				)
    _UpperCamelCase:  List[str]    						=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

    _UpperCamelCase:  int    						=			{
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''vae''': vae,
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''safety_checker''': None,
        '''feature_extractor''': None,
    }
    return components



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	List[Any]		,	_lowercase       :	int=0				):



    """simple docstring"""



    if str(_lowercase				).startswith('''mps'''				):
     _UpperCamelCase:  Tuple    						=			torch.manual_seed(_lowercase				)
    else:
     _UpperCamelCase:  List[str]    						=			torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
    _UpperCamelCase:  Dict    						=			{
        '''prompt''': '''.''',
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 1.0,
        '''sag_scale''': 1.0,
        '''output_type''': '''numpy''',
    }
    return inputs



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    super().test_inference_batch_single_identical(expected_max_diff=3E-3				)




@slow
@require_torch_gpu
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




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



    """simple docstring"""



    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''				)
    _UpperCamelCase:  Any    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Tuple    						=			'''.'''
    _UpperCamelCase:  Tuple    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Union[str, Any]    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Optional[Any]    						=			output.images

    _UpperCamelCase:  Tuple    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  Any    						=			np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 5E-2



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  List[Any]    						=			'''.'''
    _UpperCamelCase:  str    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Union[str, Any]    						=			output.images

    _UpperCamelCase:  List[str]    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  int    						=			np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 5E-2



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Any    						=			'''.'''
    _UpperCamelCase:  List[Any]    						=			torch.manual_seed(0				)
    _UpperCamelCase:  str    						=			sag_pipe(
        [prompt]		,	width=768		,	height=512		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''		,	)

    _UpperCamelCase:  str    						=			output.images

    assert image.shape == (1, 512, 768, 3)