






"""simple docstring"""







from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
 from transformers import (
     TFAutoModel,
     TFAutoModelForCausalLM,
     TFAutoModelForMaskedLM,
     TFAutoModelForPreTraining,
     TFAutoModelForQuestionAnswering,
     TFAutoModelForSeqaSeqLM,
     TFAutoModelForSequenceClassification,
     TFAutoModelForTableQuestionAnswering,
     TFAutoModelForTokenClassification,
     TFAutoModelWithLMHead,
     TFBertForMaskedLM,
     TFBertForPreTraining,
     TFBertForQuestionAnswering,
     TFBertForSequenceClassification,
     TFBertModel,
     TFFunnelBaseModel,
     TFFunnelModel,
     TFGPTaLMHeadModel,
     TFRobertaForMaskedLM,
     TFTaForConditionalGeneration,
     TFTapasForQuestionAnswering,
 )
 from transformers.models.auto.modeling_tf_auto import (
     TF_MODEL_FOR_CAUSAL_LM_MAPPING,
     TF_MODEL_FOR_MASKED_LM_MAPPING,
     TF_MODEL_FOR_PRETRAINING_MAPPING,
     TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
     TF_MODEL_MAPPING,
 )
 from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST





class   UpperCamelCase_ (       __lowercase  ):
  _A       :					Any				=	'new-model'


if is_tf_available():





 class   UpperCamelCase_ (       __lowercase  ):
   _A       :					str				=	NewModelConfig





@require_tf
class   UpperCamelCase_ (       unittest.TestCase  ):



  @slow
  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			"""bert-base-cased"""
         UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
         self.assertIsNotNone(_A				)
         self.assertIsInstance(_A ,		_A				)

         UpperCAmelCase	    =			TFAutoModel.from_pretrained(_A				)
         self.assertIsNotNone(_A				)
         self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			"""bert-base-cased"""
         UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
         self.assertIsNotNone(_A				)
         self.assertIsInstance(_A ,		_A				)

         UpperCAmelCase	    =			TFAutoModelForPreTraining.from_pretrained(_A				)
         self.assertIsNotNone(_A				)
         self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			List[Any]:



         """simple docstring"""





         for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForCausalLM.from_pretrained(_A				)
                UpperCAmelCase    ,       UpperCAmelCase	    =			TFAutoModelForCausalLM.from_pretrained(_A ,		output_loading_info=_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Optional[Any]:



         """simple docstring"""





         for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelWithLMHead.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Union[str, Any]:



         """simple docstring"""





         for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForMaskedLM.from_pretrained(_A				)
                UpperCAmelCase    ,       UpperCAmelCase	    =			TFAutoModelForMaskedLM.from_pretrained(_A ,		output_loading_info=_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Optional[Any]:



         """simple docstring"""





         for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForSeqaSeqLM.from_pretrained(_A				)
                UpperCAmelCase    ,       UpperCAmelCase	    =			TFAutoModelForSeqaSeqLM.from_pretrained(_A ,		output_loading_info=_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Dict:



         """simple docstring"""





         for model_name in ["bert-base-uncased"]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForSequenceClassification.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  def 		UpperCamelCase_				(      self				)			->			Any:



         """simple docstring"""





         for model_name in ["bert-base-uncased"]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForQuestionAnswering.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  @slow
  @require_tensorflow_probability
  def 		UpperCamelCase_				(      self				)			->			List[str]:



         """simple docstring"""





         for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
                UpperCAmelCase	    =			AutoConfig.from_pretrained(_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)

                UpperCAmelCase	    =			TFAutoModelForTableQuestionAnswering.from_pretrained(_A				)
                UpperCAmelCase    ,       UpperCAmelCase	    =			TFAutoModelForTableQuestionAnswering.from_pretrained(
                    _A ,		output_loading_info=_A				)
                self.assertIsNotNone(_A				)
                self.assertIsInstance(_A ,		_A				)



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			TFAutoModelWithLMHead.from_pretrained(_A				)
         self.assertIsInstance(_A ,		_A				)
         self.assertEqual(model.num_parameters() ,		1_44_10				)
         self.assertEqual(model.num_parameters(only_trainable=_A				) ,		1_44_10				)



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



         """simple docstring"""





         UpperCAmelCase	    =			TFAutoModelWithLMHead.from_pretrained(_A				)
         self.assertIsInstance(_A ,		_A				)
         self.assertEqual(model.num_parameters() ,		1_44_10				)
         self.assertEqual(model.num_parameters(only_trainable=_A				) ,		1_44_10				)



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



         """simple docstring"""





         UpperCAmelCase	    =			TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""				)
         self.assertIsInstance(_A ,		_A				)

         UpperCAmelCase	    =			copy.deepcopy(model.config				)
         UpperCAmelCase	    =			["""FunnelBaseModel"""]
         UpperCAmelCase	    =			TFAutoModel.from_config(_A				)
         self.assertIsInstance(_A ,		_A				)

         with tempfile.TemporaryDirectory() as tmp_dir:
                model.save_pretrained(_A				)
                UpperCAmelCase	    =			TFAutoModel.from_pretrained(_A				)
                self.assertIsInstance(_A ,		_A				)



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



         """simple docstring"""





         try:
                AutoConfig.register("""new-model""" ,		_A				)

                UpperCAmelCase	    =			[
                    TFAutoModel,
                    TFAutoModelForCausalLM,
                    TFAutoModelForMaskedLM,
                    TFAutoModelForPreTraining,
                    TFAutoModelForQuestionAnswering,
                    TFAutoModelForSequenceClassification,
                    TFAutoModelForTokenClassification,
                ]

                for auto_class in auto_classes:
                       with self.subTest(auto_class.__name__				):
                              # Wrong config class will raise an error
                              with self.assertRaises(_A				):
                                     auto_class.register(_A ,		_A				)
                              auto_class.register(_A ,		_A				)
                              # Trying to register something existing in the Transformers library will raise an error
                              with self.assertRaises(_A				):
                                     auto_class.register(_A ,		_A				)

                              # Now that the config is registered, it can be used as any other config with the auto-API
                              UpperCAmelCase	    =			BertModelTester(self				).get_config()
                              UpperCAmelCase	    =			NewModelConfig(**tiny_config.to_dict()				)
                              UpperCAmelCase	    =			auto_class.from_config(_A				)
                              self.assertIsInstance(_A ,		_A				)

                              with tempfile.TemporaryDirectory() as tmp_dir:
                                     model.save_pretrained(_A				)
                                     UpperCAmelCase	    =			auto_class.from_pretrained(_A				)
                                     self.assertIsInstance(_A ,		_A				)

         finally:
                if "new-model" in CONFIG_MAPPING._extra_content:
                       del CONFIG_MAPPING._extra_content["new-model"]
                for mapping in (
                    TF_MODEL_MAPPING,
                    TF_MODEL_FOR_PRETRAINING_MAPPING,
                    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                    TF_MODEL_FOR_MASKED_LM_MAPPING,
                ):
                       if NewModelConfig in mapping._extra_content:
                              del mapping._extra_content[NewModelConfig]



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         with self.assertRaisesRegex(
             _A ,		"""bert-base is not a local folder and is not a valid model identifier"""				):
                UpperCAmelCase	    =			TFAutoModel.from_pretrained("""bert-base"""				)



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         with self.assertRaisesRegex(
             _A ,		R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""				):
                UpperCAmelCase	    =			TFAutoModel.from_pretrained(_A ,		revision="""aaaaaa"""				)



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



         """simple docstring"""





         with self.assertRaisesRegex(
             _A ,		"""hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin""" ,		):
                UpperCAmelCase	    =			TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""				)



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



         """simple docstring"""





         with self.assertRaisesRegex(_A ,		"""Use `from_pt=True` to load this model"""				):
                UpperCAmelCase	    =			TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""				)



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



         """simple docstring"""





         UpperCAmelCase	    =			TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""				)
         with RequestCounter() as counter:
                UpperCAmelCase	    =			TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""				)
                self.assertEqual(counter.get_request_count ,		0				)
                self.assertEqual(counter.head_request_count ,		1				)
                self.assertEqual(counter.other_request_count ,		0				)

         # With a sharded checkpoint
         UpperCAmelCase	    =			TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""				)
         with RequestCounter() as counter:
                UpperCAmelCase	    =			TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""				)
                self.assertEqual(counter.get_request_count ,		0				)
                self.assertEqual(counter.head_request_count ,		1				)
                self.assertEqual(counter.other_request_count ,		0				)








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



  def 		UpperCamelCase_				(      self				)			->			Any:



         """simple docstring"""





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

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

         torch.manual_seed(0				)
         UpperCAmelCase	    =			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				)
         UpperCAmelCase	    =			DDPMScheduler(
             num_train_timesteps=10_00 ,		beta_schedule="""squaredcos_cap_v2""" ,		beta_start=0.0_001 ,		beta_end=0.02 ,		thresholding=snake_case__ ,		dynamic_thresholding_ratio=0.95 ,		sample_max_value=1.0 ,		prediction_type="""epsilon""" ,		variance_type="""learned_range""" ,		)

         torch.manual_seed(0				)
         UpperCAmelCase	    =			IFWatermarker()

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



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





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

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

         torch.manual_seed(0				)
         UpperCAmelCase	    =			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.414 ,		time_embedding_act_fn="""gelu""" ,		time_embedding_dim=32 ,		)
         unet.set_attn_processor(AttnAddedKVProcessor()				)  # For reproducibility tests

         torch.manual_seed(0				)
         UpperCAmelCase	    =			DDPMScheduler(
             num_train_timesteps=10_00 ,		beta_schedule="""squaredcos_cap_v2""" ,		beta_start=0.0_001 ,		beta_end=0.02 ,		thresholding=snake_case__ ,		dynamic_thresholding_ratio=0.95 ,		sample_max_value=1.0 ,		prediction_type="""epsilon""" ,		variance_type="""learned_range""" ,		)

         torch.manual_seed(0				)
         UpperCAmelCase	    =			DDPMScheduler(
             num_train_timesteps=10_00 ,		beta_schedule="""squaredcos_cap_v2""" ,		beta_start=0.0_001 ,		beta_end=0.02 ,		)

         torch.manual_seed(0				)
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			self.get_dummy_components()
         UpperCAmelCase	    =			self.pipeline_class(**snake_case__				)
         pipe.to(snake_case__				)
         pipe.set_progress_bar_config(disable=snake_case__				)

         UpperCAmelCase	    =			self.get_dummy_inputs(snake_case__				)

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

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

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

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

         UpperCAmelCase    ,       UpperCAmelCase	    =			pipe.encode_prompt(snake_case__				)

         # inputs with prompt converted to embeddings
         UpperCAmelCase	    =			{
             """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:
                UpperCAmelCase	    =			image

         if mask_image is not None:
                UpperCAmelCase	    =			mask_image

         if original_image is not None:
                UpperCAmelCase	    =			original_image

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

         UpperCAmelCase	    =			pipe(**snake_case__				)[0]

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

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

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

         UpperCAmelCase	    =			self.get_dummy_inputs(snake_case__				)

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

         # inputs with prompt converted to embeddings
         UpperCAmelCase	    =			{
             """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:
                UpperCAmelCase	    =			image

         if mask_image is not None:
                UpperCAmelCase	    =			mask_image

         if original_image is not None:
                UpperCAmelCase	    =			original_image

         UpperCAmelCase	    =			pipe_loaded(**snake_case__				)[0]

         UpperCAmelCase	    =			np.abs(to_np(snake_case__				) - to_np(snake_case__				)				).max()
         self.assertLess(snake_case__ ,		1e-4				)



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         UpperCAmelCase	    =			self.get_dummy_components()
         UpperCAmelCase	    =			self.pipeline_class(**snake_case__				)
         pipe.to(snake_case__				)
         pipe.set_progress_bar_config(disable=snake_case__				)

         UpperCAmelCase	    =			self.get_dummy_inputs(snake_case__				)
         UpperCAmelCase	    =			pipe(**snake_case__				)[0]

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

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

         UpperCAmelCase	    =			self.get_dummy_inputs(snake_case__				)
         UpperCAmelCase	    =			pipe_loaded(**snake_case__				)[0]

         UpperCAmelCase	    =			np.abs(to_np(snake_case__				) - to_np(snake_case__				)				).max()
         self.assertLess(snake_case__ ,		1e-4				)
