




import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


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

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)


@require_flax
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =			jnp.ones((batch_size, length)						) / length
											return scores



				def       __snake_case		( self						):
											lowerCAmelCase   =			None
											lowerCAmelCase   =			20

											lowerCAmelCase   =			self._get_uniform_logits(batch_size=2       , length=lowerCAmelCase__						)

											# tweak scores to not be uniform anymore
											lowerCAmelCase   =			scores.at[1, 5].set((1 / length) + 0.1						)  # peak, 1st batch
											lowerCAmelCase   =			scores.at[1, 10].set((1 / length) - 0.4						)  # valley, 1st batch

											# compute softmax
											lowerCAmelCase   =			jax.nn.softmax(lowerCAmelCase__       , axis=-1						)

											lowerCAmelCase   =			FlaxTemperatureLogitsWarper(temperature=0.5						)
											lowerCAmelCase   =			FlaxTemperatureLogitsWarper(temperature=1.3						)

											lowerCAmelCase   =			jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase__       , scores.copy()       , cur_len=lowerCAmelCase__						)       , axis=-1						)
											lowerCAmelCase   =			jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase__       , scores.copy()       , cur_len=lowerCAmelCase__						)       , axis=-1						)

											# uniform distribution stays uniform
											self.assertTrue(jnp.allclose(probs[0, :]       , warped_prob_sharp[0, :]       , atol=1E-3						)						)
											self.assertTrue(jnp.allclose(probs[0, :]       , warped_prob_smooth[0, :]       , atol=1E-3						)						)

											# sharp peaks get higher, valleys get lower
											self.assertLess(probs[1, :].max()       , warped_prob_sharp[1, :].max()						)
											self.assertGreater(probs[1, :].min()       , warped_prob_sharp[1, :].min()						)

											# smooth peaks get lower, valleys get higher
											self.assertGreater(probs[1, :].max()       , warped_prob_smooth[1, :].max()						)
											self.assertLess(probs[1, :].min()       , warped_prob_smooth[1, :].min()						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			None
											lowerCAmelCase   =			10
											lowerCAmelCase   =			2

											# create ramp distribution
											lowerCAmelCase   =			np.broadcast_to(np.arange(lowerCAmelCase__						)[None, :]       , (batch_size, vocab_size)						).copy()
											lowerCAmelCase   =			ramp_logits[1:, : vocab_size // 2] + vocab_size

											lowerCAmelCase   =			FlaxTopKLogitsWarper(3						)

											lowerCAmelCase   =			top_k_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)

											# check that correct tokens are filtered
											self.assertListEqual(jnp.isinf(scores[0]						).tolist()       , 7 * [True] + 3 * [False]						)
											self.assertListEqual(jnp.isinf(scores[1]						).tolist()       , 2 * [True] + 3 * [False] + 5 * [True]						)

											# check special case
											lowerCAmelCase   =			5
											lowerCAmelCase   =			FlaxTopKLogitsWarper(top_k=1       , filter_value=0.0       , min_tokens_to_keep=3						)

											lowerCAmelCase   =			np.broadcast_to(np.arange(lowerCAmelCase__						)[None, :]       , (batch_size, length)						).copy()
											lowerCAmelCase   =			top_k_warp_safety_check(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)

											# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
											self.assertListEqual((scores == 0.0).sum(axis=-1						).tolist()       , [2, 2]						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			None
											lowerCAmelCase   =			10
											lowerCAmelCase   =			2

											# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
											lowerCAmelCase   =			np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]						)						)

											lowerCAmelCase   =			FlaxTopPLogitsWarper(0.8						)
											lowerCAmelCase   =			np.exp(top_p_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)						)

											# dist should be filtered to keep min num values so that sum is >= top_p
											# exp (-inf) => 0
											lowerCAmelCase   =			np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]						)
											self.assertTrue(np.allclose(lowerCAmelCase__       , lowerCAmelCase__       , atol=1E-3						)						)

											# check edge cases with negative and extreme logits
											lowerCAmelCase   =			np.broadcast_to(np.arange(lowerCAmelCase__						)[None, :]       , (batch_size, vocab_size)						).copy() - (
											    vocab_size // 2
											)

											# make ramp_logits more extreme
											lowerCAmelCase   =			ramp_logits[1] * 100.0

											# make sure at least 2 tokens are kept
											lowerCAmelCase   =			FlaxTopPLogitsWarper(0.9       , min_tokens_to_keep=2       , filter_value=0.0						)
											lowerCAmelCase   =			top_p_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)

											# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
											self.assertListEqual((filtered_dist != 0.0).sum(axis=-1						).tolist()       , [3, 2]						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			20
											lowerCAmelCase   =			4
											lowerCAmelCase   =			0

											lowerCAmelCase   =			FlaxMinLengthLogitsProcessor(min_length=10       , eos_token_id=lowerCAmelCase__						)

											# check that min length is applied at length 5
											lowerCAmelCase   =			ids_tensor((batch_size, 20)       , vocab_size=20						)
											lowerCAmelCase   =			5
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			min_dist_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       , 4 * [-float('''inf'''						)]						)

											# check that min length is not applied anymore at length 15
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			15
											lowerCAmelCase   =			min_dist_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertFalse(jnp.isinf(lowerCAmelCase__						).any()						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			20
											lowerCAmelCase   =			4
											lowerCAmelCase   =			0

											lowerCAmelCase   =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__						)

											# check that all scores are -inf except the bos_token_id score
											lowerCAmelCase   =			ids_tensor((batch_size, 1)       , vocab_size=20						)
											lowerCAmelCase   =			1
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			logits_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]						).all()						)
											self.assertListEqual(scores[:, bos_token_id].tolist()       , 4 * [0]						)  # score for bos_token_id shold be zero

											# check that bos_token_id is not forced if current length is greater than 1
											lowerCAmelCase   =			3
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			logits_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertFalse(jnp.isinf(lowerCAmelCase__						).any()						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			20
											lowerCAmelCase   =			4
											lowerCAmelCase   =			0
											lowerCAmelCase   =			5

											lowerCAmelCase   =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       , eos_token_id=lowerCAmelCase__						)

											# check that all scores are -inf except the eos_token_id when max_length is reached
											lowerCAmelCase   =			ids_tensor((batch_size, 4)       , vocab_size=20						)
											lowerCAmelCase   =			4
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			logits_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]						).all()						)
											self.assertListEqual(scores[:, eos_token_id].tolist()       , 4 * [0]						)  # score for eos_token_id should be zero

											# check that eos_token_id is not forced if max_length is not reached
											lowerCAmelCase   =			3
											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			logits_processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											self.assertFalse(jnp.isinf(lowerCAmelCase__						).any()						)



				def       __snake_case		( self						):
											lowerCAmelCase   =			4
											lowerCAmelCase   =			10
											lowerCAmelCase   =			15
											lowerCAmelCase   =			2
											lowerCAmelCase   =			1
											lowerCAmelCase   =			15

											# dummy input_ids and scores
											lowerCAmelCase   =			ids_tensor((batch_size, sequence_length)       , lowerCAmelCase__						)
											lowerCAmelCase   =			input_ids.copy()

											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			scores.copy()

											# instantiate all dist processors
											lowerCAmelCase   =			FlaxTemperatureLogitsWarper(temperature=0.5						)
											lowerCAmelCase   =			FlaxTopKLogitsWarper(3						)
											lowerCAmelCase   =			FlaxTopPLogitsWarper(0.8						)

											# instantiate all logits processors
											lowerCAmelCase   =			FlaxMinLengthLogitsProcessor(min_length=10       , eos_token_id=lowerCAmelCase__						)
											lowerCAmelCase   =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__						)
											lowerCAmelCase   =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       , eos_token_id=lowerCAmelCase__						)

											lowerCAmelCase   =			10

											# no processor list
											lowerCAmelCase   =			temp_dist_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											lowerCAmelCase   =			top_k_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											lowerCAmelCase   =			top_p_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											lowerCAmelCase   =			min_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											lowerCAmelCase   =			bos_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
											lowerCAmelCase   =			eos_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)

											# with processor list
											lowerCAmelCase   =			FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
											lowerCAmelCase   =			processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)

											# scores should be equal
											self.assertTrue(jnp.allclose(lowerCAmelCase__       , lowerCAmelCase__       , atol=1E-3						)						)

											# input_ids should never be changed
											self.assertListEqual(input_ids.tolist()       , input_ids_comp.tolist()						)






				def       __snake_case		( self						):
											lowerCAmelCase   =			4
											lowerCAmelCase   =			10
											lowerCAmelCase   =			15
											lowerCAmelCase   =			2
											lowerCAmelCase   =			1
											lowerCAmelCase   =			15

											# dummy input_ids and scores
											lowerCAmelCase   =			ids_tensor((batch_size, sequence_length)       , lowerCAmelCase__						)
											lowerCAmelCase   =			input_ids.copy()

											lowerCAmelCase   =			self._get_uniform_logits(lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			scores.copy()

											# instantiate all dist processors
											lowerCAmelCase   =			FlaxTemperatureLogitsWarper(temperature=0.5						)
											lowerCAmelCase   =			FlaxTopKLogitsWarper(3						)
											lowerCAmelCase   =			FlaxTopPLogitsWarper(0.8						)

											# instantiate all logits processors
											lowerCAmelCase   =			FlaxMinLengthLogitsProcessor(min_length=10       , eos_token_id=lowerCAmelCase__						)
											lowerCAmelCase   =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__						)
											lowerCAmelCase   =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       , eos_token_id=lowerCAmelCase__						)

											lowerCAmelCase   =			10

											# no processor list
											def run_no_processor_list(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
																		lowerCAmelCase   =			temp_dist_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		lowerCAmelCase   =			top_k_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		lowerCAmelCase   =			top_p_warp(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		lowerCAmelCase   =			min_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		lowerCAmelCase   =			bos_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		lowerCAmelCase   =			eos_dist_proc(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		return scores

											# with processor list
											def run_processor_list(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
																		lowerCAmelCase   =			FlaxLogitsProcessorList(
																		    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
																		lowerCAmelCase   =			processor(lowerCAmelCase__       , lowerCAmelCase__       , cur_len=lowerCAmelCase__						)
																		return scores

											lowerCAmelCase   =			jax.jit(lowerCAmelCase__						)
											lowerCAmelCase   =			jax.jit(lowerCAmelCase__						)

											lowerCAmelCase   =			jitted_run_no_processor_list(lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__						)
											lowerCAmelCase   =			jitted_run_processor_list(lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__						)

											# scores should be equal
											self.assertTrue(jnp.allclose(lowerCAmelCase__       , lowerCAmelCase__       , atol=1E-3						)						)

											# input_ids should never be changed
											self.assertListEqual(input_ids.tolist()       , input_ids_comp.tolist()						)









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 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: Tuple										=IFInpaintingSuperResolutionPipeline
				__a							: Dict										=TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""width""", """height"""}
				__a							: int										=TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"""original_image"""}  )
				__a							: Union[str, Any]										=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       , )



