





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    (    lowerCAmelCase ):

			'''simple docstring'''







			def __init__(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		[]





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_init_end'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_train_begin'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_train_end'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_epoch_begin'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_epoch_end'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_step_begin'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_step_end'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_evaluate'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_predict'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_save'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_log'''							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				self.events.append('''on_prediction_step'''							)






@require_torch
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		tempfile.mkdtemp()





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				shutil.rmtree(self.output_dir							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase=0       ,  __lowerCAmelCase=0       ,  __lowerCAmelCase=6_4       ,  __lowerCAmelCase=6_4       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=False       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		RegressionDataset(length=__lowerCAmelCase							)
				lowerCamelCase__								=		RegressionDataset(length=__lowerCAmelCase							)
				lowerCamelCase__								=		RegressionModelConfig(a=__lowerCAmelCase       ,  b=__lowerCAmelCase							)
				lowerCamelCase__								=		RegressionPreTrainedModel(__lowerCAmelCase							)

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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase							):



				'''simple docstring'''







				self.assertEqual(len(__lowerCAmelCase							)       ,  len(__lowerCAmelCase							)							)

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

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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		['''on_init_end''', '''on_train_begin''']
				lowerCamelCase__								=		0
				lowerCamelCase__								=		len(trainer.get_eval_dataloader()							)
				lowerCamelCase__								=		['''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(__lowerCAmelCase							):
						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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.get_trainer()
				lowerCamelCase__								=		DEFAULT_CALLBACKS.copy() + [ProgressCallback]
				self.check_callbacks_equality(trainer.callback_handler.callbacks       ,  __lowerCAmelCase							)

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

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		DEFAULT_CALLBACKS.copy() + [ProgressCallback]
				lowerCamelCase__								=		self.get_trainer()

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

				lowerCamelCase__								=		self.get_trainer()
				lowerCamelCase__								=		trainer.pop_callback(__lowerCAmelCase							)
				self.assertEqual(cb.__class__       ,  __lowerCAmelCase							)
				self.check_callbacks_equality(trainer.callback_handler.callbacks       ,  __lowerCAmelCase							)

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

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

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

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				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=__lowerCAmelCase							)

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

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

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

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

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

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

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










import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						(UnCLIPScheduler,)





			def 		__lowerCamelCase				(    self       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		{
				    '''num_train_timesteps''': 1_0_0_0,
				    '''variance_type''': '''fixed_small_log''',
				    '''clip_sample''': True,
				    '''clip_sample_range''': 1.0,
				    '''prediction_type''': '''epsilon''',
				}

				config.update(**__lowerCAmelCase							)
				return config





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for timesteps in [1, 5, 1_0_0, 1_0_0_0]:
					self.check_over_configs(num_train_timesteps=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for variance in ["fixed_small_log", "learned_range"]:
					self.check_over_configs(variance_type=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for clip_sample in [True, False]:
					self.check_over_configs(clip_sample=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for clip_sample_range in [1, 5, 1_0, 2_0]:
					self.check_over_configs(clip_sample_range=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for prediction_type in ["epsilon", "sample"]:
					self.check_over_configs(prediction_type=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				for time_step in [0, 5_0_0, 9_9_9]:
					for prev_timestep in [None, 5, 1_0_0, 2_5_0, 5_0_0, 7_5_0]:
						if prev_timestep is not None and prev_timestep >= time_step:
							continue

						self.check_over_forward(time_step=__lowerCAmelCase       ,  prev_timestep=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config(variance_type='''fixed_small_log'''							)
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				assert torch.sum(torch.abs(scheduler._get_variance(0							) - 1.0000E-10							)							) < 1E-5
				assert torch.sum(torch.abs(scheduler._get_variance(4_8_7							) - 0.054_9625							)							) < 1E-5
				assert torch.sum(torch.abs(scheduler._get_variance(9_9_9							) - 0.999_4987							)							) < 1E-5





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config(variance_type='''learned_range'''							)
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				lowerCamelCase__								=		0.5

				assert scheduler._get_variance(1       ,  predicted_variance=__lowerCAmelCase							) - -10.171_2790 < 1E-5
				assert scheduler._get_variance(4_8_7       ,  predicted_variance=__lowerCAmelCase							) - -5.799_8052 < 1E-5
				assert scheduler._get_variance(9_9_9       ,  predicted_variance=__lowerCAmelCase							) - -0.001_0011 < 1E-5





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config()
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				lowerCamelCase__								=		scheduler.timesteps

				lowerCamelCase__								=		self.dummy_model()
				lowerCamelCase__								=		self.dummy_sample_deter
				lowerCamelCase__								=		torch.manual_seed(0							)

				for i, t in enumerate(__lowerCAmelCase							):
					# 1. predict noise residual
					lowerCamelCase__								=		model(__lowerCAmelCase       ,  __lowerCAmelCase							)

					# 2. predict previous mean of sample x_t-1
					lowerCamelCase__								=		scheduler.step(__lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  generator=__lowerCAmelCase							).prev_sample

					lowerCamelCase__								=		pred_prev_sample

				lowerCamelCase__								=		torch.sum(torch.abs(__lowerCAmelCase							)							)
				lowerCamelCase__								=		torch.mean(torch.abs(__lowerCAmelCase							)							)

				assert abs(result_sum.item() - 252.268_2495							) < 1E-2
				assert abs(result_mean.item() - 0.328_4743							) < 1E-3





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.scheduler_classes[0]
				lowerCamelCase__								=		self.get_scheduler_config()
				lowerCamelCase__								=		scheduler_class(**__lowerCAmelCase							)

				scheduler.set_timesteps(2_5							)

				lowerCamelCase__								=		scheduler.timesteps

				lowerCamelCase__								=		self.dummy_model()
				lowerCamelCase__								=		self.dummy_sample_deter
				lowerCamelCase__								=		torch.manual_seed(0							)

				for i, t in enumerate(__lowerCAmelCase							):
					# 1. predict noise residual
					lowerCamelCase__								=		model(__lowerCAmelCase       ,  __lowerCAmelCase							)

					if i + 1 == timesteps.shape[0]:
						lowerCamelCase__								=		None
					else:
						lowerCamelCase__								=		timesteps[i + 1]

					# 2. predict previous mean of sample x_t-1
					lowerCamelCase__								=		scheduler.step(
					    __lowerCAmelCase       ,  __lowerCAmelCase       ,  __lowerCAmelCase       ,  prev_timestep=__lowerCAmelCase       ,  generator=__lowerCAmelCase							).prev_sample

					lowerCamelCase__								=		pred_prev_sample

				lowerCamelCase__								=		torch.sum(torch.abs(__lowerCAmelCase							)							)
				lowerCamelCase__								=		torch.mean(torch.abs(__lowerCAmelCase							)							)

				assert abs(result_sum.item() - 258.204_4983							) < 1E-2
				assert abs(result_mean.item() - 0.336_2038							) < 1E-3





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				pass





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				pass



