



import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


a_   :Tuple        =      1e-4

if is_torch_available():
 import torch

 from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
 from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder




@require_torch
class 					snake_case__						:






     """simple docstring"""


     def __init__(       self  :  Dict, _snake_case  :  int, _snake_case  :  Optional[int]=1_6, _snake_case  :  Tuple=1_3, _snake_case  :  Any=7, _snake_case  :  Any=1_4, _snake_case  :  Dict=1_0, _snake_case  :  List[Any]=1_9, _snake_case  :  Union[str, Any]=5, _snake_case  :  Any=4, _snake_case  :  int=True, _snake_case  :  str=1_6, _snake_case  :  Dict=2, _snake_case  :  Any=4, _snake_case  :  Optional[int]=4, _snake_case  :  List[str]="gelu", _snake_case  :  Optional[int]=0.1, _snake_case  :  List[Any]=0.1, _snake_case  :  Any=[1, 2, 3, 4, 5], _snake_case  :  Optional[Any]=2_5, _snake_case  :  List[str]=5, )   ->Any:
            snake_case__							:	List[Any]   	=	d_model
            snake_case__							:	Dict   	=	parent
            snake_case__							:	int   	=	batch_size
            snake_case__							:	List[str]   	=	prediction_length
            snake_case__							:	Union[str, Any]   	=	context_length
            snake_case__							:	Dict   	=	cardinality
            snake_case__							:	Any   	=	num_time_features
            snake_case__							:	Union[str, Any]   	=	lags_sequence
            snake_case__							:	Dict   	=	embedding_dimension
            snake_case__							:	str   	=	is_training
            snake_case__							:	Union[str, Any]   	=	hidden_size
            snake_case__							:	Tuple   	=	num_hidden_layers
            snake_case__							:	str   	=	num_attention_heads
            snake_case__							:	Optional[int]   	=	intermediate_size
            snake_case__							:	Dict   	=	hidden_act
            snake_case__							:	Any   	=	hidden_dropout_prob
            snake_case__							:	Union[str, Any]   	=	attention_probs_dropout_prob

            snake_case__							:	Any   	=	context_length
            snake_case__							:	int   	=	prediction_length + label_length
            snake_case__							:	Optional[Any]   	=	label_length

            snake_case__							:	Tuple   	=	moving_average
            snake_case__							:	List[Any]   	=	autocorrelation_factor






     def    lowercase_   (       self  :  Tuple					)   ->Optional[int]:
            return AutoformerConfig(
                d_model=self.d_model, 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, prediction_length=self.prediction_length, context_length=self.context_length, label_length=self.label_length, lags_sequence=self.lags_sequence, num_time_features=self.num_time_features, num_static_categorical_features=1, cardinality=[self.cardinality], embedding_dimension=[self.embedding_dimension], moving_average=self.moving_average, )






     def    lowercase_   (       self  :  Any, _snake_case  :  Any					)   ->Union[str, Any]:
            snake_case__							:	Optional[int]   	=	config.context_length + max(config.lags_sequence					)

            snake_case__							:	str   	=	ids_tensor([self.batch_size, 1], config.cardinality[0]					)
            snake_case__							:	Union[str, Any]   	=	floats_tensor([self.batch_size, _past_length, config.num_time_features]					)
            snake_case__							:	Dict   	=	floats_tensor([self.batch_size, _past_length]					)
            snake_case__							:	Optional[int]   	=	floats_tensor([self.batch_size, _past_length]					) > 0.5

            # decoder inputs
            snake_case__							:	Tuple   	=	floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]					)
            snake_case__							:	Union[str, Any]   	=	floats_tensor([self.batch_size, config.prediction_length]					)

            snake_case__							:	Optional[Any]   	=	{
                'past_values': past_values,
                'static_categorical_features': static_categorical_features,
                'past_time_features': past_time_features,
                'past_observed_mask': past_observed_mask,
                'future_time_features': future_time_features,
                'future_values': future_values,
            }
            return inputs_dict






     def    lowercase_   (       self  :  Tuple					)   ->Any:
            snake_case__							:	Any   	=	self.get_config()
            snake_case__							:	Optional[Any]   	=	self.prepare_autoformer_inputs_dict(A_					)
            return config, inputs_dict






     def    lowercase_   (       self  :  Dict					)   ->List[Any]:
            snake_case__	,				snake_case__							:	List[str]   	=	self.prepare_config_and_inputs()
            return config, inputs_dict






     def    lowercase_   (       self  :  Optional[Any], _snake_case  :  Optional[Any], _snake_case  :  Union[str, Any]					)   ->List[str]:
            snake_case__							:	str   	=	AutoformerModel(config=A_					).to(A_					).eval()
            snake_case__							:	Union[str, Any]   	=	model(**A_					)

            snake_case__							:	str   	=	outputs.encoder_last_hidden_state
            snake_case__							:	Union[str, Any]   	=	outputs.last_hidden_state

            with tempfile.TemporaryDirectory() as tmpdirname:
                   snake_case__							:	Any   	=	model.get_encoder()
                   encoder.save_pretrained(A_					)
                   snake_case__							:	Dict   	=	AutoformerEncoder.from_pretrained(A_					).to(A_					)

            snake_case__	,				snake_case__	,				snake_case__	,				snake_case__	,				snake_case__							:	Union[str, Any]   	=	model.create_network_inputs(**A_					)
            snake_case__	,				snake_case__							:	Union[str, Any]   	=	model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]					)

            snake_case__							:	List[Any]   	=	torch.cat(
                (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...]), dim=-1, )
            snake_case__							:	Tuple   	=	encoder(inputs_embeds=A_					)[0]
            self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3					)

            snake_case__							:	Tuple   	=	(
                torch.mean(transformer_inputs[:, : config.context_length, ...], dim=1					)
                .unsqueeze(1					)
                .repeat(1, config.prediction_length, 1					)
            )
            snake_case__							:	List[str]   	=	torch.zeros(
                [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]], device=enc_input.device, )

            snake_case__							:	str   	=	torch.cat(
                (
                    torch.cat((seasonal_input[:, -config.label_length :, ...], zeros), dim=1					),
                    feature[:, config.context_length - config.label_length :, ...],
                ), dim=-1, )
            snake_case__							:	Any   	=	torch.cat(
                (
                    torch.cat((trend_input[:, -config.label_length :, ...], mean), dim=1					),
                    feature[:, config.context_length - config.label_length :, ...],
                ), dim=-1, )

            with tempfile.TemporaryDirectory() as tmpdirname:
                   snake_case__							:	List[Any]   	=	model.get_decoder()
                   decoder.save_pretrained(A_					)
                   snake_case__							:	List[Any]   	=	AutoformerDecoder.from_pretrained(A_					).to(A_					)

            snake_case__							:	Dict   	=	decoder(
                trend=A_, inputs_embeds=A_, encoder_hidden_states=A_, )[0]

            self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1e-3					)




@require_torch
class 					snake_case__						(       _SCREAMING_SNAKE_CASE				,     _SCREAMING_SNAKE_CASE				,     unittest.TestCase     ):






     """simple docstring"""


     _SCREAMING_SNAKE_CASE						=	(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
     _SCREAMING_SNAKE_CASE						=	(AutoformerForPrediction,) if is_torch_available() else ()
     _SCREAMING_SNAKE_CASE						=	{"feature-extraction": AutoformerModel} if is_torch_available() else {}
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False






     def    lowercase_   (       self  :  Optional[Any]					)   ->List[str]:
            snake_case__							:	Any   	=	AutoformerModelTester(self					)
            snake_case__							:	Tuple   	=	ConfigTester(self, config_class=A_, has_text_modality=A_					)






     def    lowercase_   (       self  :  str					)   ->List[Any]:
            self.config_tester.run_common_tests()






     def    lowercase_   (       self  :  Union[str, Any]					)   ->Optional[Any]:
            snake_case__	,				snake_case__							:	Any   	=	self.model_tester.prepare_config_and_inputs()
            for model_class in self.all_model_classes:
                   snake_case__							:	Optional[int]   	=	model_class(A_					)

                   with tempfile.TemporaryDirectory() as tmpdirname:
                          model.save_pretrained(A_					)
                          snake_case__	,				snake_case__							:	List[str]   	=	model_class.from_pretrained(A_, output_loading_info=A_					)
                   self.assertEqual(info['missing_keys'], []					)






     def    lowercase_   (       self  :  Optional[int]					)   ->int:
            snake_case__							:	Optional[int]   	=	self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_encoder_decoder_model_standalone(*A_					)






     @unittest.skip(reason='Model has no tokens embeddings'					)
     def    lowercase_   (       self  :  Optional[Any]					)   ->int:
            pass






     def    lowercase_   (       self  :  Optional[Any]					)   ->Union[str, Any]:
            snake_case__							:	int   	=	inspect.signature(getattr(A_, 'forward'					)					)
            # The main input is the name of the argument after `self`
            snake_case__							:	List[str]   	=	list(model_signature.parameters.keys()					)[1]
            self.assertEqual(AutoformerModel.main_input_name, A_					)






     def    lowercase_   (       self  :  List[str]					)   ->List[str]:
            snake_case__	,				snake_case__							:	str   	=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   snake_case__							:	Optional[Any]   	=	model_class(A_					)
                   snake_case__							:	Tuple   	=	inspect.signature(model.forward					)
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   snake_case__							:	Any   	=	[*signature.parameters.keys()]

                   snake_case__							:	Tuple   	=	[
                       'past_values',
                       'past_time_features',
                       'past_observed_mask',
                       'static_categorical_features',
                       'static_real_features',
                       'future_values',
                       'future_time_features',
                   ]

                   if model.__class__.__name__ in ["AutoformerForPrediction"]:
                          expected_arg_names.append('future_observed_mask'					)

                   expected_arg_names.extend(
                       [
                           'decoder_attention_mask',
                           'head_mask',
                           'decoder_head_mask',
                           'cross_attn_head_mask',
                           'encoder_outputs',
                           'past_key_values',
                           'output_hidden_states',
                           'output_attentions',
                           'use_cache',
                           'return_dict',
                       ]					)

                   self.assertListEqual(arg_names[: len(A_					)], A_					)






     def    lowercase_   (       self  :  Dict					)   ->Any:
            snake_case__	,				snake_case__							:	Optional[Any]   	=	self.model_tester.prepare_config_and_inputs_for_common()
            snake_case__							:	str   	=	True

            snake_case__							:	List[str]   	=	getattr(self.model_tester, 'seq_length', A_					)
            snake_case__							:	Optional[Any]   	=	getattr(self.model_tester, 'decoder_seq_length', A_					)
            snake_case__							:	Union[str, Any]   	=	getattr(self.model_tester, 'encoder_seq_length', A_					)
            snake_case__							:	int   	=	getattr(self.model_tester, 'd_model', A_					)
            snake_case__							:	int   	=	getattr(self.model_tester, 'num_attention_heads', A_					)
            snake_case__							:	Dict   	=	d_model // num_attention_heads

            for model_class in self.all_model_classes:
                   snake_case__							:	Optional[Any]   	=	True
                   snake_case__							:	List[str]   	=	False
                   snake_case__							:	Any   	=	True
                   snake_case__							:	Optional[int]   	=	model_class(A_					)
                   model.to(A_					)
                   model.eval()
                   with torch.no_grad():
                          snake_case__							:	Union[str, Any]   	=	model(**self._prepare_for_class(A_, A_					)					)
                   snake_case__							:	Optional[Any]   	=	outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                   self.assertEqual(len(A_					), self.model_tester.num_hidden_layers					)

                   # check that output_attentions also work using config
                   del inputs_dict["output_attentions"]
                   snake_case__							:	Dict   	=	True
                   snake_case__							:	int   	=	model_class(A_					)
                   model.to(A_					)
                   model.eval()
                   with torch.no_grad():
                          snake_case__							:	List[Any]   	=	model(**self._prepare_for_class(A_, A_					)					)
                   snake_case__							:	int   	=	outputs.encoder_attentions
                   self.assertEqual(len(A_					), self.model_tester.num_hidden_layers					)

                   self.assertListEqual(
                       list(attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, encoder_seq_length, dim], )
                   snake_case__							:	str   	=	len(A_					)

                   snake_case__							:	Tuple   	=	7

                   if "last_hidden_state" in outputs:
                          correct_outlen += 1

                   if "trend" in outputs:
                          correct_outlen += 1

                   if "past_key_values" in outputs:
                          correct_outlen += 1  # past_key_values have been returned

                   if "loss" in outputs:
                          correct_outlen += 1

                   if "params" in outputs:
                          correct_outlen += 1

                   self.assertEqual(A_, A_					)

                   # decoder attentions
                   snake_case__							:	Any   	=	outputs.decoder_attentions
                   self.assertIsInstance(A_, (list, tuple)					)
                   self.assertEqual(len(A_					), self.model_tester.num_hidden_layers					)
                   self.assertListEqual(
                       list(decoder_attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, decoder_seq_length, dim], )

                   # cross attentions
                   snake_case__							:	Tuple   	=	outputs.cross_attentions
                   self.assertIsInstance(A_, (list, tuple)					)
                   self.assertEqual(len(A_					), self.model_tester.num_hidden_layers					)
                   self.assertListEqual(
                       list(cross_attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, decoder_seq_length, dim], )

            # Check attention is always last and order is fine
            snake_case__							:	Union[str, Any]   	=	True
            snake_case__							:	List[str]   	=	True
            snake_case__							:	str   	=	model_class(A_					)
            model.to(A_					)
            model.eval()
            with torch.no_grad():
                   snake_case__							:	List[str]   	=	model(**self._prepare_for_class(A_, A_					)					)

            self.assertEqual(out_len + 2, len(A_					)					)

            snake_case__							:	int   	=	outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

            self.assertEqual(len(A_					), self.model_tester.num_hidden_layers					)
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, encoder_seq_length, dim], )






     @is_flaky()
     def    lowercase_   (       self  :  Optional[int]					)   ->str:
            super().test_retain_grad_hidden_states_attentions()


def        lowercase_							(A : Any="train-batch.pt"      ):
       snake_case__							:	Any   	=	hf_hub_download(repo_id='hf-internal-testing/tourism-monthly-batch'		,					filename=A		,					repo_type='dataset'      )
       snake_case__							:	Dict   	=	torch.load(A		,					map_location=A      )
       return batch





@require_torch
@slow
class 					snake_case__						(       unittest.TestCase     ):






     """simple docstring"""


     def    lowercase_   (       self  :  List[str]					)   ->Any:
            snake_case__							:	List[Any]   	=	AutoformerModel.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(A_					)
            snake_case__							:	Optional[int]   	=	prepare_batch()

            with torch.no_grad():
                   snake_case__							:	str   	=	model(
                       past_values=batch['past_values'], past_time_features=batch['past_time_features'], past_observed_mask=batch['past_observed_mask'], static_categorical_features=batch['static_categorical_features'], future_values=batch['future_values'], future_time_features=batch['future_time_features'], )[0]

            snake_case__							:	Optional[Any]   	=	torch.Size(
                (6_4, model.config.prediction_length + model.config.label_length, model.config.feature_size)					)
            self.assertEqual(output.shape, A_					)

            snake_case__							:	List[str]   	=	torch.tensor(
                [[0.3_5_9_3, -1.3_3_9_8, 0.6_3_3_0], [0.2_2_7_9, 1.5_3_9_6, -0.1_7_9_2], [0.0_4_5_0, 1.3_2_2_5, -0.2_3_3_5]], device=A_					)
            self.assertTrue(torch.allclose(output[0, :3, :3], A_, atol=A_					)					)






     def    lowercase_   (       self  :  Optional[Any]					)   ->List[str]:
            snake_case__							:	Union[str, Any]   	=	AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(A_					)
            snake_case__							:	List[str]   	=	prepare_batch('val-batch.pt'					)
            with torch.no_grad():
                   snake_case__							:	Dict   	=	model(
                       past_values=batch['past_values'], past_time_features=batch['past_time_features'], past_observed_mask=batch['past_observed_mask'], static_categorical_features=batch['static_categorical_features'], ).encoder_last_hidden_state
            snake_case__							:	List[str]   	=	torch.Size((6_4, model.config.context_length, model.config.d_model)					)
            self.assertEqual(output.shape, A_					)

            snake_case__							:	int   	=	torch.tensor(
                [[-0.0_7_3_4, -0.9_0_3_6, 0.8_3_5_8], [4.7_1_8_6, 2.4_1_1_3, 1.9_5_8_1], [1.7_9_5_3, 2.3_5_5_8, 1.2_9_7_0]], device=A_					)
            self.assertTrue(torch.allclose(output[0, :3, :3], A_, atol=A_					)					)






     def    lowercase_   (       self  :  Any					)   ->List[Any]:
            snake_case__							:	str   	=	AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(A_					)
            snake_case__							:	List[Any]   	=	prepare_batch('val-batch.pt'					)
            with torch.no_grad():
                   snake_case__							:	Dict   	=	model.generate(
                       static_categorical_features=batch['static_categorical_features'], past_time_features=batch['past_time_features'], past_values=batch['past_values'], future_time_features=batch['future_time_features'], past_observed_mask=batch['past_observed_mask'], )
            snake_case__							:	int   	=	torch.Size((6_4, model.config.num_parallel_samples, model.config.prediction_length)					)
            self.assertEqual(outputs.sequences.shape, A_					)

            snake_case__							:	Optional[Any]   	=	torch.tensor([3_1_3_0.6_7_6_3, 4_0_5_6.5_2_9_3, 7_0_5_3.0_7_8_6], device=A_					)
            snake_case__							:	Union[str, Any]   	=	outputs.sequences.mean(dim=1					)
            self.assertTrue(torch.allclose(mean_prediction[0, -3:], A_, rtol=1e-1					)					)








import random
import unittest

import torch

from diffusers import IFInpaintingPipeline
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 				lowercase					(		_SCREAMING_SNAKE_CASE  ,   _SCREAMING_SNAKE_CASE  ,   unittest.TestCase   ):
		__lowercase       :       int	     =				IFInpaintingPipeline
		__lowercase       :       str	     =				TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"width", "height"}
		__lowercase       :       Optional[int]	     =				TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		__lowercase       :       Optional[int]	     =				PipelineTesterMixin.required_optional_params - {"latents"}



		def __UpperCamelCase ( self       )  ->					List[str]:
				"""simple docstring"""
				return self._get_dummy_components()



		def __UpperCamelCase ( self				,      A_				,      A_=0       )  ->					List[Any]:
				"""simple docstring"""
				if str(A_       ).startswith('mps'       ):
						UpperCamelCase		=     torch.manual_seed(A_       )
				else:
						UpperCamelCase		=     torch.Generator(device=A_       ).manual_seed(A_       )

				UpperCamelCase		=     floats_tensor((1, 3, 32, 32)				,      rng=random.Random(A_       )       ).to(A_       )
				UpperCamelCase		=     floats_tensor((1, 3, 32, 32)				,      rng=random.Random(A_       )       ).to(A_       )

				UpperCamelCase		=     {
				    'prompt': 'A painting of a squirrel eating a burger',
				    'image': 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 __UpperCamelCase ( self       )  ->					str:
				"""simple docstring"""
				self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3       )



		def __UpperCamelCase ( self       )  ->					str:
				"""simple docstring"""
				self._test_save_load_optional_components()



		@unittest.skipIf(torch_device != 'cuda'				,      reason='float16 requires CUDA'       )
		def __UpperCamelCase ( self       )  ->					str:
				"""simple docstring"""
				# 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 __UpperCamelCase ( self       )  ->					Optional[Any]:
				"""simple docstring"""
				self._test_attention_slicing_forward_pass(expected_max_diff=1e-2       )



		def __UpperCamelCase ( self       )  ->					Dict:
				"""simple docstring"""
				self._test_save_load_local()



		def __UpperCamelCase ( self       )  ->					Dict:
				"""simple docstring"""
				self._test_inference_batch_single_identical(
				    expected_max_diff=1e-2				,      )
