





'''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	(       __magic_name__					):






       """simple docstring"""
       _snake_case      :	List[str]									=       'new-model'


if is_tf_available():
   class 							__lowerCAmelCase	(       __magic_name__					):






          """simple docstring"""
          _snake_case      :	List[Any]									=       NewModelConfig
@require_tf
class 							__lowerCAmelCase	(       unittest.TestCase					):






       """simple docstring"""





       @slow
       def snake_case__		(				self  :						Any	)       ->	int:

            '''simple docstring'''




            _UpperCamelCase							       =			'''bert-base-cased'''
            _UpperCamelCase							       =			AutoConfig.from_pretrained(lowerCAmelCase__	)
            self.assertIsNotNone(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)

            _UpperCamelCase							       =			TFAutoModel.from_pretrained(lowerCAmelCase__	)
            self.assertIsNotNone(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





       @slow
       def snake_case__		(				self  :						Dict	)       ->	Tuple:

            '''simple docstring'''




            _UpperCamelCase							       =			'''bert-base-cased'''
            _UpperCamelCase							       =			AutoConfig.from_pretrained(lowerCAmelCase__	)
            self.assertIsNotNone(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)

            _UpperCamelCase							       =			TFAutoModelForPreTraining.from_pretrained(lowerCAmelCase__	)
            self.assertIsNotNone(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





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

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForCausalLM.from_pretrained(lowerCAmelCase__	)
                 _UpperCamelCase       ,							_UpperCamelCase							       =			TFAutoModelForCausalLM.from_pretrained(lowerCAmelCase__			,					output_loading_info=lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





       @slow
       def snake_case__		(				self  :						str	)       ->	str:

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





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

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForMaskedLM.from_pretrained(lowerCAmelCase__	)
                 _UpperCamelCase       ,							_UpperCamelCase							       =			TFAutoModelForMaskedLM.from_pretrained(lowerCAmelCase__			,					output_loading_info=lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





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

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase__	)
                 _UpperCamelCase       ,							_UpperCamelCase							       =			TFAutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase__			,					output_loading_info=lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





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

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForSequenceClassification.from_pretrained(lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





       @slow
       def snake_case__		(				self  :						List[str]	)       ->	int:

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForQuestionAnswering.from_pretrained(lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





       @slow
       @require_tensorflow_probability
       def snake_case__		(				self  :						int	)       ->	Union[str, Any]:

            '''simple docstring'''




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

                 _UpperCamelCase							       =			TFAutoModelForTableQuestionAnswering.from_pretrained(lowerCAmelCase__	)
                 _UpperCamelCase       ,							_UpperCamelCase							       =			TFAutoModelForTableQuestionAnswering.from_pretrained(
                     lowerCAmelCase__			,					output_loading_info=lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)
                 self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)





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

            '''simple docstring'''




            _UpperCamelCase							       =			TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)
            self.assertEqual(model.num_parameters()			,					14410	)
            self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase__	)			,					14410	)





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

            '''simple docstring'''




            _UpperCamelCase							       =			TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)
            self.assertEqual(model.num_parameters()			,					14410	)
            self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase__	)			,					14410	)





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

            '''simple docstring'''




            _UpperCamelCase							       =			TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny'''	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)

            _UpperCamelCase							       =			copy.deepcopy(model.config	)
            _UpperCamelCase							       =			['''FunnelBaseModel''']
            _UpperCamelCase							       =			TFAutoModel.from_config(lowerCAmelCase__	)
            self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)

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





       def snake_case__		(				self  :						Optional[Any]	)       ->	str:

            '''simple docstring'''




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

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

                           # 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(lowerCAmelCase__	)
                           self.assertIsInstance(lowerCAmelCase__			,					lowerCAmelCase__	)

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

            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  :						Union[str, Any]	)       ->	Dict:

            '''simple docstring'''




            with self.assertRaisesRegex(
                lowerCAmelCase__			,					'''bert-base is not a local folder and is not a valid model identifier'''	):
                 _UpperCamelCase							       =			TFAutoModel.from_pretrained('''bert-base'''	)





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

            '''simple docstring'''




            with self.assertRaisesRegex(
                lowerCAmelCase__			,					r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''	):
                 _UpperCamelCase							       =			TFAutoModel.from_pretrained(lowerCAmelCase__			,					revision='''aaaaaa'''	)





       def snake_case__		(				self  :						Optional[Any]	)       ->	Optional[int]:

            '''simple docstring'''




            with self.assertRaisesRegex(
                lowerCAmelCase__			,					'''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 snake_case__		(				self  :						Tuple	)       ->	Optional[int]:

            '''simple docstring'''




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







       def snake_case__		(				self  :						Optional[Any]	)       ->	str:

            '''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	)









import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


SCREAMING_SNAKE_CASE__     :     Union[str, Any]		    =       logging.get_logger(__name__)



class       a_       (     SCREAMING_SNAKE_CASE__     ):





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




				"""simple docstring"""



				super().__init__()
				SCREAMING_SNAKE_CASE_								=						nn.ModuleList(SCREAMING_SNAKE_CASE      )





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = True      ,     )  ->					Union[ControlNetOutput, Tuple]:




				"""simple docstring"""



				for i, (image, scale, controlnet) in enumerate(zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     self.nets      )      ):
							SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						controlnet(
							    SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     )

							# merge samples
							if i == 0:
										SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						down_samples, mid_sample
							else:
										SCREAMING_SNAKE_CASE_								=						[
										    samples_prev + samples_curr
										    for samples_prev, samples_curr in zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      )
										]
										mid_block_res_sample += mid_sample

				return down_block_res_samples, mid_block_res_sample





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = True      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = None      ,     )  ->					int:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						save_directory
				for controlnet in self.nets:
							controlnet.save_pretrained(
							    SCREAMING_SNAKE_CASE      ,     is_main_process=SCREAMING_SNAKE_CASE      ,     save_function=SCREAMING_SNAKE_CASE      ,     safe_serialization=SCREAMING_SNAKE_CASE      ,     variant=SCREAMING_SNAKE_CASE      ,     )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						model_path_to_save + f'_{idx}'





	@classmethod
	def 						A_(							cls      ,     SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )  ->					Tuple:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						[]

				# load controlnet and append to list until no controlnet directory exists anymore
				# first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
				# second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
				SCREAMING_SNAKE_CASE_								=						pretrained_model_path
				while os.path.isdir(SCREAMING_SNAKE_CASE      ):
							SCREAMING_SNAKE_CASE_								=						ControlNetModel.from_pretrained(SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )
							controlnets.append(SCREAMING_SNAKE_CASE      )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						pretrained_model_path + f'_{idx}'

				logger.info(f'{len(SCREAMING_SNAKE_CASE      )} controlnets loaded from {pretrained_model_path}.'      )

				if len(SCREAMING_SNAKE_CASE      ) == 0:
							raise ValueError(
							    f'No ControlNets found under {os.path.dirname(SCREAMING_SNAKE_CASE      )}. Expected at least {pretrained_model_path + "_0"}.'      )

				return cls(SCREAMING_SNAKE_CASE      )
