


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 	snake_case_       (A__				):
					UpperCAmelCase__ :   Union[str, Any]   				=      'new-model'


if is_tf_available():







	class 	snake_case_       (A__				):
						UpperCAmelCase__ :   str   				=      NewModelConfig







@require_tf
class 	snake_case_       (unittest.TestCase				):
					@slow
					def 			lowerCamelCase__(				self       :List[str]				) ->   int:
												a__					=					'bert-base-cased'
												a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

												a__					=					TFAutoModel.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :str				) ->   Any:
												a__					=					'bert-base-cased'
												a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

												a__					=					TFAutoModelForPreTraining.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   int:
												for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForCausalLM.from_pretrained(lowerCamelCase__				)
																			a__				, a__					=					TFAutoModelForCausalLM.from_pretrained(lowerCamelCase__      ,output_loading_info=lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   Any:
												for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :List[str]				) ->   Any:
												for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase__				)
																			a__				, a__					=					TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase__      ,output_loading_info=lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   str:
												for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__				)
																			a__				, a__					=					TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__      ,output_loading_info=lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :int				) ->   str:
												# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												for model_name in ["bert-base-uncased"]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForSequenceClassification.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   Union[str, Any]:
												# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												for model_name in ["bert-base-uncased"]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForQuestionAnswering.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					@slow
					@require_tensorflow_probability
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[Any]:
												for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
																			a__					=					AutoConfig.from_pretrained(lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																			a__					=					TFAutoModelForTableQuestionAnswering.from_pretrained(lowerCamelCase__				)
																			a__				, a__					=					TFAutoModelForTableQuestionAnswering.from_pretrained(
																			    lowerCamelCase__      ,output_loading_info=lowerCamelCase__				)
																			self.assertIsNotNone(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   Tuple:
												a__					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
												self.assertEqual(model.num_parameters()      ,1_44_10				)
												self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase__				)      ,1_44_10				)
					def 			lowerCamelCase__(				self       :int				) ->   Optional[Any]:
												a__					=					TFAutoModelWithLMHead.from_pretrained(lowerCamelCase__				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
												self.assertEqual(model.num_parameters()      ,1_44_10				)
												self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase__				)      ,1_44_10				)
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   Any:
												# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
												a__					=					TFAutoModel.from_pretrained('sgugger/funnel-random-tiny'				)
												self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

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

												with tempfile.TemporaryDirectory() as tmp_dir:
																			model.save_pretrained(lowerCamelCase__				)
																			a__					=					TFAutoModel.from_pretrained(lowerCamelCase__				)
																			self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)
					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   List[Any]:
												try:
																			AutoConfig.register('new-model'      ,lowerCamelCase__				)

																			a__					=					[
																			    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
																																	a__					=					BertModelTester(self				).get_config()
																																	a__					=					NewModelConfig(**tiny_config.to_dict()				)
																																	a__					=					auto_class.from_config(lowerCamelCase__				)
																																	self.assertIsInstance(lowerCamelCase__      ,lowerCamelCase__				)

																																	with tempfile.TemporaryDirectory() as tmp_dir:
																																								model.save_pretrained(lowerCamelCase__				)
																																								a__					=					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 			lowerCamelCase__(				self       :List[Any]				) ->   Union[str, Any]:
												with self.assertRaisesRegex(
												    lowerCamelCase__      ,'bert-base is not a local folder and is not a valid model identifier'				):
																			a__					=					TFAutoModel.from_pretrained('bert-base'				)
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   str:
												with self.assertRaisesRegex(
												    lowerCamelCase__      ,R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'				):
																			a__					=					TFAutoModel.from_pretrained(lowerCamelCase__      ,revision='aaaaaa'				)
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   Tuple:
												with self.assertRaisesRegex(
												    lowerCamelCase__      ,'hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'      ,):
																			a__					=					TFAutoModel.from_pretrained('hf-internal-testing/config-no-model'				)
					def 			lowerCamelCase__(				self       :Tuple				) ->   Tuple:
												with self.assertRaisesRegex(lowerCamelCase__      ,'Use `from_pt=True` to load this model'				):
																			a__					=					TFAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'				)



					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   List[Any]:
												# Make sure we have cached the model.
												a__					=					TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'				)
												with RequestCounter() as counter:
																			a__					=					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
												a__					=					TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'				)
												with RequestCounter() as counter:
																			a__					=					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 gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard







@slow
@require_flax
class 	snake_case_       (unittest.TestCase				):
					def 			lowerCamelCase__(				self       :List[Any]				) ->   str:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
					def 			lowerCamelCase__(				self       :int				) ->   Optional[Any]:
												a__				, a__					=					FlaxControlNetModel.from_pretrained(
												    'lllyasviel/sd-controlnet-canny'      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__				, a__					=					FlaxStableDiffusionControlNetPipeline.from_pretrained(
												    'runwayml/stable-diffusion-v1-5'      ,controlnet=__snake_case      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__					=					controlnet_params

												a__					=					'bird'
												a__					=					jax.device_count()
												a__					=					pipe.prepare_text_inputs([prompts] * num_samples				)

												a__					=					load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'				)
												a__					=					pipe.prepare_image_inputs([canny_image] * num_samples				)

												a__					=					jax.random.PRNGKey(0				)
												a__					=					jax.random.split(__snake_case      ,jax.device_count()				)

												a__					=					replicate(__snake_case				)
												a__					=					shard(__snake_case				)
												a__					=					shard(__snake_case				)

												a__					=					pipe(
												    prompt_ids=__snake_case      ,image=__snake_case      ,params=__snake_case      ,prng_seed=__snake_case      ,num_inference_steps=50      ,jit=__snake_case      ,).images
												assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

												a__					=					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
												a__					=					images[0, 2_53:2_56, 2_53:2_56, -1]

												a__					=					jnp.asarray(jax.device_get(image_slice.flatten()				)				)
												a__					=					jnp.array(
												    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]				)
												print(F'output_slice: {output_slice}'				)
												assert jnp.abs(output_slice - expected_slice				).max() < 1E-2



					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   Optional[Any]:
												a__				, a__					=					FlaxControlNetModel.from_pretrained(
												    'lllyasviel/sd-controlnet-openpose'      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__				, a__					=					FlaxStableDiffusionControlNetPipeline.from_pretrained(
												    'runwayml/stable-diffusion-v1-5'      ,controlnet=__snake_case      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__					=					controlnet_params

												a__					=					'Chef in the kitchen'
												a__					=					jax.device_count()
												a__					=					pipe.prepare_text_inputs([prompts] * num_samples				)

												a__					=					load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'				)
												a__					=					pipe.prepare_image_inputs([pose_image] * num_samples				)

												a__					=					jax.random.PRNGKey(0				)
												a__					=					jax.random.split(__snake_case      ,jax.device_count()				)

												a__					=					replicate(__snake_case				)
												a__					=					shard(__snake_case				)
												a__					=					shard(__snake_case				)

												a__					=					pipe(
												    prompt_ids=__snake_case      ,image=__snake_case      ,params=__snake_case      ,prng_seed=__snake_case      ,num_inference_steps=50      ,jit=__snake_case      ,).images
												assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

												a__					=					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
												a__					=					images[0, 2_53:2_56, 2_53:2_56, -1]

												a__					=					jnp.asarray(jax.device_get(image_slice.flatten()				)				)
												a__					=					jnp.array(
												    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]				)
												print(F'output_slice: {output_slice}'				)
												assert jnp.abs(output_slice - expected_slice				).max() < 1E-2




