


'''simple docstring'''
import random
import unittest

import torch

from diffusers import IFImgaImgSuperResolutionPipeline
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    lowercase	(							A__	,		A__	,		unittest.TestCase					):

		"""simple docstring"""

		_a     =				IFImgaImgSuperResolutionPipeline
		_a     =				TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'width', 'height'}
		_a     =				TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'original_image'}					)
		_a     =				PipelineTesterMixin.required_optional_params - {'latents'}




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							return self._get_superresolution_dummy_components()




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_=0							):

							'''simple docstring'''


							if str(_a							).startswith('''mps'''							):
												UpperCamelCase__   :Dict												= torch.manual_seed(_a							)
							else:
												UpperCamelCase__   :Any												= torch.Generator(device=_a							).manual_seed(_a							)

							UpperCamelCase__   :Union[str, Any]												= floats_tensor((1, 3, 32, 32)    ,   rng=random.Random(_a							)							).to(_a							)
							UpperCamelCase__   :Dict												= floats_tensor((1, 3, 16, 16)    ,   rng=random.Random(_a							)							).to(_a							)

							UpperCamelCase__   :List[str]												= {
							    "prompt": "A painting of a squirrel eating a burger",
							    "image": image,
							    "original_image": original_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     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							self._test_save_load_optional_components()




		@unittest.skipIf(torch_device != '''cuda'''    ,   reason='''float16 requires CUDA'''							)
		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							super().test_save_load_floataa(expected_max_diff=1e-1							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							self._test_attention_slicing_forward_pass(expected_max_diff=1e-2							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							self._test_save_load_local()




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							self._test_inference_batch_single_identical(
							    expected_max_diff=1e-2    ,   )

from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
				import tensorflow as tf

				from transformers import TFLEDForConditionalGeneration, TFLEDModel





@require_tf
class       _snake_case					:
			UpperCamelCase__          =    LEDConfig
			UpperCamelCase__          =    {}
			UpperCamelCase__          =    'gelu'







			def __init__(				self						,   _a						,   _a=13						,   _a=7						,   _a=True						,   _a=False						,   _a=99						,   _a=32						,   _a=2						,   _a=4						,   _a=37						,   _a=0.1						,   _a=0.1						,   _a=20						,   _a=2						,   _a=1						,   _a=0						,   _a=4						,   ):
				__magic_name__		:		int	   =     parent
				__magic_name__		:		Optional[int]	   =     batch_size
				__magic_name__		:		Tuple	   =     seq_length
				__magic_name__		:		List[Any]	   =     is_training
				__magic_name__		:		Dict	   =     use_labels
				__magic_name__		:		Optional[Any]	   =     vocab_size
				__magic_name__		:		int	   =     hidden_size
				__magic_name__		:		Optional[int]	   =     num_hidden_layers
				__magic_name__		:		Optional[int]	   =     num_attention_heads
				__magic_name__		:		Tuple	   =     intermediate_size

				__magic_name__		:		Any	   =     hidden_dropout_prob
				__magic_name__		:		Optional[int]	   =     attention_probs_dropout_prob
				__magic_name__		:		List[str]	   =     max_position_embeddings
				__magic_name__		:		Any	   =     eos_token_id
				__magic_name__		:		str	   =     pad_token_id
				__magic_name__		:		int	   =     bos_token_id
				__magic_name__		:		Optional[int]	   =     attention_window

				# `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
				# [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
				# returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
				# because its local attention only attends to `self.attention_window` and one before and one after
				__magic_name__		:		Tuple	   =     self.attention_window + 2

				# because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
				# the `test_attention_outputs` and `test_hidden_states_output` tests
				__magic_name__		:		Tuple	   =     (
				    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
				)







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__		:		Union[str, Any]	   =     ids_tensor([self.batch_size, self.seq_length - 1]						,   self.vocab_size		)
				__magic_name__		:		int	   =     tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)						,   1		)
				__magic_name__		:		int	   =     tf.concat([input_ids, eos_tensor]						,   axis=1		)

				__magic_name__		:		Optional[int]	   =     ids_tensor([self.batch_size, self.seq_length]						,   self.vocab_size		)

				__magic_name__		:		Dict	   =     self.config_cls(
				    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_ids=[2]						,   bos_token_id=self.bos_token_id						,   pad_token_id=self.pad_token_id						,   decoder_start_token_id=self.pad_token_id						,   attention_window=self.attention_window						,   **self.config_updates						,   )
				__magic_name__		:		List[str]	   =     prepare_led_inputs_dict(_a						,   _a						,   _a		)
				__magic_name__		:		Union[str, Any]	   =     tf.concat(
				    [tf.zeros_like(_a		)[:, :-1], tf.ones_like(_a		)[:, -1:]]						,   axis=-1						,   )
				__magic_name__		:		List[Any]	   =     global_attention_mask
				return config, inputs_dict







			def 			SCREAMING_SNAKE_CASE    (				self						,   _a						,   _a		):
				__magic_name__		:		Dict	   =     TFLEDModel(config=_a		).get_decoder()
				__magic_name__		:		Optional[int]	   =     inputs_dict["input_ids"]

				__magic_name__		:		Union[str, Any]	   =     input_ids[:1, :]
				__magic_name__		:		str	   =     inputs_dict["attention_mask"][:1, :]
				__magic_name__		:		int	   =     1

				# first forward pass
				__magic_name__		:		Tuple	   =     model(_a						,   attention_mask=_a						,   use_cache=_a		)

				__magic_name__   ,    __magic_name__		:		str	   =     outputs.to_tuple()

				# create hypothetical next token and extent to next_input_ids
				__magic_name__		:		Optional[int]	   =     ids_tensor((self.batch_size, 3)						,   config.vocab_size		)
				__magic_name__		:		Any	   =     tf.cast(ids_tensor((self.batch_size, 3)						,   2		)						,   tf.inta		)

				# append to next input_ids and
				__magic_name__		:		Optional[Any]	   =     tf.concat([input_ids, next_tokens]						,   axis=-1		)
				__magic_name__		:		List[Any]	   =     tf.concat([attention_mask, next_attn_mask]						,   axis=-1		)

				__magic_name__		:		List[str]	   =     model(_a						,   attention_mask=_a		)[0]
				__magic_name__		:		Dict	   =     model(_a						,   attention_mask=_a						,   past_key_values=_a		)[0]

				self.parent.assertEqual(next_tokens.shape[1]						,   output_from_past.shape[1]		)

				# select random slice
				__magic_name__		:		List[Any]	   =     int(ids_tensor((1,)						,   output_from_past.shape[-1]		)		)
				__magic_name__		:		Union[str, Any]	   =     output_from_no_past[:, -3:, random_slice_idx]
				__magic_name__		:		List[str]	   =     output_from_past[:, :, random_slice_idx]

				# test that outputs are equal for slice
				tf.debugging.assert_near(_a						,   _a						,   rtol=1e-3		)







def   lowerCAmelCase_						(					_snake_case      :  Any      ,   _snake_case      :  List[Any]      ,   _snake_case      :  Any      ,   _snake_case      :  str=None      ,   _snake_case      :  List[str]=None      ,   _snake_case      :  int=None      ,   _snake_case      :  Any=None      ,   )				->		int:







	'''simple docstring'''


	if attention_mask is None:
		__magic_name__		:		str	   =     tf.cast(tf.math.not_equal(_snake_case      ,   config.pad_token_id   )      ,   tf.inta   )
	if decoder_attention_mask is None:
		__magic_name__		:		List[Any]	   =     tf.concat(
		    [
		        tf.ones(decoder_input_ids[:, :1].shape      ,   dtype=tf.inta   ),
		        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]      ,   config.pad_token_id   )      ,   tf.inta   ),
		    ]      ,   axis=-1      ,   )
	if head_mask is None:
		__magic_name__		:		Dict	   =     tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
	if decoder_head_mask is None:
		__magic_name__		:		Any	   =     tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
	return {
	    "input_ids": input_ids,
	    "attention_mask": attention_mask,
	    "decoder_input_ids": decoder_input_ids,
	    "decoder_attention_mask": decoder_attention_mask,
	    "head_mask": head_mask,
	    "decoder_head_mask": decoder_head_mask,
	}





@require_tf
class       _snake_case					(     snake_case				,						snake_case				,						unittest.TestCase				):
			UpperCamelCase__          =    (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
			UpperCamelCase__          =    (TFLEDForConditionalGeneration,) if is_tf_available() else ()
			UpperCamelCase__          =    (
			    {
			        'conversational': TFLEDForConditionalGeneration,
			        'feature-extraction': TFLEDModel,
			        'summarization': TFLEDForConditionalGeneration,
			        'text2text-generation': TFLEDForConditionalGeneration,
			        'translation': TFLEDForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			UpperCamelCase__          =    True
			UpperCamelCase__          =    False
			UpperCamelCase__          =    False
			UpperCamelCase__          =    False







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__		:		Dict	   =     TFLEDModelTester(self		)
				__magic_name__		:		List[Any]	   =     ConfigTester(self						,   config_class=_a		)







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







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__		:		List[str]	   =     self.model_tester.prepare_config_and_inputs_for_common()
				self.model_tester.check_decoder_model_past_large_inputs(*_a		)







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__   ,    __magic_name__		:		int	   =     self.model_tester.prepare_config_and_inputs_for_common()
				__magic_name__		:		List[str]	   =     tf.zeros_like(inputs_dict["attention_mask"]		)
				__magic_name__		:		Optional[Any]	   =     2
				__magic_name__		:		Tuple	   =     tf.where(
				    tf.range(self.model_tester.seq_length		)[None, :] < num_global_attn_indices						,   1						,   inputs_dict["global_attention_mask"]						,   )

				__magic_name__		:		Any	   =     True
				__magic_name__		:		str	   =     self.model_tester.seq_length
				__magic_name__		:		Dict	   =     self.model_tester.encoder_seq_length

				def check_decoder_attentions_output(_a		):
					__magic_name__		:		str	   =     outputs.decoder_attentions
					self.assertEqual(len(_a		)						,   self.model_tester.num_hidden_layers		)
					self.assertListEqual(
					    list(decoder_attentions[0].shape[-3:]		)						,   [self.model_tester.num_attention_heads, seq_length, seq_length]						,   )

				def check_encoder_attentions_output(_a		):
					__magic_name__		:		Any	   =     [t.numpy() for t in outputs.encoder_attentions]
					__magic_name__		:		Tuple	   =     [t.numpy() for t in outputs.encoder_global_attentions]
					self.assertEqual(len(_a		)						,   self.model_tester.num_hidden_layers		)
					self.assertEqual(len(_a		)						,   self.model_tester.num_hidden_layers		)
					self.assertListEqual(
					    list(attentions[0].shape[-3:]		)						,   [self.model_tester.num_attention_heads, seq_length, seq_length]						,   )
					self.assertListEqual(
					    list(global_attentions[0].shape[-3:]		)						,   [self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]						,   )

				for model_class in self.all_model_classes:
					__magic_name__		:		Union[str, Any]	   =     True
					__magic_name__		:		List[str]	   =     False
					__magic_name__		:		Tuple	   =     False
					__magic_name__		:		Optional[int]	   =     model_class(_a		)
					__magic_name__		:		str	   =     model(self._prepare_for_class(_a						,   _a		)		)
					__magic_name__		:		Any	   =     len(_a		)
					self.assertEqual(config.output_hidden_states						,   _a		)
					check_encoder_attentions_output(_a		)

					if self.is_encoder_decoder:
						__magic_name__		:		Tuple	   =     model_class(_a		)
						__magic_name__		:		Optional[Any]	   =     model(self._prepare_for_class(_a						,   _a		)		)
						self.assertEqual(config.output_hidden_states						,   _a		)
						check_decoder_attentions_output(_a		)

					# Check that output attentions can also be changed via the config
					del inputs_dict["output_attentions"]
					__magic_name__		:		Dict	   =     True
					__magic_name__		:		str	   =     model_class(_a		)
					__magic_name__		:		Any	   =     model(self._prepare_for_class(_a						,   _a		)		)
					self.assertEqual(config.output_hidden_states						,   _a		)
					check_encoder_attentions_output(_a		)

					# Check attention is always last and order is fine
					__magic_name__		:		Union[str, Any]	   =     True
					__magic_name__		:		Union[str, Any]	   =     True
					__magic_name__		:		List[str]	   =     model_class(_a		)
					__magic_name__		:		Any	   =     model(self._prepare_for_class(_a						,   _a		)		)

					self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)						,   len(_a		)		)
					self.assertEqual(model.config.output_hidden_states						,   _a		)
					check_encoder_attentions_output(_a		)







			@unittest.skip("LED keeps using potentially symbolic tensors in conditionals and breaks tracing."		)
			def 			SCREAMING_SNAKE_CASE    (				self		):
				pass







			def 			SCREAMING_SNAKE_CASE    (				self		):
				# TODO: Head-masking not yet implement
				pass







def   lowerCAmelCase_						(					_snake_case      :  int   )				->		Optional[int]:







	'''simple docstring'''


	return tf.constant(_snake_case      ,   dtype=tf.intaa   )


snake_case :		Optional[int]							      =  1E-4





@slow
@require_tf
class       _snake_case					(     unittest.TestCase				):







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__		:		List[Any]	   =     TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"		).led

				# change to intended input here
				__magic_name__		:		Optional[int]	   =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
				__magic_name__		:		str	   =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
				__magic_name__		:		Any	   =     prepare_led_inputs_dict(model.config						,   _a						,   _a		)
				__magic_name__		:		List[Any]	   =     model(**_a		)[0]
				__magic_name__		:		List[str]	   =     (1, 1_024, 768)
				self.assertEqual(output.shape						,   _a		)
				# change to expected output here
				__magic_name__		:		int	   =     tf.convert_to_tensor(
				    [[2.30_50, 2.82_79, 0.65_31], [-1.84_57, -0.14_55, -3.56_61], [-1.01_86, 0.45_86, -2.20_43]]						,   )
				tf.debugging.assert_near(output[:, :3, :3]						,   _a						,   atol=1e-3		)







			def 			SCREAMING_SNAKE_CASE    (				self		):
				__magic_name__		:		Tuple	   =     TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"		)

				# change to intended input here
				__magic_name__		:		int	   =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
				__magic_name__		:		Tuple	   =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
				__magic_name__		:		Optional[Any]	   =     prepare_led_inputs_dict(model.config						,   _a						,   _a		)
				__magic_name__		:		Union[str, Any]	   =     model(**_a		)[0]
				__magic_name__		:		Optional[int]	   =     (1, 1_024, model.config.vocab_size)
				self.assertEqual(output.shape						,   _a		)
				# change to expected output here
				__magic_name__		:		str	   =     tf.convert_to_tensor(
				    [[33.65_07, 6.45_72, 16.80_89], [5.87_39, -2.42_38, 11.29_02], [-3.21_39, -4.31_49, 4.27_83]]						,   )
				tf.debugging.assert_near(output[:, :3, :3]						,   _a						,   atol=1e-3						,   rtol=1e-3		)
