


'''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__    	=     logging.get_logger(__name__)

enable_full_determinism()






class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	Tuple          =     UNetaDModel
			lowerCAmelCase :	Optional[int]          =     "sample"




			@property
			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      4
					_UpperCAmelCase		:       Any								=      3
					_UpperCAmelCase		:       Tuple								=      (32, 32)

					_UpperCAmelCase		:       Dict								=      floats_tensor((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      torch.tensor([10] ).to(lowerCamelCase__ )

					return {"sample": noise, "timestep": time_step}




			@property
			def 							lowerCAmelCase__    (		self  :     Dict )   ->Dict:



					'''simple docstring'''







					return (3, 32, 32)




			@property
			def 							lowerCAmelCase__    (		self  :     str )   ->Optional[Any]:



					'''simple docstring'''







					return (3, 32, 32)






			def 							lowerCAmelCase__    (		self  :     str )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       str								=      {
					    "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,
					}
					_UpperCAmelCase		:       List[str]								=      self.dummy_input
					return init_dict, inputs_dict






class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	Dict          =     UNetaDModel
			lowerCAmelCase :	Any          =     "sample"




			@property
			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:       Tuple								=      4
					_UpperCAmelCase		:       Optional[int]								=      4
					_UpperCAmelCase		:       Any								=      (32, 32)

					_UpperCAmelCase		:       Dict								=      floats_tensor((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      torch.tensor([10] ).to(lowerCamelCase__ )

					return {"sample": noise, "timestep": time_step}




			@property
			def 							lowerCAmelCase__    (		self  :     List[str] )   ->Optional[Any]:



					'''simple docstring'''







					return (4, 32, 32)




			@property
			def 							lowerCAmelCase__    (		self  :     List[Any] )   ->int:



					'''simple docstring'''







					return (4, 32, 32)




			def 							lowerCAmelCase__    (		self  :     Any )   ->int:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[int]								=      {
					    "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"),
					}
					_UpperCAmelCase		:       List[Any]								=      self.dummy_input
					return init_dict, inputs_dict




			def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"	,  output_loading_info=lowerCamelCase__ )

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

					model.to(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[int]								=      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 							lowerCAmelCase__    (		self  :     Any )   ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Dict								=      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"	,  output_loading_info=lowerCamelCase__ )
					model.to(lowerCamelCase__ )
					_UpperCAmelCase		:       str								=      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 							lowerCAmelCase__    (		self  :     List[str] )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"	,  output_loading_info=lowerCamelCase__ )
					model_accelerate.to(lowerCamelCase__ )
					model_accelerate.eval()

					_UpperCAmelCase		:       List[Any]								=      torch.randn(
					    1	,  model_accelerate.config.in_channels	,  model_accelerate.config.sample_size	,  model_accelerate.config.sample_size	,  generator=torch.manual_seed(0 )	,  )
					_UpperCAmelCase		:       Tuple								=      noise.to(lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      torch.tensor([10] * noise.shape[0] ).to(lowerCamelCase__ )

					_UpperCAmelCase		:       Any								=      model_accelerate(lowerCamelCase__	,  lowerCamelCase__ )["sample"]

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

					_UpperCAmelCase					,    _UpperCAmelCase		:       int								=      UNetaDModel.from_pretrained(
					    "fusing/unet-ldm-dummy-update"	,  output_loading_info=lowerCamelCase__	,  low_cpu_mem_usage=lowerCamelCase__ )
					model_normal_load.to(lowerCamelCase__ )
					model_normal_load.eval()
					_UpperCAmelCase		:       List[str]								=      model_normal_load(lowerCamelCase__	,  lowerCamelCase__ )["sample"]

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






			def 							lowerCAmelCase__    (		self  :     int )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       Union[str, Any]								=      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update" )
					model.eval()
					model.to(lowerCamelCase__ )

					_UpperCAmelCase		:       List[str]								=      torch.randn(
					    1	,  model.config.in_channels	,  model.config.sample_size	,  model.config.sample_size	,  generator=torch.manual_seed(0 )	,  )
					_UpperCAmelCase		:       List[Any]								=      noise.to(lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      torch.tensor([10] * noise.shape[0] ).to(lowerCamelCase__ )

					with torch.no_grad():
							_UpperCAmelCase		:       Any								=      model(lowerCamelCase__	,  lowerCamelCase__ ).sample

					_UpperCAmelCase		:       List[str]								=      output[0, -1, -3:, -3:].flatten().cpu()
					# fmt: off
					_UpperCAmelCase		:       Tuple								=      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(lowerCamelCase__	,  lowerCamelCase__	,  rtol=1E-3 ) )






class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	Union[str, Any]          =     UNetaDModel
			lowerCAmelCase :	Optional[int]          =     "sample"




			@property
			def 							lowerCAmelCase__    (		self  :     Tuple	,  lowerCamelCase__  :     Dict=(32, 32) )   ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Dict								=      4
					_UpperCAmelCase		:       Union[str, Any]								=      3

					_UpperCAmelCase		:       List[Any]								=      floats_tensor((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       List[str]								=      torch.tensor(batch_size * [10] ).to(dtype=torch.intaa	,  device=lowerCamelCase__ )

					return {"sample": noise, "timestep": time_step}




			@property
			def 							lowerCAmelCase__    (		self  :     List[str] )   ->str:



					'''simple docstring'''







					return (3, 32, 32)




			@property
			def 							lowerCAmelCase__    (		self  :     Any )   ->Dict:



					'''simple docstring'''







					return (3, 32, 32)




			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      {
					    "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",
					    ],
					}
					_UpperCAmelCase		:       Optional[int]								=      self.dummy_input
					return init_dict, inputs_dict




			@slow
			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"	,  output_loading_info=lowerCamelCase__ )
					self.assertIsNotNone(lowerCamelCase__ )
					self.assertEqual(len(loading_info["missing_keys"] )	,  0 )

					model.to(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      self.dummy_input
					_UpperCAmelCase		:       Union[str, Any]								=      floats_tensor((4, 3) + (2_56, 2_56) ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       int								=      noise
					_UpperCAmelCase		:       Dict								=      model(**lowerCamelCase__ )

					assert image is not None, "Make sure output is not None"




			@slow
			def 							lowerCAmelCase__    (		self  :     str )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256" )
					model.to(lowerCamelCase__ )

					_UpperCAmelCase		:       List[str]								=      4
					_UpperCAmelCase		:       int								=      3
					_UpperCAmelCase		:       Optional[int]								=      (2_56, 2_56)

					_UpperCAmelCase		:       Optional[int]								=      torch.ones((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       int								=      torch.tensor(batch_size * [1E-4] ).to(lowerCamelCase__ )

					with torch.no_grad():
							_UpperCAmelCase		:       Tuple								=      model(lowerCamelCase__	,  lowerCamelCase__ ).sample

					_UpperCAmelCase		:       List[Any]								=      output[0, -3:, -3:, -1].flatten().cpu()
					# fmt: off
					_UpperCAmelCase		:       Union[str, Any]								=      torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8] )
					# fmt: on

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




			def 							lowerCAmelCase__    (		self  :     List[str] )   ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update" )
					model.to(lowerCamelCase__ )

					_UpperCAmelCase		:       Tuple								=      4
					_UpperCAmelCase		:       List[Any]								=      3
					_UpperCAmelCase		:       Any								=      (32, 32)

					_UpperCAmelCase		:       Optional[int]								=      torch.ones((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )
					_UpperCAmelCase		:       int								=      torch.tensor(batch_size * [1E-4] ).to(lowerCamelCase__ )

					with torch.no_grad():
							_UpperCAmelCase		:       Any								=      model(lowerCamelCase__	,  lowerCamelCase__ ).sample

					_UpperCAmelCase		:       Optional[Any]								=      output[0, -3:, -3:, -1].flatten().cpu()
					# fmt: off
					_UpperCAmelCase		:       Dict								=      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(lowerCamelCase__	,  lowerCamelCase__	,  rtol=1E-2 ) )






			def 							lowerCAmelCase__    (		self  :     Dict )   ->Optional[Any]:



					'''simple docstring'''







					pass







'''simple docstring'''


# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
lowerCamelCase__    	=     abspath(join(dirname(dirname(dirname(__file__))), 'src'))
sys.path.insert(1, git_repo_path)


# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='ignore', category=FutureWarning)






def  __lowerCAmelCase (__lowerCAmelCase   ):
		from transformers.testing_utils import pytest_addoption_shared

		pytest_addoption_shared(__lowerCAmelCase   )






def  __lowerCAmelCase (__lowerCAmelCase   ):
		from transformers.testing_utils import pytest_terminal_summary_main

		_UpperCAmelCase		:       Optional[int]								=      terminalreporter.config.getoption("--make-reports"   )
		if make_reports:
				pytest_terminal_summary_main(__lowerCAmelCase      , id=__lowerCAmelCase   )



