


from __future__ import annotations

import random
import unittest

from transformers import TransfoXLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers import (
	    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFTransfoXLForSequenceClassification,
	    TFTransfoXLLMHeadModel,
	    TFTransfoXLModel,
	)







class 	snake_case_       :
					def __init__(				self       :Union[str, Any]      ,__snake_case       :List[str]      ,) ->   Union[str, Any]:
												a__					=					parent
												a__					=					13
												a__					=					7
												a__					=					30
												a__					=					self.seq_length + self.mem_len
												a__					=					15
												a__					=					True
												a__					=					True
												a__					=					99
												a__					=					[10, 50, 80]
												a__					=					32
												a__					=					32
												a__					=					4
												a__					=					8
												a__					=					1_28
												a__					=					2
												a__					=					2
												a__					=					None
												a__					=					1
												a__					=					0
												a__					=					3
												a__					=					self.vocab_size - 1
												a__					=					0.01
					def 			lowerCamelCase__(				self       :List[str]				) ->   Dict:
												a__					=					ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size				)
												a__					=					ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size				)

												a__					=					None
												if self.use_labels:
																			a__					=					ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size				)

												a__					=					TransfoXLConfig(
												    vocab_size=self.vocab_size      ,mem_len=self.mem_len      ,clamp_len=self.clamp_len      ,cutoffs=self.cutoffs      ,d_model=self.hidden_size      ,d_embed=self.d_embed      ,n_head=self.num_attention_heads      ,d_head=self.d_head      ,d_inner=self.d_inner      ,div_val=self.div_val      ,n_layer=self.num_hidden_layers      ,eos_token_id=self.eos_token_id      ,pad_token_id=self.vocab_size - 1      ,init_range=self.init_range      ,num_labels=self.num_labels      ,)

												return (config, input_ids_a, input_ids_a, lm_labels)
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   List[str]:
												random.seed(self.seed				)
												tf.random.set_seed(self.seed				)
					def 			lowerCamelCase__(				self       :List[Any]      ,__snake_case       :Dict      ,__snake_case       :List[str]      ,__snake_case       :List[str]      ,__snake_case       :str				) ->   str:
												a__					=					TFTransfoXLModel(_A				)

												a__					=					model(_A				).to_tuple()

												a__					=					{'input_ids': input_ids_a, 'mems': mems_a}

												a__					=					model(_A				).to_tuple()

												self.parent.assertEqual(hidden_states_a.shape      ,(self.batch_size, self.seq_length, self.hidden_size)				)
												self.parent.assertEqual(hidden_states_a.shape      ,(self.batch_size, self.seq_length, self.hidden_size)				)
												self.parent.assertListEqual(
												    [mem.shape for mem in mems_a]      ,[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,)
												self.parent.assertListEqual(
												    [mem.shape for mem in mems_a]      ,[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,)
					def 			lowerCamelCase__(				self       :Optional[Any]      ,__snake_case       :Dict      ,__snake_case       :List[Any]      ,__snake_case       :Union[str, Any]      ,__snake_case       :Tuple				) ->   int:
												a__					=					TFTransfoXLLMHeadModel(_A				)

												a__					=					model(_A				).to_tuple()

												a__					=					{'input_ids': input_ids_a, 'labels': lm_labels}
												a__					=					model(_A				).to_tuple()

												a__					=					model([input_ids_a, mems_a]				).to_tuple()

												a__					=					{'input_ids': input_ids_a, 'mems': mems_a, 'labels': lm_labels}

												a__					=					model(_A				).to_tuple()

												self.parent.assertEqual(lm_logits_a.shape      ,(self.batch_size, self.seq_length, self.vocab_size)				)
												self.parent.assertListEqual(
												    [mem.shape for mem in mems_a]      ,[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,)

												self.parent.assertEqual(lm_logits_a.shape      ,(self.batch_size, self.seq_length, self.vocab_size)				)
												self.parent.assertListEqual(
												    [mem.shape for mem in mems_a]      ,[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,)
					def 			lowerCamelCase__(				self       :Tuple      ,__snake_case       :Union[str, Any]      ,__snake_case       :Dict      ,__snake_case       :Optional[Any]      ,__snake_case       :Optional[int]				) ->   Union[str, Any]:
												a__					=					TFTransfoXLForSequenceClassification(_A				)
												a__					=					model(_A				)
												self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels)				)



					def 			lowerCamelCase__(				self       :int				) ->   Optional[Any]:
												a__					=					self.prepare_config_and_inputs()
												(a__)					=					config_and_inputs
												a__					=					{'input_ids': input_ids_a}
												return config, inputs_dict







@require_tf
class 	snake_case_       (lowerCamelCase_     ,						lowerCamelCase_     ,						unittest.TestCase				):
					UpperCAmelCase__ :   Dict   				=      (
					    (TFTransfoXLModel, TFTransfoXLLMHeadModel, TFTransfoXLForSequenceClassification) if is_tf_available() else ()
					)
					UpperCAmelCase__ :   Union[str, Any]   				=      () if is_tf_available() else ()
					UpperCAmelCase__ :   Tuple   				=      (
					    {
					        """feature-extraction""": TFTransfoXLModel,
					        """text-classification""": TFTransfoXLForSequenceClassification,
					        """text-generation""": TFTransfoXLLMHeadModel,
					        """zero-shot""": TFTransfoXLForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					# TODO: add this test when TFTransfoXLLMHead has a linear output layer implemented
					UpperCAmelCase__ :   Optional[int]   				=      False
					UpperCAmelCase__ :   Optional[Any]   				=      False
					UpperCAmelCase__ :   Tuple   				=      False
					UpperCAmelCase__ :   Optional[int]   				=      False
					def 			lowerCamelCase__(				self       :str      ,__snake_case       :Union[str, Any]      ,__snake_case       :List[str]      ,__snake_case       :Optional[int]      ,__snake_case       :str      ,__snake_case       :Union[str, Any]				) ->   Tuple:
												if pipeline_test_casse_name == "TextGenerationPipelineTests":
																			# Get `ValueError: AttributeError: 'NoneType' object has no attribute 'new_ones'` or `AssertionError`.
																			# `TransfoXLConfig` was never used in pipeline tests: cannot create a simple
																			# tokenizer.
																			return True

												return False
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   int:
												a__					=					TFTransfoXLModelTester(self				)
												a__					=					ConfigTester(self      ,config_class=_A      ,d_embed=37				)
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   Optional[Any]:
												self.config_tester.run_common_tests()
					def 			lowerCamelCase__(				self       :str				) ->   Dict:
												self.model_tester.set_seed()
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_transfo_xl_model(*_A				)
					def 			lowerCamelCase__(				self       :Any				) ->   Union[str, Any]:
												self.model_tester.set_seed()
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_transfo_xl_lm_head(*_A				)
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[int]:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_transfo_xl_for_sequence_classification(*_A				)
					def 			lowerCamelCase__(				self       :List[str]				) ->   List[Any]:
												a__					=					self.model_tester.prepare_config_and_inputs_for_common()
												a__					=					[TFTransfoXLForSequenceClassification]

												for model_class in self.all_model_classes:
																			a__					=					model_class(_A				)
																			assert isinstance(model.get_input_embeddings()      ,tf.keras.layers.Layer				)
																			if model_class in list_other_models_with_output_ebd:
																										a__					=					model.get_output_embeddings()
																										assert isinstance(_A      ,tf.keras.layers.Layer				)
																										a__					=					model.get_bias()
																										assert name is None
																			else:
																										a__					=					model.get_output_embeddings()
																										assert x is None
																										a__					=					model.get_bias()
																										assert name is None
					def 			lowerCamelCase__(				self       :Any				) ->   Dict:
												# TODO JP: Make TransfoXL XLA compliant
												pass
					@slow
					def 			lowerCamelCase__(				self       :int				) ->   Optional[int]:
												for model_name in TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					TFTransfoXLModel.from_pretrained(_A				)
																			self.assertIsNotNone(_A				)



					@unittest.skip(reason='This model doesn\'t play well with fit() due to not returning a single loss.'				)
					def 			lowerCamelCase__(				self       :List[Any]				) ->   List[Any]:
												pass







@require_tf
class 	snake_case_       (unittest.TestCase				):
					@unittest.skip('Skip test until #12651 is resolved.'				)
					@slow
					def 			lowerCamelCase__(				self       :int				) ->   Optional[int]:
												a__					=					TFTransfoXLLMHeadModel.from_pretrained('transfo-xl-wt103'				)
												# fmt: off
												a__					=					tf.convert_to_tensor([[33,12_97,2,1,10_09,4,11_09,1_17_39,47_62,3_58,5,25,2_45,22,17_06,17,2_00_98,5,32_15,21,37,11_10,3,13,10_41,4,24,6_03,4_90,2,7_14_77,2_00_98,10_44_47,2,2_09_61,1,26_04,4,1,3_29,3,62_24,8_31,1_60_02,2,8,6_03,7_89_67,2_95_46,23,8_03,20,25,4_16,5,8,2_32,4,2_77,6,18_55,46_01,3,2_95_46,54,8,36_09,5,5_72_11,49,4,1,2_77,18,8,17_55,1_56_91,3,3_41,25,4_16,6_93,4_25_73,71,17,4_01,94,31,1_79_19,2,2_95_46,78_73,18,1,4_35,23,1_10_11,7_55,5,51_67,3,79_83,98,84,2,2_95_46,32_67,8,36_09,4,1,48_65,10_75,2,60_87,71,6,3_46,8,58_54,3,2_95_46,8_24,14_00,18_68,2,19,1_60,2,3_11,8,54_96,2,2_09_20,17,25,1_50_97,3,24,24,0]]      ,dtype=tf.intaa				)  # noqa: E231
												# fmt: on
												#  In 1991 , the remains of Russian Tsar Nicholas II and his family
												#  ( except for Alexei and Maria ) are discovered .
												#  The voice of Nicholas's young son , Tsarevich Alexei Nikolaevich , narrates the
												#  remainder of the story . 1883 Western Siberia ,
												#  a young Grigori Rasputin is asked by his father and a group of men to perform magic .
												#  Rasputin has a vision and denounces one of the men as a horse thief . Although his
												#  father initially slaps him for making such an accusation , Rasputin watches as the
												#  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
												#  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
												#  with people , even a bishop , begging for his blessing . <eod> </s> <eos>

												# fmt: off
												a__					=					[33,12_97,2,1,10_09,4,11_09,1_17_39,47_62,3_58,5,25,2_45,22,17_06,17,2_00_98,5,32_15,21,37,11_10,3,13,10_41,4,24,6_03,4_90,2,7_14_77,2_00_98,10_44_47,2,2_09_61,1,26_04,4,1,3_29,3,62_24,8_31,1_60_02,2,8,6_03,7_89_67,2_95_46,23,8_03,20,25,4_16,5,8,2_32,4,2_77,6,18_55,46_01,3,2_95_46,54,8,36_09,5,5_72_11,49,4,1,2_77,18,8,17_55,1_56_91,3,3_41,25,4_16,6_93,4_25_73,71,17,4_01,94,31,1_79_19,2,2_95_46,78_73,18,1,4_35,23,1_10_11,7_55,5,51_67,3,79_83,98,84,2,2_95_46,32_67,8,36_09,4,1,48_65,10_75,2,60_87,71,6,3_46,8,58_54,3,2_95_46,8_24,14_00,18_68,2,19,1_60,2,3_11,8,54_96,2,2_09_20,17,25,1_50_97,3,24,24,0,33,1,18_57,2,1,10_09,4,11_09,1_17_39,47_62,3_58,5,25,2_45,28,11_10,3,13,10_41,4,24,6_03,4_90,2,7_14_77,2_00_98,10_44_47,2,2_09_61,1,26_04,4,1,3_29,3,0]  # noqa: E231
												# fmt: on
												#  In 1991, the remains of Russian Tsar Nicholas II and his family (
												#  except for Alexei and Maria ) are discovered. The voice of young son,
												#  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
												#  1883 Western Siberia, a young Grigori Rasputin is asked by his father
												#  and a group of men to perform magic. Rasputin has a vision and
												#  denounces one of the men as a horse thief. Although his father initially
												#  slaps him for making such an accusation, Rasputin watches as the man
												#  is chased outside and beaten. Twenty years later, Rasputin sees a vision
												#  of the Virgin Mary, prompting him to become a priest.
												#  Rasputin quickly becomes famous, with people, even a bishop, begging for
												#  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
												# Nicholas II and his family were discovered. The voice of <unk> young son,
												# Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>

												a__					=					model.generate(_A      ,max_length=2_00      ,do_sample=_A				)
												self.assertListEqual(output_ids[0].numpy().tolist()      ,_A				)








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




