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


snake_case__				= logging.get_logger(__name__)

enable_full_determinism()



class  UpperCamelCase				( __lowercase   ,  __lowercase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  UNetaDModel
     A_										=  'sample'



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






           """simple docstring"""
           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				3
           _lowerCamelCase			    =				(32, 32)

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

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



     @property
     def 		UpperCamelCase_						(					self						) ->				str:






           """simple docstring"""
           return (3, 32, 32)



     @property
     def 		UpperCamelCase_						(					self						) ->				str:






           """simple docstring"""
           return (3, 32, 32)





     def 		UpperCamelCase_						(					self						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				{
               '''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,
           }
           _lowerCamelCase			    =				self.dummy_input
           return init_dict, inputs_dict





class  UpperCamelCase				( __lowercase   ,  __lowercase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  UNetaDModel
     A_										=  'sample'



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






           """simple docstring"""
           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				(32, 32)

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

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



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






           """simple docstring"""
           return (4, 32, 32)



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






           """simple docstring"""
           return (4, 32, 32)



     def 		UpperCamelCase_						(					self						) ->				List[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				{
               '''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'''),
           }
           _lowerCamelCase			    =				self.dummy_input
           return init_dict, inputs_dict



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase			    =				UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''				,				output_loading_info=A_						)

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

           model.to(A_						)
           _lowerCamelCase			    =				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						) ->				Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase			    =				UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''				,				output_loading_info=A_						)
           model.to(A_						)
           _lowerCamelCase			    =				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"""
           # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
           _lowerCamelCase    ,					_lowerCamelCase			    =				UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''				,				output_loading_info=A_						)
           model_accelerate.to(A_						)
           model_accelerate.eval()

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

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

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





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






           """simple docstring"""
           _lowerCamelCase			    =				UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''						)
           model.eval()
           model.to(A_						)

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

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				A_						).sample

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





class  UpperCamelCase				( __lowercase   ,  __lowercase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  UNetaDModel
     A_										=  'sample'



     @property
     def 		UpperCamelCase_						(					self				,				A_=(32, 32)						) ->				Optional[int]:






           """simple docstring"""
           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				3

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

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



     @property
     def 		UpperCamelCase_						(					self						) ->				Union[str, Any]:






           """simple docstring"""
           return (3, 32, 32)



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






           """simple docstring"""
           return (3, 32, 32)



     def 		UpperCamelCase_						(					self						) ->				Dict:






           """simple docstring"""
           _lowerCamelCase			    =				{
               '''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''',
               ],
           }
           _lowerCamelCase			    =				self.dummy_input
           return init_dict, inputs_dict



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






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase			    =				UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''				,				output_loading_info=A_						)
           self.assertIsNotNone(A_						)
           self.assertEqual(len(loading_info['''missing_keys''']						)				,				0						)

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

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



     @slow
     def 		UpperCamelCase_						(					self						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''						)
           model.to(A_						)

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

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

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				A_						).sample

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



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






           """simple docstring"""
           _lowerCamelCase			    =				UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''						)
           model.to(A_						)

           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				3
           _lowerCamelCase			    =				(32, 32)

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

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				A_						).sample

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





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






           """simple docstring"""
           # not required for this model
           pass
import warnings

from ...utils import logging
from .image_processing_dpt import DPTImageProcessor


snake_case__				= logging.get_logger(__name__)



class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     def __init__(					self				,				*A_				,				**A_						) ->				None:






           """simple docstring"""
           warnings.warn(
               '''The class DPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'''
               ''' use DPTImageProcessor instead.'''				,				A_				,				)
           super().__init__(*A_				,				**A_						)