





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

enable_full_determinism()






class 					lowercase__ (			__lowerCamelCase    ,       __lowerCamelCase    ,       unittest.TestCase       ):

			'''simple docstring'''


			a		:							Any     	=	UNetaDModel
			a		:							str     	=	"sample"




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






							"""simple docstring"""
							UpperCamelCase__							:       Optional[int]      =   4
							UpperCamelCase__							:       str      =   3
							UpperCamelCase__							:       Dict      =   (32, 32)

							UpperCamelCase__							:       Optional[int]      =   floats_tensor((batch_size, num_channels) + sizes							).to(__magic_name__							)
							UpperCamelCase__							:       List[str]      =   torch.tensor([10]							).to(__magic_name__							)

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




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






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




			@property
			def    UpperCamelCase__						(					self							)				->   Tuple:






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




			def    UpperCamelCase__						(					self							)				->   Optional[int]:






							"""simple docstring"""
							UpperCamelCase__							:       int      =   {
							    '''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__							:       Dict      =   self.dummy_input
							return init_dict, inputs_dict






class 					lowercase__ (			__lowerCamelCase    ,       __lowerCamelCase    ,       unittest.TestCase       ):

			'''simple docstring'''


			a		:							str     	=	UNetaDModel
			a		:							Any     	=	"sample"




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






							"""simple docstring"""
							UpperCamelCase__							:       Dict      =   4
							UpperCamelCase__							:       List[str]      =   4
							UpperCamelCase__							:       Optional[int]      =   (32, 32)

							UpperCamelCase__							:       Optional[Any]      =   floats_tensor((batch_size, num_channels) + sizes							).to(__magic_name__							)
							UpperCamelCase__							:       Any      =   torch.tensor([10]							).to(__magic_name__							)

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




			@property
			def    UpperCamelCase__						(					self							)				->   str:






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




			@property
			def    UpperCamelCase__						(					self							)				->   Any:






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




			def    UpperCamelCase__						(					self							)				->   Dict:






							"""simple docstring"""
							UpperCamelCase__							:       Tuple      =   {
							    '''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__							:       str      =   self.dummy_input
							return init_dict, inputs_dict




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






							"""simple docstring"""
							UpperCamelCase__							:       Optional[Any]      =   UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update''',	output_loading_info=__magic_name__							)

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

							model.to(__magic_name__							)
							UpperCamelCase__							:       List[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    UpperCamelCase__						(					self							)				->   Any:






							"""simple docstring"""
							UpperCamelCase__							:       List[Any]      =   UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update''',	output_loading_info=__magic_name__							)
							model.to(__magic_name__							)
							UpperCamelCase__							:       List[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    UpperCamelCase__						(					self							)				->   int:






							"""simple docstring"""
							UpperCamelCase__							:       List[Any]      =   UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update''',	output_loading_info=__magic_name__							)
							model_accelerate.to(__magic_name__							)
							model_accelerate.eval()

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

							UpperCamelCase__							:       Optional[Any]      =   model_accelerate(__magic_name__,	__magic_name__							)['''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__							:       Optional[Any]      =   UNetaDModel.from_pretrained(
							    '''fusing/unet-ldm-dummy-update''',	output_loading_info=__magic_name__,	low_cpu_mem_usage=__magic_name__							)
							model_normal_load.to(__magic_name__							)
							model_normal_load.eval()
							UpperCamelCase__							:       List[str]      =   model_normal_load(__magic_name__,	__magic_name__							)['''sample''']

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




			def    UpperCamelCase__						(					self							)				->   Any:






							"""simple docstring"""
							UpperCamelCase__							:       List[str]      =   UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''							)
							model.eval()
							model.to(__magic_name__							)

							UpperCamelCase__							:       Tuple      =   torch.randn(
							    1,	model.config.in_channels,	model.config.sample_size,	model.config.sample_size,	generator=torch.manual_seed(0							),	)
							UpperCamelCase__							:       Tuple      =   noise.to(__magic_name__							)
							UpperCamelCase__							:       Optional[int]      =   torch.tensor([10] * noise.shape[0]							).to(__magic_name__							)

							with torch.no_grad():
											UpperCamelCase__							:       Any      =   model(__magic_name__,	__magic_name__							).sample

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






class 					lowercase__ (			__lowerCamelCase    ,       __lowerCamelCase    ,       unittest.TestCase       ):

			'''simple docstring'''


			a		:							Optional[Any]     	=	UNetaDModel
			a		:							Optional[Any]     	=	"sample"




			@property
			def    UpperCamelCase__						(					self,	__magic_name__=(32, 32)							)				->   Any:






							"""simple docstring"""
							UpperCamelCase__							:       List[str]      =   4
							UpperCamelCase__							:       Optional[Any]      =   3

							UpperCamelCase__							:       List[Any]      =   floats_tensor((batch_size, num_channels) + sizes							).to(__magic_name__							)
							UpperCamelCase__							:       Optional[int]      =   torch.tensor(batch_size * [10]							).to(dtype=torch.intaa,	device=__magic_name__							)

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




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






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




			@property
			def    UpperCamelCase__						(					self							)				->   int:






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




			def    UpperCamelCase__						(					self							)				->   str:






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




			@slow
			def    UpperCamelCase__						(					self							)				->   Union[str, Any]:






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

							model.to(__magic_name__							)
							UpperCamelCase__							:       Dict      =   self.dummy_input
							UpperCamelCase__							:       Dict      =   floats_tensor((4, 3) + (256, 256)							).to(__magic_name__							)
							UpperCamelCase__							:       str      =   noise
							UpperCamelCase__							:       str      =   model(**__magic_name__							)

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




			@slow
			def    UpperCamelCase__						(					self							)				->   str:






							"""simple docstring"""
							UpperCamelCase__							:       List[Any]      =   UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''							)
							model.to(__magic_name__							)

							UpperCamelCase__							:       List[str]      =   4
							UpperCamelCase__							:       Dict      =   3
							UpperCamelCase__							:       Any      =   (256, 256)

							UpperCamelCase__							:       Optional[Any]      =   torch.ones((batch_size, num_channels) + sizes							).to(__magic_name__							)
							UpperCamelCase__							:       List[str]      =   torch.tensor(batch_size * [1E-4]							).to(__magic_name__							)

							with torch.no_grad():
											UpperCamelCase__							:       List[str]      =   model(__magic_name__,	__magic_name__							).sample

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




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






							"""simple docstring"""
							UpperCamelCase__							:       str      =   UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''							)
							model.to(__magic_name__							)

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

							UpperCamelCase__							:       Any      =   torch.ones((batch_size, num_channels) + sizes							).to(__magic_name__							)
							UpperCamelCase__							:       List[Any]      =   torch.tensor(batch_size * [1E-4]							).to(__magic_name__							)

							with torch.no_grad():
											UpperCamelCase__							:       List[str]      =   model(__magic_name__,	__magic_name__							).sample

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




			def    UpperCamelCase__						(					self							)				->   str:






							"""simple docstring"""
							pass








import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


UpperCAmelCase_  							=				False






class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   pass






@nightly
@require_torch_gpu
class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


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






       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




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






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   VersatileDiffusionTextToImagePipeline.from_pretrained('''shi-labs/versatile-diffusion'''							)
       # remove text_unet
       pipe.remove_unused_weights()
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       int      =   '''A painting of a squirrel eating a burger '''
       UpperCamelCase__							:       str      =   torch.manual_seed(0							)
       UpperCamelCase__							:       int      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=2,	output_type='''numpy'''							).images

       with tempfile.TemporaryDirectory() as tmpdirname:
           pipe.save_pretrained(__magic_name__							)
           UpperCamelCase__							:       Union[str, Any]      =   VersatileDiffusionTextToImagePipeline.from_pretrained(__magic_name__							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Optional[Any]      =   generator.manual_seed(0							)
       UpperCamelCase__							:       str      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=2,	output_type='''numpy'''							).images

       assert np.abs(image - new_image							).sum() < 1E-5, "Models don't have the same forward pass"




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       str      =   VersatileDiffusionTextToImagePipeline.from_pretrained(
           '''shi-labs/versatile-diffusion''',	torch_dtype=torch.floataa							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Optional[int]      =   '''A painting of a squirrel eating a burger '''
       UpperCamelCase__							:       str      =   torch.manual_seed(0							)
       UpperCamelCase__							:       Optional[int]      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=50,	output_type='''numpy'''							).images

       UpperCamelCase__							:       int      =   image[0, 253:256, 253:256, -1]

       assert image.shape == (1, 512, 512, 3)
       UpperCamelCase__							:       int      =   np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]							)

       assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2

