



'''simple docstring'''

import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

class 							a__( snake_case__   ,      unittest.TestCase     ):
			a_	: Tuple									=      KandinskyVaaPriorPipeline
			a_	: List[Any]									=      ['''prompt''']
			a_	: Dict									=      ['''prompt''', '''negative_prompt''']
			a_	: Optional[int]									=      [
			    '''num_images_per_prompt''',
			    '''generator''',
			    '''num_inference_steps''',
			    '''latents''',
			    '''negative_prompt''',
			    '''guidance_scale''',
			    '''output_type''',
			    '''return_dict''',
			]
			a_	: Optional[Any]									=      False






			@property
			def 						_lowercase  (     self						)							->							str:
					return 32






			@property
			def 						_lowercase  (     self						)							->							int:
					return 32






			@property
			def 						_lowercase  (     self						)							->							List[Any]:
					return self.time_input_dim






			@property
			def 						_lowercase  (     self						)							->							str:
					return self.time_input_dim * 4






			@property
			def 						_lowercase  (     self						)							->							Dict:
					return 100






			@property
			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__	     =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)
					return tokenizer






			@property
			def 						_lowercase  (     self						)							->							int:
					torch.manual_seed(0						)
					snake_case__	     =CLIPTextConfig(
					    bos_token_id=0					,   eos_token_id=2					,   hidden_size=self.text_embedder_hidden_size					,   projection_dim=self.text_embedder_hidden_size					,   intermediate_size=37					,   layer_norm_eps=1E-05					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1000					,   )
					return CLIPTextModelWithProjection(_UpperCAmelCase						)






			@property
			def 						_lowercase  (     self						)							->							List[str]:
					torch.manual_seed(0						)

					snake_case__	     ={
					    'num_attention_heads': 2,
					    'attention_head_dim': 12,
					    'embedding_dim': self.text_embedder_hidden_size,
					    'num_layers': 1,
					}

					snake_case__	     =PriorTransformer(**_UpperCAmelCase						)
					# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
					snake_case__	     =nn.Parameter(torch.ones(model.clip_std.shape						)						)
					return model






			@property
			def 						_lowercase  (     self						)							->							List[str]:
					torch.manual_seed(0						)
					snake_case__	     =CLIPVisionConfig(
					    hidden_size=self.text_embedder_hidden_size					,   image_size=224					,   projection_dim=self.text_embedder_hidden_size					,   intermediate_size=37					,   num_attention_heads=4					,   num_channels=3					,   num_hidden_layers=5					,   patch_size=14					,   )

					snake_case__	     =CLIPVisionModelWithProjection(_UpperCAmelCase						)
					return model






			@property
			def 						_lowercase  (     self						)							->							List[str]:
					snake_case__	     =CLIPImageProcessor(
					    crop_size=224					,   do_center_crop=_UpperCAmelCase					,   do_normalize=_UpperCAmelCase					,   do_resize=_UpperCAmelCase					,   image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]					,   image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]					,   resample=3					,   size=224					,   )

					return image_processor






			def 						_lowercase  (     self						)							->							List[str]:
					snake_case__	     =self.dummy_prior
					snake_case__	     =self.dummy_image_encoder
					snake_case__	     =self.dummy_text_encoder
					snake_case__	     =self.dummy_tokenizer
					snake_case__	     =self.dummy_image_processor

					snake_case__	     =UnCLIPScheduler(
					    variance_type='fixed_small_log'					,   prediction_type='sample'					,   num_train_timesteps=1000					,   clip_sample=_UpperCAmelCase					,   clip_sample_range=10.0					,   )

					snake_case__	     ={
					    'prior': prior,
					    'image_encoder': image_encoder,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'scheduler': scheduler,
					    'image_processor': image_processor,
					}

					return components






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase=0						)							->							List[str]:
					if str(_UpperCAmelCase						).startswith('mps'						):
							snake_case__	     =torch.manual_seed(_UpperCAmelCase						)
					else:
							snake_case__	     =torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)
					snake_case__	     ={
					    'prompt': 'horse',
					    'generator': generator,
					    'guidance_scale': 4.0,
					    'num_inference_steps': 2,
					    'output_type': 'np',
					}
					return inputs






			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__	     ='cpu'

					snake_case__	     =self.get_dummy_components()

					snake_case__	     =self.pipeline_class(**_UpperCAmelCase						)
					snake_case__	     =pipe.to(_UpperCAmelCase						)

					pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

					snake_case__	     =pipe(**self.get_dummy_inputs(_UpperCAmelCase						)						)
					snake_case__	     =output.image_embeds

					snake_case__	     =pipe(
					    **self.get_dummy_inputs(_UpperCAmelCase						)					,   return_dict=_UpperCAmelCase					,   )[0]

					snake_case__	     =image[0, -10:]
					snake_case__	     =image_from_tuple[0, -10:]

					assert image.shape == (1, 32)

					snake_case__	     =np.array(
					    [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]						)

					assert np.abs(image_slice.flatten() - expected_slice						).max() < 1E-2
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < 1E-2






			@skip_mps
			def 						_lowercase  (     self						)							->							Union[str, Any]:
					snake_case__	     =torch_device == 'cpu'
					snake_case__	     =True
					snake_case__	     =False

					self._test_inference_batch_single_identical(
					    test_max_difference=_UpperCAmelCase					,   relax_max_difference=_UpperCAmelCase					,   test_mean_pixel_difference=_UpperCAmelCase					,   )






			@skip_mps
			def 						_lowercase  (     self						)							->							int:
					snake_case__	     =torch_device == 'cpu'
					snake_case__	     =False

					self._test_attention_slicing_forward_pass(
					    test_max_difference=_UpperCAmelCase					,   test_mean_pixel_difference=_UpperCAmelCase					,   )









'''simple docstring'''

import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
			from transformers.trainer import DEFAULT_CALLBACKS

			from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel

class 							a__( snake_case__     ):






			def __init__(     self						)							->							Dict:
					snake_case__	     =[]






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							List[str]:
					self.events.append('on_init_end'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Dict:
					self.events.append('on_train_begin'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Optional[Any]:
					self.events.append('on_train_end'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							int:
					self.events.append('on_epoch_begin'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							List[str]:
					self.events.append('on_epoch_end'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Optional[int]:
					self.events.append('on_step_begin'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							str:
					self.events.append('on_step_end'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							List[Any]:
					self.events.append('on_evaluate'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Dict:
					self.events.append('on_predict'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Any:
					self.events.append('on_save'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							Tuple:
					self.events.append('on_log'						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   **_UpperCAmelCase						)							->							List[Any]:
					self.events.append('on_prediction_step'						)





@require_torch
class 							a__( unittest.TestCase     ):






			def 						_lowercase  (     self						)							->							Optional[Any]:
					snake_case__	     =tempfile.mkdtemp()






			def 						_lowercase  (     self						)							->							List[str]:
					shutil.rmtree(self.output_dir						)






			def 						_lowercase  (     self					,   _UpperCAmelCase=0					,   _UpperCAmelCase=0					,   _UpperCAmelCase=64					,   _UpperCAmelCase=64					,   _UpperCAmelCase=None					,   _UpperCAmelCase=False					,   **_UpperCAmelCase						)							->							int:
					# disable_tqdm in TrainingArguments has a flaky default since it depends on the level of logging. We make sure
					# its set to False since the tests later on depend on its value.
					snake_case__	     =RegressionDataset(length=_UpperCAmelCase						)
					snake_case__	     =RegressionDataset(length=_UpperCAmelCase						)
					snake_case__	     =RegressionModelConfig(a=_UpperCAmelCase					,   b=_UpperCAmelCase						)
					snake_case__	     =RegressionPreTrainedModel(_UpperCAmelCase						)

					snake_case__	     =TrainingArguments(self.output_dir					,   disable_tqdm=_UpperCAmelCase					,   report_to=[]					,   **_UpperCAmelCase						)
					return Trainer(
					    _UpperCAmelCase					,   _UpperCAmelCase					,   train_dataset=_UpperCAmelCase					,   eval_dataset=_UpperCAmelCase					,   callbacks=_UpperCAmelCase					,   )






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Optional[Any]:
					self.assertEqual(len(_UpperCAmelCase						)					,   len(_UpperCAmelCase						)						)

					# Order doesn't matter
					snake_case__	     =sorted(_UpperCAmelCase					,   key=lambda _UpperCAmelCase						: cb.__name__ if isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) else cb.__class__.__name__						)
					snake_case__	     =sorted(_UpperCAmelCase					,   key=lambda _UpperCAmelCase						: cb.__name__ if isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) else cb.__class__.__name__						)

					for cba, cba in zip(_UpperCAmelCase					,   _UpperCAmelCase						):
							if isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) and isinstance(_UpperCAmelCase					,   _UpperCAmelCase						):
									self.assertEqual(_UpperCAmelCase					,   _UpperCAmelCase						)
							elif isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) and not isinstance(_UpperCAmelCase					,   _UpperCAmelCase						):
									self.assertEqual(_UpperCAmelCase					,   cba.__class__						)
							elif not isinstance(_UpperCAmelCase					,   _UpperCAmelCase						) and isinstance(_UpperCAmelCase					,   _UpperCAmelCase						):
									self.assertEqual(cba.__class__					,   _UpperCAmelCase						)
							else:
									self.assertEqual(_UpperCAmelCase					,   _UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Any:
					snake_case__	     =['on_init_end', 'on_train_begin']
					snake_case__	     =0
					snake_case__	     =len(trainer.get_eval_dataloader()						)
					snake_case__	     =['on_prediction_step'] * len(trainer.get_eval_dataloader()						) + ['on_log', 'on_evaluate']
					for _ in range(trainer.state.num_train_epochs						):
							expected_events.append('on_epoch_begin'						)
							for _ in range(_UpperCAmelCase						):
									step += 1
									expected_events += ["on_step_begin", "on_step_end"]
									if step % trainer.args.logging_steps == 0:
											expected_events.append('on_log'						)
									if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
											expected_events += evaluation_events.copy()
									if step % trainer.args.save_steps == 0:
											expected_events.append('on_save'						)
							expected_events.append('on_epoch_end'						)
							if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
									expected_events += evaluation_events.copy()
					expected_events += ["on_log", "on_train_end"]
					return expected_events






			def 						_lowercase  (     self						)							->							Optional[int]:
					snake_case__	     =self.get_trainer()
					snake_case__	     =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					# Callbacks passed at init are added to the default callbacks
					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]						)
					expected_callbacks.append(_UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
					snake_case__	     =self.get_trainer(disable_tqdm=_UpperCAmelCase						)
					snake_case__	     =DEFAULT_CALLBACKS.copy() + [PrinterCallback]
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							Optional[Any]:
					snake_case__	     =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
					snake_case__	     =self.get_trainer()

					# We can add, pop, or remove by class name
					trainer.remove_callback(_UpperCAmelCase						)
					expected_callbacks.remove(_UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					snake_case__	     =self.get_trainer()
					snake_case__	     =trainer.pop_callback(_UpperCAmelCase						)
					self.assertEqual(cb.__class__					,   _UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					trainer.add_callback(_UpperCAmelCase						)
					expected_callbacks.insert(0					,   _UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					# We can also add, pop, or remove by instance
					snake_case__	     =self.get_trainer()
					snake_case__	     =trainer.callback_handler.callbacks[0]
					trainer.remove_callback(_UpperCAmelCase						)
					expected_callbacks.remove(_UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					snake_case__	     =self.get_trainer()
					snake_case__	     =trainer.callback_handler.callbacks[0]
					snake_case__	     =trainer.pop_callback(_UpperCAmelCase						)
					self.assertEqual(_UpperCAmelCase					,   _UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)

					trainer.add_callback(_UpperCAmelCase						)
					expected_callbacks.insert(0					,   _UpperCAmelCase						)
					self.check_callbacks_equality(trainer.callback_handler.callbacks					,   _UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							Dict:
					import warnings

					# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
					warnings.simplefilter(action='ignore'					,   category=_UpperCAmelCase						)

					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]						)
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					# Independent log/save/eval
					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]					,   logging_steps=5						)
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]					,   save_steps=5						)
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]					,   eval_steps=5					,   evaluation_strategy='steps'						)
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					snake_case__	     =self.get_trainer(callbacks=[MyTestTrainerCallback]					,   evaluation_strategy='epoch'						)
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					# A bit of everything
					snake_case__	     =self.get_trainer(
					    callbacks=[MyTestTrainerCallback]					,   logging_steps=3					,   save_steps=10					,   eval_steps=5					,   evaluation_strategy='steps'					,   )
					trainer.train()
					snake_case__	     =trainer.callback_handler.callbacks[-2].events
					self.assertEqual(_UpperCAmelCase					,   self.get_expected_events(_UpperCAmelCase						)						)

					# warning should be emitted for duplicated callbacks
					with patch('transformers.trainer_callback.logger.warning'						) as warn_mock:
							snake_case__	     =self.get_trainer(
							    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]					,   )
							assert str(_UpperCAmelCase						) in warn_mock.call_args[0][0]




