





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

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


if is_flax_available():
							import os

							# The slow tests are often failing with OOM error on GPU
							# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
							# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
							UpperCAmelCase__  :     List[str]   			='''platform'''

							import jax
							import jax.numpy as jnp

							from transformers import BlenderbotTokenizer
							from transformers.models.blenderbot.modeling_flax_blenderbot import (
							    FlaxBlenderbotForConditionalGeneration,
							    FlaxBlenderbotModel,
							    shift_tokens_right,
							)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase=None     , _UpperCAmelCase=None     , _UpperCAmelCase=None     , _UpperCAmelCase=None     , _UpperCAmelCase=None     , _UpperCAmelCase=None     , )  ->					List[Any]:

				if attention_mask is None:
								lowerCamelCase			=np.where(input_ids != config.pad_token_id     , 1     , 0			)
				if decoder_attention_mask is None:
								lowerCamelCase			=np.where(decoder_input_ids != config.pad_token_id     , 1     , 0			)
				if head_mask is None:
								lowerCamelCase			=np.ones((config.encoder_layers, config.encoder_attention_heads)			)
				if decoder_head_mask is None:
								lowerCamelCase			=np.ones((config.decoder_layers, config.decoder_attention_heads)			)
				if cross_attn_head_mask is None:
								lowerCamelCase			=np.ones((config.decoder_layers, config.decoder_attention_heads)			)
				return {
				    "input_ids": input_ids,
				    "decoder_input_ids": decoder_input_ids,
				    "attention_mask": attention_mask,
				    "decoder_attention_mask": attention_mask,
				}

class  __A       :


				def __init__(						self				,	UpperCAmelCase_				,	UpperCAmelCase_=13				,	UpperCAmelCase_=7				,	UpperCAmelCase_=True				,	UpperCAmelCase_=False				,	UpperCAmelCase_=99				,	UpperCAmelCase_=16				,	UpperCAmelCase_=2				,	UpperCAmelCase_=4				,	UpperCAmelCase_=4				,	UpperCAmelCase_="gelu"				,	UpperCAmelCase_=0.1				,	UpperCAmelCase_=0.1				,	UpperCAmelCase_=32				,	UpperCAmelCase_=2				,	UpperCAmelCase_=1				,	UpperCAmelCase_=0				,	UpperCAmelCase_=0.0_2				,	):
								lowerCamelCase			=parent
								lowerCamelCase			=batch_size
								lowerCamelCase			=seq_length
								lowerCamelCase			=is_training
								lowerCamelCase			=use_labels
								lowerCamelCase			=vocab_size
								lowerCamelCase			=hidden_size
								lowerCamelCase			=num_hidden_layers
								lowerCamelCase			=num_attention_heads
								lowerCamelCase			=intermediate_size
								lowerCamelCase			=hidden_act
								lowerCamelCase			=hidden_dropout_prob
								lowerCamelCase			=attention_probs_dropout_prob
								lowerCamelCase			=max_position_embeddings
								lowerCamelCase			=eos_token_id
								lowerCamelCase			=pad_token_id
								lowerCamelCase			=bos_token_id
								lowerCamelCase			=initializer_range


				def _snake_case    (						self							):
								lowerCamelCase			=np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,	self.vocab_size							)				,	3				,	self.vocab_size							)
								lowerCamelCase			=np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,	dtype=np.intaa							))				,	-1							)

								lowerCamelCase			=shift_tokens_right(lowercase_				,	1				,	2							)

								lowerCamelCase			=BlenderbotConfig(
								    vocab_size=self.vocab_size				,	d_model=self.hidden_size				,	encoder_layers=self.num_hidden_layers				,	decoder_layers=self.num_hidden_layers				,	encoder_attention_heads=self.num_attention_heads				,	decoder_attention_heads=self.num_attention_heads				,	encoder_ffn_dim=self.intermediate_size				,	decoder_ffn_dim=self.intermediate_size				,	dropout=self.hidden_dropout_prob				,	attention_dropout=self.attention_probs_dropout_prob				,	max_position_embeddings=self.max_position_embeddings				,	eos_token_id=self.eos_token_id				,	bos_token_id=self.bos_token_id				,	pad_token_id=self.pad_token_id				,	initializer_range=self.initializer_range				,	use_cache=lowercase_				,	)
								lowerCamelCase			=prepare_blenderbot_inputs_dict(lowercase_				,	lowercase_				,	lowercase_							)
								return config, inputs_dict


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase			=self.prepare_config_and_inputs()
								return config, inputs_dict


				def _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
								lowerCamelCase			=20
								lowerCamelCase			=model_class_name(lowercase_							)

								lowerCamelCase			=model.encode(inputs_dict["""input_ids"""]							)

								lowerCamelCase	,					lowerCamelCase			=(
								    inputs_dict["""decoder_input_ids"""],
								    inputs_dict["""decoder_attention_mask"""],
								)

								lowerCamelCase			=model.init_cache(decoder_input_ids.shape[0]				,	lowercase_				,	lowercase_							)
								lowerCamelCase			=jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,	dtype="""i4"""							)

								lowerCamelCase			=jnp.broadcast_to(
								    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]				,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,	)
								lowerCamelCase			=model.decode(
								    decoder_input_ids[:, :-1]				,	lowercase_				,	decoder_attention_mask=lowercase_				,	past_key_values=lowercase_				,	decoder_position_ids=lowercase_				,	)

								lowerCamelCase			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,	dtype="""i4"""							)
								lowerCamelCase			=model.decode(
								    decoder_input_ids[:, -1:]				,	lowercase_				,	decoder_attention_mask=lowercase_				,	past_key_values=outputs_cache.past_key_values				,	decoder_position_ids=lowercase_				,	)

								lowerCamelCase			=model.decode(lowercase_				,	lowercase_							)

								lowerCamelCase			=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 _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
								lowerCamelCase			=20
								lowerCamelCase			=model_class_name(lowercase_							)

								lowerCamelCase			=model.encode(inputs_dict["""input_ids"""]							)

								lowerCamelCase	,					lowerCamelCase			=(
								    inputs_dict["""decoder_input_ids"""],
								    inputs_dict["""decoder_attention_mask"""],
								)

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

								lowerCamelCase			=model.init_cache(decoder_input_ids.shape[0]				,	lowercase_				,	lowercase_							)
								lowerCamelCase			=jnp.broadcast_to(
								    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]				,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,	)

								lowerCamelCase			=model.decode(
								    decoder_input_ids[:, :-1]				,	lowercase_				,	decoder_attention_mask=lowercase_				,	past_key_values=lowercase_				,	decoder_position_ids=lowercase_				,	)
								lowerCamelCase			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,	dtype="""i4"""							)
								lowerCamelCase			=model.decode(
								    decoder_input_ids[:, -1:]				,	lowercase_				,	past_key_values=outputs_cache.past_key_values				,	decoder_attention_mask=lowercase_				,	decoder_position_ids=lowercase_				,	)

								lowerCamelCase			=model.decode(lowercase_				,	lowercase_				,	decoder_attention_mask=lowercase_							)

								lowerCamelCase			=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_flax
class  __A       (  unittest.TestCase      ):
				UpperCAmelCase_								=			99


				def _snake_case    (						self							):
								lowerCamelCase			=np.array(
								    [
								        [71, 82, 18, 33, 46, 91, 2],
								        [68, 34, 26, 58, 30, 82, 2],
								        [5, 97, 17, 39, 94, 40, 2],
								        [76, 83, 94, 25, 70, 78, 2],
								        [87, 59, 41, 35, 48, 66, 2],
								        [55, 13, 16, 58, 5, 2, 1],  # note padding
								        [64, 27, 31, 51, 12, 75, 2],
								        [52, 64, 86, 17, 83, 39, 2],
								        [48, 61, 9, 24, 71, 82, 2],
								        [26, 1, 60, 48, 22, 13, 2],
								        [21, 5, 62, 28, 14, 76, 2],
								        [45, 98, 37, 86, 59, 48, 2],
								        [70, 70, 50, 9, 28, 0, 2],
								    ]				,	dtype=np.intaa				,	)

								lowerCamelCase			=input_ids.shape[0]
								lowerCamelCase			=BlenderbotConfig(
								    vocab_size=self.vocab_size				,	d_model=24				,	encoder_layers=2				,	decoder_layers=2				,	encoder_attention_heads=2				,	decoder_attention_heads=2				,	encoder_ffn_dim=32				,	decoder_ffn_dim=32				,	max_position_embeddings=48				,	eos_token_id=2				,	pad_token_id=1				,	bos_token_id=0				,	)
								return config, input_ids, batch_size


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase	,					lowerCamelCase			=self._get_config_and_data()
								lowerCamelCase			=FlaxBlenderbotForConditionalGeneration(lowercase_							)
								lowerCamelCase			=lm_model(input_ids=lowercase_							)
								lowerCamelCase			=(batch_size, input_ids.shape[1], config.vocab_size)
								self.assertEqual(outputs["""logits"""].shape				,	lowercase_							)


				def _snake_case    (						self							):
								lowerCamelCase			=BlenderbotConfig(
								    vocab_size=self.vocab_size				,	d_model=14				,	encoder_layers=2				,	decoder_layers=2				,	encoder_attention_heads=2				,	decoder_attention_heads=2				,	encoder_ffn_dim=8				,	decoder_ffn_dim=8				,	max_position_embeddings=48				,	)
								lowerCamelCase			=FlaxBlenderbotForConditionalGeneration(lowercase_							)
								lowerCamelCase			=np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,	dtype=np.intaa							)
								lowerCamelCase			=np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,	dtype=np.intaa							)
								lowerCamelCase			=lm_model(input_ids=lowercase_				,	decoder_input_ids=lowercase_							)
								lowerCamelCase			=(*summary.shape, config.vocab_size)
								self.assertEqual(outputs["""logits"""].shape				,	lowercase_							)


				def _snake_case    (						self							):
								lowerCamelCase			=np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,	dtype=np.intaa							)
								lowerCamelCase			=shift_tokens_right(lowercase_				,	1				,	2							)
								lowerCamelCase			=np.equal(lowercase_				,	1							).astype(np.floataa							).sum()
								lowerCamelCase			=np.equal(lowercase_				,	1							).astype(np.floataa							).sum()
								self.assertEqual(shifted.shape				,	input_ids.shape							)
								self.assertEqual(lowercase_				,	n_pad_before - 1							)
								self.assertTrue(np.equal(shifted[:, 0]				,	2							).all()							)




@require_flax
class  __A       (  a      ,   unittest.TestCase      ,   a      ):
				UpperCAmelCase_								=			True
				UpperCAmelCase_								=			(
				    (
				        FlaxBlenderbotModel,
				        FlaxBlenderbotForConditionalGeneration,
				    )
				    if is_flax_available()
				    else ()
				)
				UpperCAmelCase_								=			(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()


				def _snake_case    (						self							):
								lowerCamelCase			=FlaxBlenderbotModelTester(self							)


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase			=self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
												self.model_tester.check_use_cache_forward(lowercase_				,	lowercase_				,	lowercase_							)


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase			=self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
												self.model_tester.check_use_cache_forward_with_attn_mask(lowercase_				,	lowercase_				,	lowercase_							)


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase			=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__							):
																lowerCamelCase			=self._prepare_for_class(lowercase_				,	lowercase_							)
																lowerCamelCase			=model_class(lowercase_							)

																@jax.jit
																def encode_jitted(UpperCAmelCase_				,	UpperCAmelCase_=None				,	**UpperCAmelCase_							):
																				return model.encode(input_ids=lowercase_				,	attention_mask=lowercase_							)

																with self.subTest("""JIT Enabled"""							):
																				lowerCamelCase			=encode_jitted(**lowercase_							).to_tuple()

																with self.subTest("""JIT Disabled"""							):
																				with jax.disable_jit():
																								lowerCamelCase			=encode_jitted(**lowercase_							).to_tuple()

																self.assertEqual(len(lowercase_							)				,	len(lowercase_							)							)
																for jitted_output, output in zip(lowercase_				,	lowercase_							):
																				self.assertEqual(jitted_output.shape				,	output.shape							)


				def _snake_case    (						self							):
								lowerCamelCase	,					lowerCamelCase			=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__							):
																lowerCamelCase			=model_class(lowercase_							)
																lowerCamelCase			=model.encode(inputs_dict["""input_ids"""]				,	inputs_dict["""attention_mask"""]							)

																lowerCamelCase			={
																    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
																    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
																    """encoder_outputs""": encoder_outputs,
																}

																@jax.jit
																def decode_jitted(UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
																				return model.decode(
																				    decoder_input_ids=lowercase_				,	decoder_attention_mask=lowercase_				,	encoder_outputs=lowercase_				,	)

																with self.subTest("""JIT Enabled"""							):
																				lowerCamelCase			=decode_jitted(**lowercase_							).to_tuple()

																with self.subTest("""JIT Disabled"""							):
																				with jax.disable_jit():
																								lowerCamelCase			=decode_jitted(**lowercase_							).to_tuple()

																self.assertEqual(len(lowercase_							)				,	len(lowercase_							)							)
																for jitted_output, output in zip(lowercase_				,	lowercase_							):
																				self.assertEqual(jitted_output.shape				,	output.shape							)


				@slow
				def _snake_case    (						self							):







								for model_class_name in self.all_model_classes:
												lowerCamelCase			=model_class_name.from_pretrained("""facebook/blenderbot-400M-distill"""							)
												# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
												lowerCamelCase			=np.ones((1, 1)							) * model.config.eos_token_id
												lowerCamelCase			=model(lowercase_							)
												self.assertIsNotNone(lowercase_							)


				@unittest.skipUnless(jax_device != """cpu"""				,	"""3B test too slow on CPU."""							)
				@slow
				def _snake_case    (						self							):
								lowerCamelCase			={"""num_beams""": 1, """early_stopping""": True, """min_length""": 15, """max_length""": 25}
								lowerCamelCase			={"""skip_special_tokens""": True, """clean_up_tokenization_spaces""": True}

								lowerCamelCase			=FlaxBlenderbotForConditionalGeneration.from_pretrained("""facebook/blenderbot-3B"""				,	from_pt=lowercase_							)
								lowerCamelCase			=BlenderbotTokenizer.from_pretrained("""facebook/blenderbot-3B"""							)

								lowerCamelCase			=["""Sam"""]
								lowerCamelCase			=tokenizer(lowercase_				,	return_tensors="""jax"""							)

								lowerCamelCase			=model.generate(**lowercase_				,	**lowercase_							)
								lowerCamelCase			="""Sam is a great name. It means \"sun\" in Gaelic."""

								lowerCamelCase			=tokenizer.batch_decode(lowercase_				,	**lowercase_							)
								assert generated_txt[0].strip() == tgt_text












import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
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_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class  __A       (  a      ,   a      ,   unittest.TestCase      ):
				__A								=			IFInpaintingSuperResolutionPipeline
				__A								=			TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""width""", """height"""}
				__A								=			TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"""original_image"""}      )
				__A								=			PipelineTesterMixin.required_optional_params - {"""latents"""}


				def _snake_case    (						self							):
								return self._get_superresolution_dummy_components()


				def _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_=0							):
								if str(UpperCAmelCase_							).startswith("""mps"""							):
												lowerCamelCase			=torch.manual_seed(UpperCAmelCase_							)
								else:
												lowerCamelCase			=torch.Generator(device=UpperCAmelCase_							).manual_seed(UpperCAmelCase_							)

								lowerCamelCase			=floats_tensor((1, 3, 16, 16)				,	rng=random.Random(UpperCAmelCase_							)							).to(UpperCAmelCase_							)
								lowerCamelCase			=floats_tensor((1, 3, 32, 32)				,	rng=random.Random(UpperCAmelCase_							)							).to(UpperCAmelCase_							)
								lowerCamelCase			=floats_tensor((1, 3, 32, 32)				,	rng=random.Random(UpperCAmelCase_							)							).to(UpperCAmelCase_							)

								lowerCamelCase			={
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """image""": image,
								    """original_image""": original_image,
								    """mask_image""": mask_image,
								    """generator""": generator,
								    """num_inference_steps""": 2,
								    """output_type""": """numpy""",
								}

								return inputs


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


				def _snake_case    (						self							):
								self._test_save_load_optional_components()


				@unittest.skipIf(torch_device != """cuda"""				,	reason="""float16 requires CUDA"""							)
				def _snake_case    (						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							)


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


				def _snake_case    (						self							):
								self._test_save_load_local()


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





