

"""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


__UpperCamelCase    : Any				 =						logging.get_logger(__name__)

enable_full_determinism()







class 							SCREAMING_SNAKE_CASE   (						a_       ,				a_       ,				unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			UNetaDModel
					lowercase__												=			"sample"







					@property
					def 			__lowerCAmelCase					(			self     :							int						):
						lowerCAmelCase__      :		Any    							=      4
						lowerCAmelCase__      :		Optional[Any]    							=      3
						lowerCAmelCase__      :		Any    							=      (3_2, 3_2)

						lowerCAmelCase__      :		Tuple    							=      floats_tensor((batch_size, num_channels) + sizes						).to(lowercase_						)
						lowerCAmelCase__      :		List[str]    							=      torch.tensor([1_0]						).to(lowercase_						)

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







					@property
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						return (3, 3_2, 3_2)







					@property
					def 			__lowerCAmelCase					(			self     :							Any						):
						return (3, 3_2, 3_2)







					def 			__lowerCAmelCase					(			self     :							List[Any]						):
						lowerCAmelCase__      :		Any    							=      {
						    '''block_out_channels''': (3_2, 6_4),
						    '''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''': 3_2,
						}
						lowerCAmelCase__      :		Optional[int]    							=      self.dummy_input
						return init_dict, inputs_dict







class 							SCREAMING_SNAKE_CASE   (						a_       ,				a_       ,				unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			UNetaDModel
					lowercase__												=			"sample"







					@property
					def 			__lowerCAmelCase					(			self     :							Tuple						):
						lowerCAmelCase__      :		Optional[int]    							=      4
						lowerCAmelCase__      :		str    							=      4
						lowerCAmelCase__      :		List[Any]    							=      (3_2, 3_2)

						lowerCAmelCase__      :		Tuple    							=      floats_tensor((batch_size, num_channels) + sizes						).to(lowercase_						)
						lowerCAmelCase__      :		Tuple    							=      torch.tensor([1_0]						).to(lowercase_						)

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







					@property
					def 			__lowerCAmelCase					(			self     :							Tuple						):
						return (4, 3_2, 3_2)







					@property
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						return (4, 3_2, 3_2)







					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						lowerCAmelCase__      :		Dict    							=      {
						    '''sample_size''': 3_2,
						    '''in_channels''': 4,
						    '''out_channels''': 4,
						    '''layers_per_block''': 2,
						    '''block_out_channels''': (3_2, 6_4),
						    '''attention_head_dim''': 3_2,
						    '''down_block_types''': ('''DownBlock2D''', '''DownBlock2D'''),
						    '''up_block_types''': ('''UpBlock2D''', '''UpBlock2D'''),
						}
						lowerCAmelCase__      :		Optional[int]    							=      self.dummy_input
						return init_dict, inputs_dict







					def 			__lowerCAmelCase					(			self     :							Tuple						):
						lowerCAmelCase__							,lowerCAmelCase__      :		List[Any]    							=      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_						)
						lowerCAmelCase__      :		Dict    							=      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     :							int						):
						lowerCAmelCase__							,lowerCAmelCase__      :		str    							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''  ,output_loading_info=lowercase_						)
						model.to(lowercase_						)
						lowerCAmelCase__      :		Union[str, Any]    							=      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     :							str						):
						# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
						lowerCAmelCase__							,lowerCAmelCase__      :		Any    							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''  ,output_loading_info=lowercase_						)
						model_accelerate.to(lowercase_						)
						model_accelerate.eval()

						lowerCAmelCase__      :		Optional[Any]    							=      torch.randn(
						    1  ,model_accelerate.config.in_channels  ,model_accelerate.config.sample_size  ,model_accelerate.config.sample_size  ,generator=torch.manual_seed(0						)  ,)
						lowerCAmelCase__      :		Tuple    							=      noise.to(lowercase_						)
						lowerCAmelCase__      :		Optional[int]    							=      torch.tensor([1_0] * noise.shape[0]						).to(lowercase_						)

						lowerCAmelCase__      :		Any    							=      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()

						lowerCAmelCase__							,lowerCAmelCase__      :		Union[str, Any]    							=      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()
						lowerCAmelCase__      :		Tuple    							=      model_normal_load(lowercase_  ,lowercase_						)['''sample''']

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







					def 			__lowerCAmelCase					(			self     :							List[Any]						):
						lowerCAmelCase__      :		str    							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''						)
						model.eval()
						model.to(lowercase_						)

						lowerCAmelCase__      :		Optional[Any]    							=      torch.randn(
						    1  ,model.config.in_channels  ,model.config.sample_size  ,model.config.sample_size  ,generator=torch.manual_seed(0						)  ,)
						lowerCAmelCase__      :		List[Any]    							=      noise.to(lowercase_						)
						lowerCAmelCase__      :		Union[str, Any]    							=      torch.tensor([1_0] * noise.shape[0]						).to(lowercase_						)

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

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







class 							SCREAMING_SNAKE_CASE   (						a_       ,				a_       ,				unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			UNetaDModel
					lowercase__												=			"sample"







					@property
					def 			__lowerCAmelCase					(			self     :							str  ,lowercase_     :							str=(3_2, 3_2)						):
						lowerCAmelCase__      :		List[str]    							=      4
						lowerCAmelCase__      :		Union[str, Any]    							=      3

						lowerCAmelCase__      :		List[Any]    							=      floats_tensor((batch_size, num_channels) + sizes						).to(lowercase_						)
						lowerCAmelCase__      :		List[Any]    							=      torch.tensor(batch_size * [1_0]						).to(dtype=torch.intaa  ,device=lowercase_						)

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







					@property
					def 			__lowerCAmelCase					(			self     :							Dict						):
						return (3, 3_2, 3_2)







					@property
					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						return (3, 3_2, 3_2)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						lowerCAmelCase__      :		int    							=      {
						    '''block_out_channels''': [3_2, 6_4, 6_4, 6_4],
						    '''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''',
						    ],
						}
						lowerCAmelCase__      :		Tuple    							=      self.dummy_input
						return init_dict, inputs_dict







					@slow
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						lowerCAmelCase__							,lowerCAmelCase__      :		Dict    							=      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_						)
						lowerCAmelCase__      :		Any    							=      self.dummy_input
						lowerCAmelCase__      :		str    							=      floats_tensor((4, 3) + (2_5_6, 2_5_6)						).to(lowercase_						)
						lowerCAmelCase__      :		List[str]    							=      noise
						lowerCAmelCase__      :		List[str]    							=      model(**lowercase_						)

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







					@slow
					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						lowerCAmelCase__      :		List[str]    							=      UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''						)
						model.to(lowercase_						)

						lowerCAmelCase__      :		str    							=      4
						lowerCAmelCase__      :		Any    							=      3
						lowerCAmelCase__      :		int    							=      (2_5_6, 2_5_6)

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

						with torch.no_grad():
							lowerCAmelCase__      :		List[str]    							=      model(lowercase_  ,lowercase_						).sample

						lowerCAmelCase__      :		str    							=      output[0, -3:, -3:, -1].flatten().cpu()
						# fmt: off
						lowerCAmelCase__      :		Any    							=      torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -1_0980.7129, -2_0028.8535, 8148.2822, 2342.2905, 567.7608]						)
						# fmt: on

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







					def 			__lowerCAmelCase					(			self     :							int						):
						lowerCAmelCase__      :		Optional[int]    							=      UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''						)
						model.to(lowercase_						)

						lowerCAmelCase__      :		Tuple    							=      4
						lowerCAmelCase__      :		Optional[int]    							=      3
						lowerCAmelCase__      :		List[Any]    							=      (3_2, 3_2)

						lowerCAmelCase__      :		Tuple    							=      torch.ones((batch_size, num_channels) + sizes						).to(lowercase_						)
						lowerCAmelCase__      :		List[str]    							=      torch.tensor(batch_size * [1E-4]						).to(lowercase_						)

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

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







					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						# not required for this model
						pass







"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    : str				 =						{
    '''configuration_swinv2''': ['''SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Swinv2Config'''],
}


try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__UpperCamelCase    : List[Any]				 =						[
							    '''SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST''',
							    '''Swinv2ForImageClassification''',
							    '''Swinv2ForMaskedImageModeling''',
							    '''Swinv2Model''',
							    '''Swinv2PreTrainedModel''',
							]


if TYPE_CHECKING:
							from .configuration_swinva import SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinvaConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_swinva import (
														    SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST,
														    SwinvaForImageClassification,
														    SwinvaForMaskedImageModeling,
														    SwinvaModel,
														    SwinvaPreTrainedModel,
														)


else:
							import sys

							__UpperCamelCase    : List[Any]				 =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




