





'''simple docstring'''

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     __lowerCAmelCase					:






       """simple docstring"""
       _snake_case     :					Optional[Any]						    =   XGLMConfig
       _snake_case     :					int						    =   {}
       _snake_case     :					Any						    =   'gelu'





       def __init__(							self       :		List[Any]	,			lowerCAmelCase__       :		Optional[int]	,			lowerCAmelCase__       :		Optional[int]=14	,			lowerCAmelCase__       :		Optional[int]=7	,			lowerCAmelCase__       :		List[str]=True	,			lowerCAmelCase__       :		Optional[int]=True	,			lowerCAmelCase__       :		Union[str, Any]=True	,			lowerCAmelCase__       :		Optional[int]=99	,			lowerCAmelCase__       :		Optional[int]=32	,			lowerCAmelCase__       :		Tuple=2	,			lowerCAmelCase__       :		int=4	,			lowerCAmelCase__       :		List[Any]=37	,			lowerCAmelCase__       :		Any="gelu"	,			lowerCAmelCase__       :		Optional[int]=0.1	,			lowerCAmelCase__       :		List[str]=0.1	,			lowerCAmelCase__       :		str=512	,			lowerCAmelCase__       :		Tuple=0.02	,			)   ->					List[Any]:

            '''simple docstring'''




            _UpperCamelCase       =	parent
            _UpperCamelCase       =	batch_size
            _UpperCamelCase       =	seq_length
            _UpperCamelCase       =	is_training
            _UpperCamelCase       =	use_input_mask
            _UpperCamelCase       =	use_labels
            _UpperCamelCase       =	vocab_size
            _UpperCamelCase       =	d_model
            _UpperCamelCase       =	num_hidden_layers
            _UpperCamelCase       =	num_attention_heads
            _UpperCamelCase       =	ffn_dim
            _UpperCamelCase       =	activation_function
            _UpperCamelCase       =	activation_dropout
            _UpperCamelCase       =	attention_dropout
            _UpperCamelCase       =	max_position_embeddings
            _UpperCamelCase       =	initializer_range
            _UpperCamelCase       =	None
            _UpperCamelCase       =	0
            _UpperCamelCase       =	2
            _UpperCamelCase       =	1





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

            '''simple docstring'''




            return XGLMConfig.from_pretrained('''facebook/xglm-564M'''					)





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

            '''simple docstring'''




            _UpperCamelCase       =	tf.clip_by_value(
                ids_tensor([self.batch_size, self.seq_length]	,			self.vocab_size					)	,			clip_value_min=0	,			clip_value_max=3					)

            _UpperCamelCase       =	None
            if self.use_input_mask:
                 _UpperCamelCase       =	random_attention_mask([self.batch_size, self.seq_length]					)

            _UpperCamelCase       =	self.get_config()

            _UpperCamelCase       =	floats_tensor([self.num_hidden_layers, self.num_attention_heads]	,			2					)

            return (
                config,
                input_ids,
                input_mask,
                head_mask,
            )





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

            '''simple docstring'''




            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 snake_case__						(							self       :		Dict					)   ->					Any:

            '''simple docstring'''




            _UpperCamelCase       =	self.prepare_config_and_inputs()

            (
                (
                _UpperCamelCase
            )   ,    (
                _UpperCamelCase
            )   ,    (
                _UpperCamelCase
            )   ,    (
                _UpperCamelCase
            )   ,    
            )       =	config_and_inputs

            _UpperCamelCase       =	{
                '''input_ids''': input_ids,
                '''head_mask''': head_mask,
            }

            return config, inputs_dict

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






       """simple docstring"""
       _snake_case     :					Optional[int]						    =   (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
       _snake_case     :					str						    =   (TFXGLMForCausalLM,) if is_tf_available() else ()
       _snake_case     :					Optional[int]						    =   (
           {'feature-extraction': TFXGLMModel, 'text-generation': TFXGLMForCausalLM} if is_tf_available() else {}
       )
       _snake_case     :					Any						    =   False
       _snake_case     :					Dict						    =   False
       _snake_case     :					Dict						    =   False





       def snake_case__						(							self       :		Tuple					)   ->					str:

            '''simple docstring'''




            _UpperCamelCase       =	TFXGLMModelTester(self					)
            _UpperCamelCase       =	ConfigTester(self	,			config_class=__A	,			n_embd=37					)





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

            '''simple docstring'''




            self.config_tester.run_common_tests()





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

            '''simple docstring'''




            for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _UpperCamelCase       =	TFXGLMModel.from_pretrained(__A					)
                 self.assertIsNotNone(__A					)







       @unittest.skip(reason='''Currently, model embeddings are going to undergo a major refactor.'''					)
       def snake_case__						(							self       :		Tuple					)   ->					Optional[int]:

            '''simple docstring'''




            super().test_resize_token_embeddings()

@require_tf
class     __lowerCAmelCase					(   unittest.TestCase			):






       """simple docstring"""





       @slow
       def snake_case__						(							self       :		Any	,			lowerCAmelCase__       :		str=True					)   ->					Dict:

            '''simple docstring'''




            _UpperCamelCase       =	TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''					)
            _UpperCamelCase       =	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
            _UpperCamelCase       =	[2, 268, 9865, 67, 11, 1988, 57252, 9865, 5, 984, 67, 1988, 213838, 1658, 53, 70446, 33, 6657, 278, 1581]
            # fmt: on
            _UpperCamelCase       =	model.generate(__A	,			do_sample=__A	,			num_beams=1					)
            if verify_outputs:
                 self.assertListEqual(output_ids[0].numpy().tolist()	,			__A					)





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

            '''simple docstring'''




            _UpperCamelCase       =	XGLMTokenizer.from_pretrained('''facebook/xglm-564M'''					)
            _UpperCamelCase       =	TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''					)

            tf.random.set_seed(0					)
            _UpperCamelCase       =	tokenizer('''Today is a nice day and'''	,			return_tensors='''tf'''					)
            _UpperCamelCase       =	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'''					):
                 _UpperCamelCase       =	model.generate(__A	,			do_sample=__A	,			seed=[7, 0]					)
            _UpperCamelCase       =	tokenizer.decode(output_ids[0]	,			skip_special_tokens=__A					)

            _UpperCamelCase       =	(
                '''Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'''
            )
            self.assertEqual(__A	,			__A					)







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

            '''simple docstring'''




            _UpperCamelCase       =	TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''					)
            _UpperCamelCase       =	XGLMTokenizer.from_pretrained('''facebook/xglm-564M'''					)

            _UpperCamelCase       =	'''left'''

            # use different length sentences to test batching
            _UpperCamelCase       =	[
                '''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''',
            ]

            _UpperCamelCase       =	tokenizer(__A	,			return_tensors='''tf'''	,			padding=__A					)
            _UpperCamelCase       =	inputs['''input_ids''']

            _UpperCamelCase       =	model.generate(input_ids=__A	,			attention_mask=inputs['''attention_mask''']	,			max_new_tokens=12					)

            _UpperCamelCase       =	tokenizer(sentences[0]	,			return_tensors='''tf'''					).input_ids
            _UpperCamelCase       =	model.generate(input_ids=__A	,			max_new_tokens=12					)

            _UpperCamelCase       =	tokenizer(sentences[1]	,			return_tensors='''tf'''					).input_ids
            _UpperCamelCase       =	model.generate(input_ids=__A	,			max_new_tokens=12					)

            _UpperCamelCase       =	tokenizer.batch_decode(__A	,			skip_special_tokens=__A					)
            _UpperCamelCase       =	tokenizer.decode(output_non_padded[0]	,			skip_special_tokens=__A					)
            _UpperCamelCase       =	tokenizer.decode(output_padded[0]	,			skip_special_tokens=__A					)

            _UpperCamelCase       =	[
                '''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]					)













'''simple docstring'''

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     __lowerCAmelCase					(   unittest.TestCase			):






       """simple docstring"""





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

            '''simple docstring'''




            super().tearDown()
            gc.collect()





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

            '''simple docstring'''




            _UpperCamelCase   ,    _UpperCamelCase       =	FlaxControlNetModel.from_pretrained(
                '''lllyasviel/sd-controlnet-canny'''	,			from_pt=lowerCAmelCase__	,			dtype=jnp.bfloataa					)
            _UpperCamelCase   ,    _UpperCamelCase       =	FlaxStableDiffusionControlNetPipeline.from_pretrained(
                '''runwayml/stable-diffusion-v1-5'''	,			controlnet=lowerCAmelCase__	,			from_pt=lowerCAmelCase__	,			dtype=jnp.bfloataa					)
            _UpperCamelCase       =	controlnet_params

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

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

            _UpperCamelCase       =	jax.random.PRNGKey(0					)
            _UpperCamelCase       =	jax.random.split(lowerCAmelCase__	,			jax.device_count()					)

            _UpperCamelCase       =	replicate(lowerCAmelCase__					)
            _UpperCamelCase       =	shard(lowerCAmelCase__					)
            _UpperCamelCase       =	shard(lowerCAmelCase__					)

            _UpperCamelCase       =	pipe(
                prompt_ids=lowerCAmelCase__	,			image=lowerCAmelCase__	,			params=lowerCAmelCase__	,			prng_seed=lowerCAmelCase__	,			num_inference_steps=50	,			jit=lowerCAmelCase__	,			).images
            assert images.shape == (jax.device_count(), 1, 768, 512, 3)

            _UpperCamelCase       =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
            _UpperCamelCase       =	images[0, 253:256, 253:256, -1]

            _UpperCamelCase       =	jnp.asarray(jax.device_get(image_slice.flatten()					)					)
            _UpperCamelCase       =	jnp.array(
                [0.167969, 0.116699, 0.081543, 0.154297, 0.132812, 0.108887, 0.169922, 0.169922, 0.205078]					)
            print(f"""output_slice: {output_slice}"""					)
            assert jnp.abs(output_slice - expected_slice					).max() < 1e-2







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

            '''simple docstring'''




            _UpperCamelCase   ,    _UpperCamelCase       =	FlaxControlNetModel.from_pretrained(
                '''lllyasviel/sd-controlnet-openpose'''	,			from_pt=lowerCAmelCase__	,			dtype=jnp.bfloataa					)
            _UpperCamelCase   ,    _UpperCamelCase       =	FlaxStableDiffusionControlNetPipeline.from_pretrained(
                '''runwayml/stable-diffusion-v1-5'''	,			controlnet=lowerCAmelCase__	,			from_pt=lowerCAmelCase__	,			dtype=jnp.bfloataa					)
            _UpperCamelCase       =	controlnet_params

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

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

            _UpperCamelCase       =	jax.random.PRNGKey(0					)
            _UpperCamelCase       =	jax.random.split(lowerCAmelCase__	,			jax.device_count()					)

            _UpperCamelCase       =	replicate(lowerCAmelCase__					)
            _UpperCamelCase       =	shard(lowerCAmelCase__					)
            _UpperCamelCase       =	shard(lowerCAmelCase__					)

            _UpperCamelCase       =	pipe(
                prompt_ids=lowerCAmelCase__	,			image=lowerCAmelCase__	,			params=lowerCAmelCase__	,			prng_seed=lowerCAmelCase__	,			num_inference_steps=50	,			jit=lowerCAmelCase__	,			).images
            assert images.shape == (jax.device_count(), 1, 768, 512, 3)

            _UpperCamelCase       =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
            _UpperCamelCase       =	images[0, 253:256, 253:256, -1]

            _UpperCamelCase       =	jnp.asarray(jax.device_get(image_slice.flatten()					)					)
            _UpperCamelCase       =	jnp.array(
                [[0.271484, 0.261719, 0.275391, 0.277344, 0.279297, 0.291016, 0.294922, 0.302734, 0.302734]]					)
            print(f"""output_slice: {output_slice}"""					)
            assert jnp.abs(output_slice - expected_slice					).max() < 1e-2



