





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       __A							(   lowerCamelCase__	):




      """simple docstring"""







      UpperCAmelCase__							=   """new-model"""


if is_tf_available():







       class       __A							(   lowerCamelCase__	):




             """simple docstring"""







             UpperCAmelCase__							=   NewModelConfig







@require_tf
class       __A							(   unittest.TestCase	):




      """simple docstring"""







      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Union[str, Any]						=		'''bert-base-cased'''
             _lowerCamelCase				:  Optional[Any]						=		AutoConfig.from_pretrained(a__)
             self.assertIsNotNone(a__)
             self.assertIsInstance(a__   ,				a__)

             _lowerCamelCase				:  List[str]						=		TFAutoModel.from_pretrained(a__)
             self.assertIsNotNone(a__)
             self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  str						=		'''bert-base-cased'''
             _lowerCamelCase				:  List[str]						=		AutoConfig.from_pretrained(a__)
             self.assertIsNotNone(a__)
             self.assertIsInstance(a__   ,				a__)

             _lowerCamelCase				:  int						=		TFAutoModelForPreTraining.from_pretrained(a__)
             self.assertIsNotNone(a__)
             self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _lowerCamelCase				:  Optional[Any]						=		AutoConfig.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)

                    _lowerCamelCase				:  Tuple						=		TFAutoModelForCausalLM.from_pretrained(a__)
                    _lowerCamelCase,	_lowerCamelCase				:  Any						=		TFAutoModelForCausalLM.from_pretrained(a__   ,				output_loading_info=a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _lowerCamelCase				:  Dict						=		AutoConfig.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)

                    _lowerCamelCase				:  Tuple						=		TFAutoModelWithLMHead.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _lowerCamelCase				:  Optional[Any]						=		AutoConfig.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)

                    _lowerCamelCase				:  List[str]						=		TFAutoModelForMaskedLM.from_pretrained(a__)
                    _lowerCamelCase,	_lowerCamelCase				:  Dict						=		TFAutoModelForMaskedLM.from_pretrained(a__   ,				output_loading_info=a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





             for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _lowerCamelCase				:  Union[str, Any]						=		AutoConfig.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)

                    _lowerCamelCase				:  str						=		TFAutoModelForSeqaSeqLM.from_pretrained(a__)
                    _lowerCamelCase,	_lowerCamelCase				:  int						=		TFAutoModelForSeqaSeqLM.from_pretrained(a__   ,				output_loading_info=a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





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

                    _lowerCamelCase				:  Optional[Any]						=		TFAutoModelForSequenceClassification.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      def 					__snake_case      (		self):



             """simple docstring"""





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

                    _lowerCamelCase				:  Tuple						=		TFAutoModelForQuestionAnswering.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      @slow
      @require_tensorflow_probability
      def 					__snake_case      (		self):



             """simple docstring"""





             for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
                    _lowerCamelCase				:  str						=		AutoConfig.from_pretrained(a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)

                    _lowerCamelCase				:  int						=		TFAutoModelForTableQuestionAnswering.from_pretrained(a__)
                    _lowerCamelCase,	_lowerCamelCase				:  Dict						=		TFAutoModelForTableQuestionAnswering.from_pretrained(
                        a__   ,				output_loading_info=a__)
                    self.assertIsNotNone(a__)
                    self.assertIsInstance(a__   ,				a__)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		TFAutoModelWithLMHead.from_pretrained(a__)
             self.assertIsInstance(a__   ,				a__)
             self.assertEqual(model.num_parameters()   ,				1_4410)
             self.assertEqual(model.num_parameters(only_trainable=a__)   ,				1_4410)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		TFAutoModelWithLMHead.from_pretrained(a__)
             self.assertIsInstance(a__   ,				a__)
             self.assertEqual(model.num_parameters()   ,				1_4410)
             self.assertEqual(model.num_parameters(only_trainable=a__)   ,				1_4410)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Tuple						=		TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny''')
             self.assertIsInstance(a__   ,				a__)

             _lowerCamelCase				:  Any						=		copy.deepcopy(model.config)
             _lowerCamelCase				:  int						=		['''FunnelBaseModel''']
             _lowerCamelCase				:  str						=		TFAutoModel.from_config(a__)
             self.assertIsInstance(a__   ,				a__)

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


      def 					__snake_case      (		self):



             """simple docstring"""





             try:
                    AutoConfig.register('''new-model'''   ,				a__)

                    _lowerCamelCase				:  Optional[int]						=		[
                        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
                                  _lowerCamelCase				:  Optional[Any]						=		BertModelTester(self).get_config()
                                  _lowerCamelCase				:  Dict						=		NewModelConfig(**tiny_config.to_dict())
                                  _lowerCamelCase				:  Any						=		auto_class.from_config(a__)
                                  self.assertIsInstance(a__   ,				a__)

                                  with tempfile.TemporaryDirectory() as tmp_dir:
                                         model.save_pretrained(a__)
                                         _lowerCamelCase				:  int						=		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 					__snake_case      (		self):



             """simple docstring"""





             with self.assertRaisesRegex(
                 a__   ,				'''bert-base is not a local folder and is not a valid model identifier'''):
                    _lowerCamelCase				:  List[Any]						=		TFAutoModel.from_pretrained('''bert-base''')


      def 					__snake_case      (		self):



             """simple docstring"""





             with self.assertRaisesRegex(
                 a__   ,				r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''):
                    _lowerCamelCase				:  List[str]						=		TFAutoModel.from_pretrained(a__   ,				revision='''aaaaaa''')


      def 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             with self.assertRaisesRegex(a__   ,				'''Use `from_pt=True` to load this model'''):
                    _lowerCamelCase				:  List[Any]						=		TFAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only''')



      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Dict						=		TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert''')
             with RequestCounter() as counter:
                    _lowerCamelCase				:  int						=		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				:  Dict						=		TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded''')
             with RequestCounter() as counter:
                    _lowerCamelCase				:  Optional[Any]						=		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)








import inspect
from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel, VQModel
from ...schedulers import DDIMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__):



             """simple docstring"""





             super().__init__()
             self.register_modules(vqvae=a__   ,				unet=a__   ,				scheduler=a__)



      @torch.no_grad()
      def __call__(		self   ,				a__ = 1   ,				a__ = None   ,				a__ = 0.0   ,				a__ = 50   ,				a__ = "pil"   ,				a__ = True   ,				**a__   ,				):



             """simple docstring"""





             _lowerCamelCase				:  Tuple						=		randn_tensor(
                 (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)   ,				generator=a__   ,				)
             _lowerCamelCase				:  Optional[int]						=		latents.to(self.device)

             # scale the initial noise by the standard deviation required by the scheduler
             _lowerCamelCase				:  Tuple						=		latents * self.scheduler.init_noise_sigma

             self.scheduler.set_timesteps(a__)

             # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
             _lowerCamelCase				:  Optional[int]						=		'''eta''' in set(inspect.signature(self.scheduler.step).parameters.keys())

             _lowerCamelCase				:  Union[str, Any]						=		{}
             if accepts_eta:
                    _lowerCamelCase				:  List[Any]						=		eta

             for t in self.progress_bar(self.scheduler.timesteps):
                    _lowerCamelCase				:  List[Any]						=		self.scheduler.scale_model_input(a__   ,				a__)
                    # predict the noise residual
                    _lowerCamelCase				:  Optional[int]						=		self.unet(a__   ,				a__).sample
                    # compute the previous noisy sample x_t -> x_t-1
                    _lowerCamelCase				:  List[str]						=		self.scheduler.step(a__   ,				a__   ,				a__   ,				**a__).prev_sample

             # decode the image latents with the VAE
             _lowerCamelCase				:  Optional[int]						=		self.vqvae.decode(a__).sample

             _lowerCamelCase				:  Dict						=		(image / 2 + 0.5).clamp(0   ,				1)
             _lowerCamelCase				:  List[str]						=		image.cpu().permute(0   ,				2   ,				3   ,				1).numpy()
             if output_type == "pil":
                    _lowerCamelCase				:  List[Any]						=		self.numpy_to_pil(a__)

             if not return_dict:
                    return (image,)

             return ImagePipelineOutput(images=a__)

