


import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   import jax
   import jax.numpy as jnp
   import numpy as np

   from transformers.modeling_flax_pytorch_utils import (
       convert_pytorch_state_dict_to_flax,
       load_flax_weights_in_pytorch_model,
   )
   from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
   import torch







@require_flax
class lowerCAmelCase_		:

 '''simple docstring'''

 def __init__(      self						, _UpperCAmelCase						, _UpperCAmelCase=14						, _UpperCAmelCase=7						, _UpperCAmelCase=True						, _UpperCAmelCase=True						, _UpperCAmelCase=True						, _UpperCAmelCase=99						, _UpperCAmelCase=32						, _UpperCAmelCase=5						, _UpperCAmelCase=4						, _UpperCAmelCase=37						, _UpperCAmelCase="gelu"						, _UpperCAmelCase=0.1						, _UpperCAmelCase=0.1						, _UpperCAmelCase=5_12						, _UpperCAmelCase=0.02						, _UpperCAmelCase=None						, ):
       snake_case_			= parent
       snake_case_			= batch_size
       snake_case_			= seq_length
       snake_case_			= is_training
       snake_case_			= use_input_mask
       snake_case_			= use_labels
       snake_case_			= vocab_size
       snake_case_			= d_model
       snake_case_			= num_hidden_layers
       snake_case_			= num_attention_heads
       snake_case_			= ffn_dim
       snake_case_			= activation_function
       snake_case_			= activation_dropout
       snake_case_			= attention_dropout
       snake_case_			= max_position_embeddings
       snake_case_			= initializer_range
       snake_case_			= None
       snake_case_			= 0
       snake_case_			= 2
       snake_case_			= 1



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= np.clip(ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size   )						, 3						, self.vocab_size   )

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

       snake_case_			= 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=_UpperCAmelCase						, bos_token_id=self.bos_token_id						, eos_token_id=self.eos_token_id						, pad_token_id=self.pad_token_id						, )

       return (config, input_ids, input_mask)



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.prepare_config_and_inputs()
       snake_case_			,       snake_case_			,       snake_case_			= config_and_inputs
       snake_case_			= {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict



 def 		UpperCamelCase__	(      self   ):
       snake_case_			,       snake_case_			,       snake_case_			= self.prepare_config_and_inputs()

       snake_case_			= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
       snake_case_			= ids_tensor([self.batch_size, self.seq_length]						, vocab_size=2   )

       return (
           config,
           input_ids,
           attention_mask,
           encoder_hidden_states,
           encoder_attention_mask,
       )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= 20
       snake_case_			= model_class_name(_UpperCAmelCase   )

       snake_case_			= model.init_cache(input_ids.shape[0]						, _UpperCAmelCase   )
       snake_case_			= jnp.ones((input_ids.shape[0], max_decoder_length)						, dtype='''i4'''   )

       snake_case_			= jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1   )[None, :]						, (input_ids.shape[0], input_ids.shape[-1] - 1)   )
       snake_case_			= model(
           input_ids[:, :-1]						, attention_mask=_UpperCAmelCase						, past_key_values=_UpperCAmelCase						, position_ids=_UpperCAmelCase						, )

       snake_case_			= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						, dtype='''i4'''   )
       snake_case_			= model(
           input_ids[:, -1:]						, attention_mask=_UpperCAmelCase						, past_key_values=outputs_cache.past_key_values						, position_ids=_UpperCAmelCase						, )

       snake_case_			= model(_UpperCAmelCase   )

       snake_case_			= np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
       self.parent.assertTrue(diff < 1E-3						, msg=F'''Max diff is {diff}'''   )





 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= 20
       snake_case_			= model_class_name(_UpperCAmelCase   )

       snake_case_			= jnp.concatenate(
           [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])   )]						, axis=-1						, )

       snake_case_			= model.init_cache(input_ids.shape[0]						, _UpperCAmelCase   )
       snake_case_			= jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1   )[None, :]						, (input_ids.shape[0], input_ids.shape[-1] - 1)   )

       snake_case_			= model(
           input_ids[:, :-1]						, attention_mask=_UpperCAmelCase						, past_key_values=_UpperCAmelCase						, position_ids=_UpperCAmelCase						, )
       snake_case_			= jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						, dtype='''i4'''   )
       snake_case_			= model(
           input_ids[:, -1:]						, past_key_values=outputs_cache.past_key_values						, attention_mask=_UpperCAmelCase						, position_ids=_UpperCAmelCase						, )

       snake_case_			= model(_UpperCAmelCase						, attention_mask=_UpperCAmelCase   )
       snake_case_			= np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
       self.parent.assertTrue(diff < 1E-3						, msg=F'''Max diff is {diff}'''   )







@require_sentencepiece
@require_flax
class lowerCAmelCase_		(					lowerCamelCase__   ,		lowerCamelCase__   ,		unittest.TestCase			):

 '''simple docstring'''

 __snake_case				=    (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
 __snake_case				=    (FlaxXGLMForCausalLM,) if is_flax_available() else ()



 def 		UpperCamelCase__	(      self   ):
       snake_case_			= FlaxXGLMModelTester(self   )



 def 		UpperCamelCase__	(      self   ):
       for model_class_name in self.all_model_classes:
             snake_case_			,       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(_UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )



 def 		UpperCamelCase__	(      self   ):
       for model_class_name in self.all_model_classes:
             snake_case_			,       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward_with_attn_mask(
                 _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )



 @slow
 def 		UpperCamelCase__	(      self   ):
       snake_case_			= XGLMTokenizer.from_pretrained('''XGLM'''						, padding_side='''left'''   )
       snake_case_			= tokenizer(['''Hello this is a long string''', '''Hey''']						, return_tensors='''np'''						, padding=_UpperCAmelCase						, truncation=_UpperCAmelCase   )

       snake_case_			= FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''   )
       snake_case_			= 1
       snake_case_			= False

       snake_case_			= jax.jit(model.generate   )

       snake_case_			= jit_generate(inputs['''input_ids''']						, attention_mask=inputs['''attention_mask''']   ).sequences

       snake_case_			= tokenizer.batch_decode(_UpperCAmelCase						, skip_special_tokens=_UpperCAmelCase   )

       snake_case_			= [
           '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
           '''Hey, I\'m a newbie to the forum and I\'''',
       ]

       self.assertListEqual(_UpperCAmelCase						, _UpperCAmelCase   )



 @is_pt_flax_cross_test
 def 		UpperCamelCase__	(      self   ):
       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__   ):
                   # prepare inputs
                   snake_case_			= self._prepare_for_class(_UpperCAmelCase						, _UpperCAmelCase   )
                   snake_case_			= {k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   snake_case_			= model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   snake_case_			= getattr(_UpperCAmelCase						, _UpperCAmelCase   )

                   snake_case_			,       snake_case_			= pt_inputs['''input_ids'''].shape
                   snake_case_			= np.random.randint(0						, seq_length - 1						, size=(batch_size,)   )
                   for batch_idx, start_index in enumerate(_UpperCAmelCase   ):
                         snake_case_			= 0
                         snake_case_			= 1
                         snake_case_			= 0
                         snake_case_			= 1
                   snake_case_			= pt_model_class(_UpperCAmelCase   ).eval()
                   # Flax models don't use the `use_cache` option and cache is not returned as a default.
                   # So we disable `use_cache` here for PyTorch model.
                   snake_case_			= False
                   snake_case_			= model_class(_UpperCAmelCase						, dtype=jnp.floataa   )
                   snake_case_			= convert_pytorch_state_dict_to_flax(pt_model.state_dict()						, _UpperCAmelCase   )
                   snake_case_			= fx_state

                   with torch.no_grad():
                         snake_case_			= pt_model(**_UpperCAmelCase   ).to_tuple()

                   snake_case_			= fx_model(**_UpperCAmelCase   ).to_tuple()
                   self.assertEqual(len(_UpperCAmelCase   )						, len(_UpperCAmelCase   )						, '''Output lengths differ between Flax and PyTorch'''   )
                   for fx_output, pt_output in zip(_UpperCAmelCase						, _UpperCAmelCase   ):
                         self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4E-2   )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                         pt_model.save_pretrained(_UpperCAmelCase   )
                         snake_case_			= model_class.from_pretrained(_UpperCAmelCase						, from_pt=_UpperCAmelCase   )

                   snake_case_			= fx_model_loaded(**_UpperCAmelCase   ).to_tuple()
                   self.assertEqual(
                       len(_UpperCAmelCase   )						, len(_UpperCAmelCase   )						, '''Output lengths differ between Flax and PyTorch'''   )
                   for fx_output_loaded, pt_output in zip(_UpperCAmelCase						, _UpperCAmelCase   ):
                         self.assert_almost_equals(fx_output_loaded[:, -1]						, pt_output[:, -1].numpy()						, 4E-2   )



 @is_pt_flax_cross_test
 def 		UpperCamelCase__	(      self   ):
       snake_case_			,       snake_case_			= self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__   ):
                   # prepare inputs
                   snake_case_			= self._prepare_for_class(_UpperCAmelCase						, _UpperCAmelCase   )
                   snake_case_			= {k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   snake_case_			= model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   snake_case_			= getattr(_UpperCAmelCase						, _UpperCAmelCase   )

                   snake_case_			= pt_model_class(_UpperCAmelCase   ).eval()
                   snake_case_			= False
                   snake_case_			= model_class(_UpperCAmelCase						, dtype=jnp.floataa   )

                   snake_case_			= load_flax_weights_in_pytorch_model(_UpperCAmelCase						, fx_model.params   )
                   snake_case_			,       snake_case_			= pt_inputs['''input_ids'''].shape
                   snake_case_			= np.random.randint(0						, seq_length - 1						, size=(batch_size,)   )
                   for batch_idx, start_index in enumerate(_UpperCAmelCase   ):
                         snake_case_			= 0
                         snake_case_			= 1
                         snake_case_			= 0
                         snake_case_			= 1

                   # make sure weights are tied in PyTorch
                   pt_model.tie_weights()

                   with torch.no_grad():
                         snake_case_			= pt_model(**_UpperCAmelCase   ).to_tuple()

                   snake_case_			= fx_model(**_UpperCAmelCase   ).to_tuple()
                   self.assertEqual(len(_UpperCAmelCase   )						, len(_UpperCAmelCase   )						, '''Output lengths differ between Flax and PyTorch'''   )
                   for fx_output, pt_output in zip(_UpperCAmelCase						, _UpperCAmelCase   ):
                         self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4E-2   )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                         fx_model.save_pretrained(_UpperCAmelCase   )
                         snake_case_			= pt_model_class.from_pretrained(_UpperCAmelCase						, from_flax=_UpperCAmelCase   )

                   with torch.no_grad():
                         snake_case_			= pt_model_loaded(**_UpperCAmelCase   ).to_tuple()

                   self.assertEqual(
                       len(_UpperCAmelCase   )						, len(_UpperCAmelCase   )						, '''Output lengths differ between Flax and PyTorch'''   )
                   for fx_output, pt_output in zip(_UpperCAmelCase						, _UpperCAmelCase   ):
                         self.assert_almost_equals(fx_output[:, -1]						, pt_output[:, -1].numpy()						, 4E-2   )





 @slow
 def 		UpperCamelCase__	(      self   ):
       for model_class_name in self.all_model_classes:
             snake_case_			= model_class_name.from_pretrained('''facebook/xglm-564M'''   )
             snake_case_			= model(np.ones((1, 1)   )   )
             self.assertIsNotNone(_UpperCAmelCase   )



import random
import unittest

import torch

from diffusers import IFImgaImgPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin







@skip_mps
class lowerCAmelCase_		(					lowerCamelCase__   ,		lowerCamelCase__   ,		unittest.TestCase			):

 '''simple docstring'''

 __snake_case				=    IFImgaImgPipeline
 __snake_case				=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
 __snake_case				=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 __snake_case				=    PipelineTesterMixin.required_optional_params - {"latents"}



 def 		UpperCamelCase__	(      self   ):
       return self._get_dummy_components()



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase=0   ):
       if str(_UpperCAmelCase   ).startswith('''mps'''   ):
             snake_case_			= torch.manual_seed(_UpperCAmelCase   )
       else:
             snake_case_			= torch.Generator(device=_UpperCAmelCase   ).manual_seed(_UpperCAmelCase   )

       snake_case_			= floats_tensor((1, 3, 32, 32)						, rng=random.Random(_UpperCAmelCase   )   ).to(_UpperCAmelCase   )

       snake_case_			= {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }

       return inputs



 def 		UpperCamelCase__	(      self   ):
       self._test_save_load_optional_components()



 @unittest.skipIf(
     torch_device != '''cuda''' or not is_xformers_available()						, reason='''XFormers attention is only available with CUDA and `xformers` installed'''						, )
 def 		UpperCamelCase__	(      self   ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3   )



 @unittest.skipIf(torch_device != '''cuda'''						, reason='''float16 requires CUDA'''   )
 def 		UpperCamelCase__	(      self   ):
       # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
       super().test_save_load_floataa(expected_max_diff=1E-1   )



 @unittest.skipIf(torch_device != '''cuda'''						, reason='''float16 requires CUDA'''   )
 def 		UpperCamelCase__	(      self   ):
       super().test_floataa_inference(expected_max_diff=1E-1   )



 def 		UpperCamelCase__	(      self   ):
       self._test_attention_slicing_forward_pass(expected_max_diff=1E-2   )



 def 		UpperCamelCase__	(      self   ):
       self._test_save_load_local()





 def 		UpperCamelCase__	(      self   ):
       self._test_inference_batch_single_identical(
           expected_max_diff=1E-2						, )