





'''simple docstring'''

import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__lowerCAmelCase							:	Union[str, Any]			=logging.get_logger(__name__)

enable_full_determinism()






class    UpperCAmelCase	(       __lowerCamelCase			,       __lowerCamelCase			,       unittest.TestCase				):
 __lowercase   		=				UNetaDModel
 __lowercase   		=				"""sample"""
 @property
 def      UpperCAmelCase_							(				self      :List[str]							)->							int:
     A__				=			4
     A__				=			3
     A__				=			(32, 32)

     A__				=			floats_tensor((batch_size, num_channels) + sizes							).to(lowercase_							)
     A__				=			torch.tensor([10]							).to(lowercase_							)

     return {"sample": noise, "timestep": time_step}
 @property
 def      UpperCAmelCase_							(				self      :int							)->							Union[str, Any]:
     return (3, 32, 32)
 @property
 def      UpperCAmelCase_							(				self      :str							)->							Dict:
     return (3, 32, 32)



 def      UpperCAmelCase_							(				self      :List[str]							)->							Dict:
     A__				=			{
         "block_out_channels": (32, 64),
         "down_block_types": ("DownBlock2D", "AttnDownBlock2D"),
         "up_block_types": ("AttnUpBlock2D", "UpBlock2D"),
         "attention_head_dim": 3,
         "out_channels": 3,
         "in_channels": 3,
         "layers_per_block": 2,
         "sample_size": 32,
     }
     A__				=			self.dummy_input
     return init_dict, inputs_dict






class    UpperCAmelCase	(       __lowerCamelCase			,       __lowerCamelCase			,       unittest.TestCase				):
 __lowercase   		=				UNetaDModel
 __lowercase   		=				"""sample"""
 @property
 def      UpperCAmelCase_							(				self      :int							)->							List[str]:
     A__				=			4
     A__				=			4
     A__				=			(32, 32)

     A__				=			floats_tensor((batch_size, num_channels) + sizes							).to(lowercase_							)
     A__				=			torch.tensor([10]							).to(lowercase_							)

     return {"sample": noise, "timestep": time_step}
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Optional[int]:
     return (4, 32, 32)
 @property
 def      UpperCAmelCase_							(				self      :Any							)->							int:
     return (4, 32, 32)
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							List[str]:
     A__				=			{
         "sample_size": 32,
         "in_channels": 4,
         "out_channels": 4,
         "layers_per_block": 2,
         "block_out_channels": (32, 64),
         "attention_head_dim": 32,
         "down_block_types": ("DownBlock2D", "DownBlock2D"),
         "up_block_types": ("UpBlock2D", "UpBlock2D"),
     }
     A__				=			self.dummy_input
     return init_dict, inputs_dict
 def      UpperCAmelCase_							(				self      :str							)->							List[str]:
     A__,	A__				=			UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,       output_loading_info=lowercase_							)

     self.assertIsNotNone(lowercase_							)
     self.assertEqual(len(loading_info["missing_keys"]							)						,       0							)

     model.to(lowercase_							)
     A__				=			model(**self.dummy_input							).sample

     assert image is not None, "Make sure output is not None"
 @unittest.skipIf(torch_device != "cuda"						,       "This test is supposed to run on GPU"							)
 def      UpperCAmelCase_							(				self      :int							)->							Union[str, Any]:
     A__,	A__				=			UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,       output_loading_info=lowercase_							)
     model.to(lowercase_							)
     A__				=			model(**self.dummy_input							).sample

     assert image is not None, "Make sure output is not None"
 @unittest.skipIf(torch_device != "cuda"						,       "This test is supposed to run on GPU"							)
 def      UpperCAmelCase_							(				self      :Any							)->							int:
     # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
     A__,	A__				=			UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,       output_loading_info=lowercase_							)
     model_accelerate.to(lowercase_							)
     model_accelerate.eval()

     A__				=			torch.randn(
         1						,       model_accelerate.config.in_channels						,       model_accelerate.config.sample_size						,       model_accelerate.config.sample_size						,       generator=torch.manual_seed(0							)						,       )
     A__				=			noise.to(lowercase_							)
     A__				=			torch.tensor([10] * noise.shape[0]							).to(lowercase_							)

     A__				=			model_accelerate(lowercase_						,       lowercase_							)["sample"]

     # two models don't need to stay in the device at the same time
     del model_accelerate
     torch.cuda.empty_cache()
     gc.collect()

     A__,	A__				=			UNetaDModel.from_pretrained(
         "fusing/unet-ldm-dummy-update"						,       output_loading_info=lowercase_						,       low_cpu_mem_usage=lowercase_							)
     model_normal_load.to(lowercase_							)
     model_normal_load.eval()
     A__				=			model_normal_load(lowercase_						,       lowercase_							)["sample"]

     assert torch_all_close(lowercase_						,       lowercase_						,       rtol=1E-3							)



 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Optional[Any]:
     A__				=			UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"							)
     model.eval()
     model.to(lowercase_							)

     A__				=			torch.randn(
         1						,       model.config.in_channels						,       model.config.sample_size						,       model.config.sample_size						,       generator=torch.manual_seed(0							)						,       )
     A__				=			noise.to(lowercase_							)
     A__				=			torch.tensor([10] * noise.shape[0]							).to(lowercase_							)

     with torch.no_grad():
         A__				=			model(lowercase_						,       lowercase_							).sample

     A__				=			output[0, -1, -3:, -3:].flatten().cpu()
     # fmt: off
     A__				=			torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]							)
     # fmt: on

     self.assertTrue(torch_all_close(lowercase_						,       lowercase_						,       rtol=1E-3							)							)






class    UpperCAmelCase	(       __lowerCamelCase			,       __lowerCamelCase			,       unittest.TestCase				):
 __lowercase   		=				UNetaDModel
 __lowercase   		=				"""sample"""
 @property
 def      UpperCAmelCase_							(				self      :Tuple						,       lowercase_      :List[Any]=(32, 32)							)->							int:
     A__				=			4
     A__				=			3

     A__				=			floats_tensor((batch_size, num_channels) + sizes							).to(lowercase_							)
     A__				=			torch.tensor(batch_size * [10]							).to(dtype=torch.intaa						,       device=lowercase_							)

     return {"sample": noise, "timestep": time_step}
 @property
 def      UpperCAmelCase_							(				self      :Tuple							)->							str:
     return (3, 32, 32)
 @property
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Optional[Any]:
     return (3, 32, 32)
 def      UpperCAmelCase_							(				self      :Any							)->							Union[str, Any]:
     A__				=			{
         "block_out_channels": [32, 64, 64, 64],
         "in_channels": 3,
         "layers_per_block": 1,
         "out_channels": 3,
         "time_embedding_type": "fourier",
         "norm_eps": 1E-6,
         "mid_block_scale_factor": math.sqrt(2.0							),
         "norm_num_groups": None,
         "down_block_types": [
             "SkipDownBlock2D",
             "AttnSkipDownBlock2D",
             "SkipDownBlock2D",
             "SkipDownBlock2D",
         ],
         "up_block_types": [
             "SkipUpBlock2D",
             "SkipUpBlock2D",
             "AttnSkipUpBlock2D",
             "SkipUpBlock2D",
         ],
     }
     A__				=			self.dummy_input
     return init_dict, inputs_dict
 @slow
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Optional[Any]:
     A__,	A__				=			UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"						,       output_loading_info=lowercase_							)
     self.assertIsNotNone(lowercase_							)
     self.assertEqual(len(loading_info["missing_keys"]							)						,       0							)

     model.to(lowercase_							)
     A__				=			self.dummy_input
     A__				=			floats_tensor((4, 3) + (2_56, 2_56)							).to(lowercase_							)
     A__				=			noise
     A__				=			model(**lowercase_							)

     assert image is not None, "Make sure output is not None"
 @slow
 def      UpperCAmelCase_							(				self      :Tuple							)->							str:
     A__				=			UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"							)
     model.to(lowercase_							)

     A__				=			4
     A__				=			3
     A__				=			(2_56, 2_56)

     A__				=			torch.ones((batch_size, num_channels) + sizes							).to(lowercase_							)
     A__				=			torch.tensor(batch_size * [1E-4]							).to(lowercase_							)

     with torch.no_grad():
         A__				=			model(lowercase_						,       lowercase_							).sample

     A__				=			output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     A__				=			torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_6_7.7_6_0_8]							)
     # fmt: on

     self.assertTrue(torch_all_close(lowercase_						,       lowercase_						,       rtol=1E-2							)							)
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Tuple:
     A__				=			UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update"							)
     model.to(lowercase_							)

     A__				=			4
     A__				=			3
     A__				=			(32, 32)

     A__				=			torch.ones((batch_size, num_channels) + sizes							).to(lowercase_							)
     A__				=			torch.tensor(batch_size * [1E-4]							).to(lowercase_							)

     with torch.no_grad():
         A__				=			model(lowercase_						,       lowercase_							).sample

     A__				=			output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     A__				=			torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]							)
     # fmt: on

     self.assertTrue(torch_all_close(lowercase_						,       lowercase_						,       rtol=1E-2							)							)



 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Optional[Any]:
     # not required for this model
     pass





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       ):
		lowercase       =    []





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_init_end'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_train_begin'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_train_end'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_epoch_begin'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_epoch_end'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_step_begin'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_step_end'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_evaluate'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_predict'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_save'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_log'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       ):
		self.events.append('on_prediction_step'       )







@require_torch
class        A_   (			unittest.TestCase   ):




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    tempfile.mkdtemp()





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		shutil.rmtree(self.output_dir       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case=0		,    snake_case=0		,    snake_case=64		,    snake_case=64		,    snake_case=None		,    snake_case=False		,    **snake_case       ):
		# 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.
		lowercase       =    RegressionDataset(length=snake_case       )
		lowercase       =    RegressionDataset(length=snake_case       )
		lowercase       =    RegressionModelConfig(a=snake_case		,    b=snake_case       )
		lowercase       =    RegressionPreTrainedModel(snake_case       )

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





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		self.assertEqual(len(snake_case       )		,    len(snake_case       )       )

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

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





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		lowercase       =    ['on_init_end', 'on_train_begin']
		lowercase       =    0
		lowercase       =    len(trainer.get_eval_dataloader()       )
		lowercase       =    ['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(snake_case       ):
				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 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.get_trainer()
		lowercase       =    DEFAULT_CALLBACKS.copy() + [ProgressCallback]
		self.check_callbacks_equality(trainer.callback_handler.callbacks		,    snake_case       )

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

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





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    DEFAULT_CALLBACKS.copy() + [ProgressCallback]
		lowercase       =    self.get_trainer()

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

		lowercase       =    self.get_trainer()
		lowercase       =    trainer.pop_callback(snake_case       )
		self.assertEqual(cb.__class__		,    snake_case       )
		self.check_callbacks_equality(trainer.callback_handler.callbacks		,    snake_case       )

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

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

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

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





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		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=snake_case       )

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

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

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

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

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

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

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

