



'''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 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_					):
  '''simple docstring'''




  __lowercase							:				Any           =				'''new-model'''


if is_tf_available():






 class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_					):
   '''simple docstring'''




   __lowercase							:				List[Any]           =				NewModelConfig






@require_tf
class 	lowerCAmelCase_(    unittest.TestCase					):
  '''simple docstring'''




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      lowerCAmelCase__			:  Tuple											= """bert-base-cased"""
      lowerCAmelCase__			:  Optional[Any]											= AutoConfig.from_pretrained(__UpperCAmelCase				)
      self.assertIsNotNone(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  List[Any]											= TFAutoModel.from_pretrained(__UpperCAmelCase				)
      self.assertIsNotNone(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		Optional[int]:
      lowerCAmelCase__			:  Optional[Any]											= """bert-base-cased"""
      lowerCAmelCase__			:  str											= AutoConfig.from_pretrained(__UpperCAmelCase				)
      self.assertIsNotNone(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  List[Any]											= TFAutoModelForPreTraining.from_pretrained(__UpperCAmelCase				)
      self.assertIsNotNone(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		List[str]:
      for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase__			:  str											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  Tuple											= TFAutoModelForCausalLM.from_pretrained(__UpperCAmelCase				)
          lowerCAmelCase__				,       lowerCAmelCase__			:  Any											= TFAutoModelForCausalLM.from_pretrained(__UpperCAmelCase  ,output_loading_info=__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		Any:
      for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase__			:  int											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  int											= TFAutoModelWithLMHead.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		Any:
      for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase__			:  Union[str, Any]											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  List[Any]											= TFAutoModelForMaskedLM.from_pretrained(__UpperCAmelCase				)
          lowerCAmelCase__				,       lowerCAmelCase__			:  List[Any]											= TFAutoModelForMaskedLM.from_pretrained(__UpperCAmelCase  ,output_loading_info=__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		Tuple:
      for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase__			:  str											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  int											= TFAutoModelForSeqaSeqLM.from_pretrained(__UpperCAmelCase				)
          lowerCAmelCase__				,       lowerCAmelCase__			:  Union[str, Any]											= TFAutoModelForSeqaSeqLM.from_pretrained(__UpperCAmelCase  ,output_loading_info=__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		List[str]:
      # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["bert-base-uncased"]:
          lowerCAmelCase__			:  Tuple											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  List[str]											= TFAutoModelForSequenceClassification.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["bert-base-uncased"]:
          lowerCAmelCase__			:  Dict											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  Any											= TFAutoModelForQuestionAnswering.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  @slow
  @require_tensorflow_probability
  def 							UpperCAmelCase_     (   self				)   ->		str:
      for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
          lowerCAmelCase__			:  int											= AutoConfig.from_pretrained(__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  str											= TFAutoModelForTableQuestionAnswering.from_pretrained(__UpperCAmelCase				)
          lowerCAmelCase__				,       lowerCAmelCase__			:  Optional[int]											= TFAutoModelForTableQuestionAnswering.from_pretrained(
              __UpperCAmelCase  ,output_loading_info=__UpperCAmelCase				)
          self.assertIsNotNone(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		Union[str, Any]:
      lowerCAmelCase__			:  str											= TFAutoModelWithLMHead.from_pretrained(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)
      self.assertEqual(model.num_parameters()  ,1_4410				)
      self.assertEqual(model.num_parameters(only_trainable=__UpperCAmelCase				)  ,1_4410				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      lowerCAmelCase__			:  int											= TFAutoModelWithLMHead.from_pretrained(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)
      self.assertEqual(model.num_parameters()  ,1_4410				)
      self.assertEqual(model.num_parameters(only_trainable=__UpperCAmelCase				)  ,1_4410				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
      lowerCAmelCase__			:  Union[str, Any]											= TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  str											= copy.deepcopy(model.config				)
      lowerCAmelCase__			:  int											= ["""FunnelBaseModel"""]
      lowerCAmelCase__			:  Dict											= TFAutoModel.from_config(__UpperCAmelCase				)
      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

      with tempfile.TemporaryDirectory() as tmp_dir:
          model.save_pretrained(__UpperCAmelCase				)
          lowerCAmelCase__			:  Union[str, Any]											= TFAutoModel.from_pretrained(__UpperCAmelCase				)
          self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		List[str]:
      try:
          AutoConfig.register("""new-model"""  ,__UpperCAmelCase				)

          lowerCAmelCase__			:  Any											= [
              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(__UpperCAmelCase				):
                      auto_class.register(__UpperCAmelCase  ,__UpperCAmelCase				)
                  auto_class.register(__UpperCAmelCase  ,__UpperCAmelCase				)
                  # Trying to register something existing in the Transformers library will raise an error
                  with self.assertRaises(__UpperCAmelCase				):
                      auto_class.register(__UpperCAmelCase  ,__UpperCAmelCase				)

                  # Now that the config is registered, it can be used as any other config with the auto-API
                  lowerCAmelCase__			:  Union[str, Any]											= BertModelTester(self				).get_config()
                  lowerCAmelCase__			:  Dict											= NewModelConfig(**tiny_config.to_dict()				)
                  lowerCAmelCase__			:  List[str]											= auto_class.from_config(__UpperCAmelCase				)
                  self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

                  with tempfile.TemporaryDirectory() as tmp_dir:
                      model.save_pretrained(__UpperCAmelCase				)
                      lowerCAmelCase__			:  Optional[Any]											= auto_class.from_pretrained(__UpperCAmelCase				)
                      self.assertIsInstance(__UpperCAmelCase  ,__UpperCAmelCase				)

      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				)   ->		List[Any]:
      with self.assertRaisesRegex(
          __UpperCAmelCase  ,"""bert-base is not a local folder and is not a valid model identifier"""				):
          lowerCAmelCase__			:  Optional[Any]											= TFAutoModel.from_pretrained("""bert-base"""				)




  def 							UpperCAmelCase_     (   self				)   ->		Tuple:
      with self.assertRaisesRegex(
          __UpperCAmelCase  ,R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""				):
          lowerCAmelCase__			:  Optional[int]											= TFAutoModel.from_pretrained(__UpperCAmelCase  ,revision="""aaaaaa"""				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      with self.assertRaisesRegex(
          __UpperCAmelCase  ,"""hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""  ,):
          lowerCAmelCase__			:  List[Any]											= TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""				)




  def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
      with self.assertRaisesRegex(__UpperCAmelCase  ,"""Use `from_pt=True` to load this model"""				):
          lowerCAmelCase__			:  str											= TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""				)






  def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
      # Make sure we have cached the model.
      lowerCAmelCase__			:  Any											= TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""				)
      with RequestCounter() as counter:
          lowerCAmelCase__			:  List[Any]											= 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
      lowerCAmelCase__			:  Tuple											= TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""				)
      with RequestCounter() as counter:
          lowerCAmelCase__			:  Dict											= 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 random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin






@skip_mps
class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_						,			unittest.TestCase					):
  '''simple docstring'''




  __lowercase							:				List[str]           =				IFInpaintingSuperResolutionPipeline
  __lowercase							:				int           =				TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'''width''', '''height'''}
  __lowercase							:				List[str]           =				TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({'''original_image'''}					)
  __lowercase							:				Optional[Any]           =				PipelineTesterMixin.required_optional_params - {'''latents'''}




  def 							UpperCAmelCase_     (   self				)   ->		Any:
      return self._get_superresolution_dummy_components()




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase=0				)   ->		List[Any]:
      if str(__UpperCAmelCase				).startswith("""mps"""				):
          lowerCAmelCase__			:  Any											= torch.manual_seed(__UpperCAmelCase				)
      else:
          lowerCAmelCase__			:  Dict											= torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)

      lowerCAmelCase__			:  Tuple											= floats_tensor((1, 3, 16, 16)  ,rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= floats_tensor((1, 3, 32, 32)  ,rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)
      lowerCAmelCase__			:  Any											= floats_tensor((1, 3, 32, 32)  ,rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)

      lowerCAmelCase__			:  Dict											= {
          """prompt""": """A painting of a squirrel eating a burger""",
          """image""": image,
          """original_image""": original_image,
          """mask_image""": mask_image,
          """generator""": generator,
          """num_inference_steps""": 2,
          """output_type""": """numpy""",
      }

      return inputs




  @unittest.skipIf(
      torch_device != """cuda""" or not is_xformers_available()  ,reason="""XFormers attention is only available with CUDA and `xformers` installed"""  ,)
  def 							UpperCAmelCase_     (   self				)   ->		Union[str, Any]:
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)




  def 							UpperCAmelCase_     (   self				)   ->		Optional[int]:
      self._test_save_load_optional_components()




  @unittest.skipIf(torch_device != """cuda"""  ,reason="""float16 requires CUDA"""				)
  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
      super().test_save_load_floataa(expected_max_diff=1E-1				)




  def 							UpperCAmelCase_     (   self				)   ->		str:
      self._test_attention_slicing_forward_pass(expected_max_diff=1E-2				)




  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      self._test_save_load_local()






  def 							UpperCAmelCase_     (   self				)   ->		int:
      self._test_inference_batch_single_identical(
          expected_max_diff=1E-2  ,)


