

from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
	import tensorflow as tf

	from transformers.models.xglm.modeling_tf_xglm import (
	    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFXGLMForCausalLM,
	    TFXGLMModel,
	)



@require_tf
class        _a		:


					'''simple docstring'''






					UpperCAmelCase__:  Dict    				=			XGLMConfig
					UpperCAmelCase__:  List[Any]    				=			{}
					UpperCAmelCase__:  List[str]    				=			'''gelu'''






					def __init__(					self       ,			A__       ,			A__=14       ,			A__=7       ,			A__=True       ,			A__=True       ,			A__=True       ,			A__=99       ,			A__=32       ,			A__=2       ,			A__=4       ,			A__=37       ,			A__="gelu"       ,			A__=0.1       ,			A__=0.1       ,			A__=512       ,			A__=0.0_2       ,			):
						A__				:       List[str]	  =						parent
						A__				:       Dict	  =						batch_size
						A__				:       Optional[Any]	  =						seq_length
						A__				:       Union[str, Any]	  =						is_training
						A__				:       int	  =						use_input_mask
						A__				:       List[Any]	  =						use_labels
						A__				:       Tuple	  =						vocab_size
						A__				:       Optional[int]	  =						d_model
						A__				:       Optional[Any]	  =						num_hidden_layers
						A__				:       List[Any]	  =						num_attention_heads
						A__				:       Optional[Any]	  =						ffn_dim
						A__				:       Optional[Any]	  =						activation_function
						A__				:       Dict	  =						activation_dropout
						A__				:       Tuple	  =						attention_dropout
						A__				:       Tuple	  =						max_position_embeddings
						A__				:       List[str]	  =						initializer_range
						A__				:       int	  =						None
						A__				:       Optional[Any]	  =						0
						A__				:       Optional[int]	  =						2
						A__				:       Any	  =						1






					def     __A     (					self     ):
						return XGLMConfig.from_pretrained("""facebook/xglm-564M"""     )






					def     __A     (					self     ):
						A__				:       int	  =						tf.clip_by_value(
						    ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size     )       ,			clip_value_min=0       ,			clip_value_max=3     )

						A__				:       Optional[int]	  =						None
						if self.use_input_mask:
							A__				:       List[Any]	  =						random_attention_mask([self.batch_size, self.seq_length]     )

						A__				:       str	  =						self.get_config()

						A__				:       Union[str, Any]	  =						floats_tensor([self.num_hidden_layers, self.num_attention_heads]       ,			2     )

						return (
						    config,
						    input_ids,
						    input_mask,
						    head_mask,
						)






					def     __A     (					self     ):
						return XGLMConfig(
						    vocab_size=self.vocab_size       ,			d_model=self.hidden_size       ,			num_layers=self.num_hidden_layers       ,			attention_heads=self.num_attention_heads       ,			ffn_dim=self.ffn_dim       ,			activation_function=self.activation_function       ,			activation_dropout=self.activation_dropout       ,			attention_dropout=self.attention_dropout       ,			max_position_embeddings=self.max_position_embeddings       ,			initializer_range=self.initializer_range       ,			use_cache=A__       ,			bos_token_id=self.bos_token_id       ,			eos_token_id=self.eos_token_id       ,			pad_token_id=self.pad_token_id       ,			return_dict=A__       ,			)






					def     __A     (					self     ):
						A__				:       Union[str, Any]	  =						self.prepare_config_and_inputs()

						(
						    (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  
						)				:       Optional[int]	  =						config_and_inputs

						A__				:       List[Any]	  =						{
						    """input_ids""": input_ids,
						    """head_mask""": head_mask,
						}

						return config, inputs_dict



@require_tf
class        _a		(__magic_name__						,  __magic_name__						,  unittest.TestCase							):


					'''simple docstring'''






					UpperCAmelCase__:  int    				=			(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
					UpperCAmelCase__:  Union[str, Any]    				=			(TFXGLMForCausalLM,) if is_tf_available() else ()
					UpperCAmelCase__:  str    				=			(
					    {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
					)
					UpperCAmelCase__:  Optional[int]    				=			False
					UpperCAmelCase__:  Any    				=			False
					UpperCAmelCase__:  int    				=			False






					def     __A     (					self     ):
						A__				:       Optional[int]	  =						TFXGLMModelTester(self     )
						A__				:       Optional[int]	  =						ConfigTester(self       ,			config_class=A__       ,			n_embd=37     )






					def     __A     (					self     ):
						self.config_tester.run_common_tests()






					@slow
					def     __A     (					self     ):
						for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__				:       Optional[int]	  =						TFXGLMModel.from_pretrained(A__     )
							self.assertIsNotNone(A__     )






					@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""     )
					def     __A     (					self     ):
						super().test_resize_token_embeddings()



@require_tf
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					@slow
					def     __A     (					self       ,			A__=True     ):
						A__				:       Optional[Any]	  =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""     )
						A__				:       str	  =						tf.convert_to_tensor([[2, 268, 9865]]       ,			dtype=tf.intaa     )  # The dog
						# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
						# fmt: off
						A__				:       List[str]	  =						[2, 268, 9865, 67, 11, 1988, 5_7252, 9865, 5, 984, 67, 1988, 21_3838, 1658, 53, 7_0446, 33, 6657, 278, 1581]
						# fmt: on
						A__				:       List[Any]	  =						model.generate(A__       ,			do_sample=A__       ,			num_beams=1     )
						if verify_outputs:
							self.assertListEqual(output_ids[0].numpy().tolist()       ,			A__     )






					@slow
					def     __A     (					self     ):
						A__				:       Optional[int]	  =						XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""     )
						A__				:       Optional[int]	  =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""     )

						tf.random.set_seed(0     )
						A__				:       int	  =						tokenizer("""Today is a nice day and"""       ,			return_tensors="""tf"""     )
						A__				:       List[Any]	  =						tokenized.input_ids
						# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
						with tf.device(""":/CPU:0"""     ):
							A__				:       List[Any]	  =						model.generate(A__       ,			do_sample=A__       ,			seed=[7, 0]     )
						A__				:       List[Any]	  =						tokenizer.decode(output_ids[0]       ,			skip_special_tokens=A__     )

						A__				:       Dict	  =						(
						    """Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"""
						)
						self.assertEqual(A__       ,			A__     )






					@slow
					def     __A     (					self     ):
						A__				:       Tuple	  =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""     )
						A__				:       Tuple	  =						XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""     )

						A__				:       Optional[int]	  =						"""left"""

						# use different length sentences to test batching
						A__				:       Union[str, Any]	  =						[
						    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
						    """left-padding, such as in batched generation. The output for the sequence below should be the same """
						    """regardless of whether left padding is applied or not. When""",
						    """Hello, my dog is a little""",
						]

						A__				:       Dict	  =						tokenizer(A__       ,			return_tensors="""tf"""       ,			padding=A__     )
						A__				:       int	  =						inputs["""input_ids"""]

						A__				:       Dict	  =						model.generate(input_ids=A__       ,			attention_mask=inputs["""attention_mask"""]       ,			max_new_tokens=12     )

						A__				:       Optional[int]	  =						tokenizer(sentences[0]       ,			return_tensors="""tf"""     ).input_ids
						A__				:       Dict	  =						model.generate(input_ids=A__       ,			max_new_tokens=12     )

						A__				:       List[Any]	  =						tokenizer(sentences[1]       ,			return_tensors="""tf"""     ).input_ids
						A__				:       Dict	  =						model.generate(input_ids=A__       ,			max_new_tokens=12     )

						A__				:       Optional[Any]	  =						tokenizer.batch_decode(A__       ,			skip_special_tokens=A__     )
						A__				:       Tuple	  =						tokenizer.decode(output_non_padded[0]       ,			skip_special_tokens=A__     )
						A__				:       List[Any]	  =						tokenizer.decode(output_padded[0]       ,			skip_special_tokens=A__     )

						A__				:       str	  =						[
						    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
						    """left-padding, such as in batched generation. The output for the sequence below should be the same """
						    """regardless of whether left padding is applied or not. When left padding is applied, the sequence will be """
						    """a single""",
						    """Hello, my dog is a little bit of a shy one, but he is very friendly""",
						]
						self.assertListEqual(A__       ,			A__     )
						self.assertListEqual(A__       ,			[non_padded_sentence, padded_sentence]     )






import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


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

	from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline



@require_flax
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					def     __A     (					self     ):
						with tempfile.TemporaryDirectory() as tmpdirname:
							# pipeline has Flax weights
							A__				:       Optional[int]	  =						FlaxDiffusionPipeline.from_pretrained(
							    """hf-internal-testing/tiny-stable-diffusion-pipe"""       ,			safety_checker=A__       ,			cache_dir=A__     )

							A__				:       int	  =						[t[-1] for t in os.walk(os.path.join(A__       ,			os.listdir(A__     )[0]       ,			"""snapshots"""     )     )]
							A__				:       str	  =						[item for sublist in all_root_files for item in sublist]

							# None of the downloaded files should be a PyTorch file even if we have some here:
							# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
							assert not any(f.endswith(""".bin"""     ) for f in files     )



@slow
@require_flax
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					def     __A     (					self     ):
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """hf-internal-testing/tiny-stable-diffusion-pipe"""       ,			safety_checker=A__     )

						A__				:       Optional[int]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       int	  =						4

						A__				:       Optional[Any]	  =						jax.device_count()
						A__				:       Union[str, Any]	  =						num_samples * [prompt]
						A__				:       str	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       str	  =						replicate(A__     )
						A__				:       Any	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       int	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 64, 64, 3)
						if jax.device_count() == 8:
							assert np.abs(np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 4.1_5_1_4_7_4_5     ) < 1e-3
							assert np.abs(np.abs(A__       ,			dtype=np.floataa     ).sum() - 4_9_9_4_7.8_7_5     ) < 5e-1

						A__				:       str	  =						pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]     )     )     )
						assert len(A__     ) == num_samples






					def     __A     (					self     ):
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""flax"""       ,			safety_checker=A__     )

						A__				:       Tuple	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       Optional[Any]	  =						50

						A__				:       List[str]	  =						jax.device_count()
						A__				:       Dict	  =						num_samples * [prompt]
						A__				:       Union[str, Any]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       int	  =						replicate(A__     )
						A__				:       List[str]	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       Optional[Any]	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_5_6_5_2_4_0_1)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_8_3_8_0_8.2)     ) < 5e-1






					def     __A     (					self     ):
						A__  ,  A__				:       int	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__     )

						A__				:       Any	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Dict	  =						jax.random.PRNGKey(0     )
						A__				:       str	  =						50

						A__				:       Any	  =						jax.device_count()
						A__				:       List[Any]	  =						num_samples * [prompt]
						A__				:       Any	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Any	  =						replicate(A__     )
						A__				:       str	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       Dict	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_0_0_3_9_0_6)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_7_3_5_1_6.7_5)     ) < 5e-1






					def     __A     (					self     ):
						A__  ,  A__				:       Tuple	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa     )

						A__				:       List[str]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       int	  =						jax.random.PRNGKey(0     )
						A__				:       Optional[int]	  =						50

						A__				:       Optional[Any]	  =						jax.device_count()
						A__				:       Optional[int]	  =						num_samples * [prompt]
						A__				:       Optional[int]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Optional[int]	  =						replicate(A__     )
						A__				:       Tuple	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[int]	  =						shard(A__     )

						A__				:       Optional[int]	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_0_0_3_9_0_6)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_7_3_5_1_6.7_5)     ) < 5e-1






					def     __A     (					self     ):
						A__				:       int	  =						FlaxDDIMScheduler(
						    beta_start=0.0_0_0_8_5       ,			beta_end=0.0_1_2       ,			beta_schedule="""scaled_linear"""       ,			set_alpha_to_one=A__       ,			steps_offset=1       ,			)

						A__  ,  A__				:       Tuple	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			scheduler=A__       ,			safety_checker=A__       ,			)
						A__				:       Union[str, Any]	  =						scheduler.create_state()

						A__				:       Optional[Any]	  =						scheduler_state

						A__				:       Union[str, Any]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       List[Any]	  =						50

						A__				:       Optional[int]	  =						jax.device_count()
						A__				:       Dict	  =						num_samples * [prompt]
						A__				:       List[Any]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Tuple	  =						replicate(A__     )
						A__				:       Dict	  =						jax.random.split(A__       ,			A__     )
						A__				:       Dict	  =						shard(A__     )

						A__				:       int	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_5_0_4_3_9_4_5)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_4_7_6_9_3.5)     ) < 5e-1






					def     __A     (					self     ):
						A__				:       Optional[int]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Optional[int]	  =						jax.device_count()
						A__				:       str	  =						num_samples * [prompt]
						A__				:       Any	  =						jax.random.split(jax.random.PRNGKey(0     )       ,			A__     )

						A__  ,  A__				:       Optional[Any]	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__       ,			)

						A__				:       Tuple	  =						replicate(A__     )
						A__				:       Dict	  =						pipeline.prepare_inputs(A__     )
						A__				:       str	  =						shard(A__     )
						A__				:       Tuple	  =						pipeline(A__       ,			A__       ,			A__       ,			jit=A__     ).images
						assert images.shape == (num_samples, 1, 512, 512, 3)
						A__				:       Optional[Any]	  =						images[2, 0, 256, 10:17, 1]

						# With memory efficient attention
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__       ,			use_memory_efficient_attention=A__       ,			)

						A__				:       Any	  =						replicate(A__     )
						A__				:       Optional[int]	  =						pipeline.prepare_inputs(A__     )
						A__				:       Dict	  =						shard(A__     )
						A__				:       Tuple	  =						pipeline(A__       ,			A__       ,			A__       ,			jit=A__     ).images
						assert images_eff.shape == (num_samples, 1, 512, 512, 3)
						A__				:       Union[str, Any]	  =						images[2, 0, 256, 10:17, 1]

						# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
						# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
						assert abs(slice_eff - slice     ).max() < 1e-2



