


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__  :List[Any]             =							logging.get_logger(__name__)

enable_full_determinism()

class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  Union[str, Any]           =     UNetaDModel
		_a     :  List[Any]           =     'sample'
		@property
		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      3
									_UpperCAmelCase					=      (32, 32)

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

									return {"sample": noise, "timestep": time_step}
		@property
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									return (3, 32, 32)
		@property
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									return (3, 32, 32)






		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      {
									    '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					=      self.dummy_input
									return init_dict, inputs_dict



class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  Optional[Any]           =     UNetaDModel
		_a     :  str           =     'sample'
		@property
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      (32, 32)

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

									return {"sample": noise, "timestep": time_step}
		@property
		def      UpperCAmelCase__    (      self						)			->	Dict:




									"""simple docstring"""





									return (4, 32, 32)
		@property
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									return (4, 32, 32)
		def      UpperCAmelCase__    (      self						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      {
									    '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					=      self.dummy_input
									return init_dict, inputs_dict
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'      ,   output_loading_info=_SCREAMING_SNAKE_CASE						)

									self.assertIsNotNone(_SCREAMING_SNAKE_CASE						)
									self.assertEqual(len(loading_info['missing_keys']						)      ,   0						)

									model.to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      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						)			->	Tuple:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'      ,   output_loading_info=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      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						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase		,	_UpperCAmelCase					=      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'      ,   output_loading_info=_SCREAMING_SNAKE_CASE						)
									model_accelerate.to(_SCREAMING_SNAKE_CASE						)
									model_accelerate.eval()

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

									_UpperCAmelCase					=      model_accelerate(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)['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					=      UNetaDModel.from_pretrained(
									    'fusing/unet-ldm-dummy-update'      ,   output_loading_info=_SCREAMING_SNAKE_CASE      ,   low_cpu_mem_usage=_SCREAMING_SNAKE_CASE						)
									model_normal_load.to(_SCREAMING_SNAKE_CASE						)
									model_normal_load.eval()
									_UpperCAmelCase					=      model_normal_load(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)['sample']

									assert torch_all_close(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rtol=1e-3						)






		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'						)
									model.eval()
									model.to(_SCREAMING_SNAKE_CASE						)

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

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

									_UpperCAmelCase					=      output[0, -1, -3:, -3:].flatten().cpu()
									# fmt: off
									_UpperCAmelCase					=      torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]						)
									# fmt: on

									self.assertTrue(torch_all_close(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rtol=1e-3						)						)



class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  str           =     UNetaDModel
		_a     :  Any           =     'sample'
		@property
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE=(32, 32)						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      3

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

									return {"sample": noise, "timestep": time_step}
		@property
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									return (3, 32, 32)
		@property
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									return (3, 32, 32)
		def      UpperCAmelCase__    (      self						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      {
									    '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					=      self.dummy_input
									return init_dict, inputs_dict
		@slow
		def      UpperCAmelCase__    (      self						)			->	Optional[int]:




									"""simple docstring"""





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

									model.to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.dummy_input
									_UpperCAmelCase					=      floats_tensor((4, 3) + (256, 256)						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      noise
									_UpperCAmelCase					=      model(**_SCREAMING_SNAKE_CASE						)

									assert image is not None, "Make sure output is not None"
		@slow
		def      UpperCAmelCase__    (      self						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'						)
									model.to(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      3
									_UpperCAmelCase					=      (256, 256)

									_UpperCAmelCase					=      torch.ones((batch_size, num_channels) + sizes						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      torch.tensor(batch_size * [1e-4]						).to(_SCREAMING_SNAKE_CASE						)

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

									_UpperCAmelCase					=      output[0, -3:, -3:, -1].flatten().cpu()
									# fmt: off
									_UpperCAmelCase					=      torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -10980.7129, -20028.8535, 8148.2822, 2342.2905, 567.7608]						)
									# fmt: on

									self.assertTrue(torch_all_close(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rtol=1e-2						)						)
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'						)
									model.to(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      4
									_UpperCAmelCase					=      3
									_UpperCAmelCase					=      (32, 32)

									_UpperCAmelCase					=      torch.ones((batch_size, num_channels) + sizes						).to(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      torch.tensor(batch_size * [1e-4]						).to(_SCREAMING_SNAKE_CASE						)

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

									_UpperCAmelCase					=      output[0, -3:, -3:, -1].flatten().cpu()
									# fmt: off
									_UpperCAmelCase					=      torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]						)
									# fmt: on

									self.assertTrue(torch_all_close(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rtol=1e-2						)						)






		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									pass




from importlib import import_module

from .logging import get_logger


lowerCAmelCase__  :Optional[Any]             =							get_logger(__name__)

class 		__a				:
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=None						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      attrs or []
									if module is not None:
																for key in module.__dict__:
																							if key in attrs or not key.startswith('__'						):
																														setattr(self      ,   _SCREAMING_SNAKE_CASE      ,   getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)						)
									_UpperCAmelCase					=      module._original_module if isinstance(_SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj						) else module



class 		__a				:
		_a     :  Any           =     []
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=None						)			->	Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      obj
									_UpperCAmelCase					=      target
									_UpperCAmelCase					=      new
									_UpperCAmelCase					=      target.split('.'						)[0]
									_UpperCAmelCase					=      {}
									_UpperCAmelCase					=      attrs or []
		def __enter__(      self						)			->	int:




									"""simple docstring"""





									*_UpperCAmelCase		,	_UpperCAmelCase					=      self.target.split('.'						)

									# Patch modules:
									# it's used to patch attributes of submodules like "os.path.join";
									# in this case we need to patch "os" and "os.path"

									for i in range(len(_SCREAMING_SNAKE_CASE						)						):
																try:
																							_UpperCAmelCase					=      import_module('.'.join(submodules[: i + 1]						)						)
																except ModuleNotFoundError:
																							continue
																# We iterate over all the globals in self.obj in case we find "os" or "os.path"
																for attr in self.obj.__dir__():
																							_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																							# We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
																							# This allows to patch renamed modules like "from os import path as ospath".
																							if obj_attr is submodule or (
																							    (isinstance(_SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj						) and obj_attr._original_module is submodule)
																							):
																														_UpperCAmelCase					=      obj_attr
																														# patch at top level
																														setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj(_SCREAMING_SNAKE_CASE      ,   attrs=self.attrs						)						)
																														_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																														# construct lower levels patches
																														for key in submodules[i + 1 :]:
																																					setattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj(getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)      ,   attrs=self.attrs						)						)
																																					_UpperCAmelCase					=      getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
																														# finally set the target attribute
																														setattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   self.new						)

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

									if submodules:  # if it's an attribute of a submodule like "os.path.join"
																try:
																							_UpperCAmelCase					=      getattr(import_module('.'.join(_SCREAMING_SNAKE_CASE						)						)      ,   _SCREAMING_SNAKE_CASE						)
																except (AttributeError, ModuleNotFoundError):
																							return
																# We iterate over all the globals in self.obj in case we find "os.path.join"
																for attr in self.obj.__dir__():
																							# We don't check for the name of the global, but rather if its value *is* "os.path.join".
																							# This allows to patch renamed attributes like "from os.path import join as pjoin".
																							if getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						) is attr_value:
																														_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																														setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.new						)
									elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
																_UpperCAmelCase					=      globals()['__builtins__'][target_attr]
																setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.new						)
									else:
																raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''						)
		def __exit__(      self      ,   *_SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





									for attr in list(self.original						):
																setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.original.pop(_SCREAMING_SNAKE_CASE						)						)
		def      UpperCAmelCase__    (      self						)			->	Optional[int]:




									"""simple docstring"""





									self.__enter__()
									self._active_patches.append(self						)






		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									try:
																self._active_patches.remove(self						)
									except ValueError:
																# If the patch hasn't been started this will fail
																return None

									return self.__exit__()
