


'''simple docstring'''
import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()
class    lowercase	(							unittest.TestCase					):

		"""simple docstring"""

		@property
		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							torch.manual_seed(0							)
							UpperCamelCase__   :Tuple												= UNetaDModel(
							    block_out_channels=(32, 64)    ,   layers_per_block=2    ,   sample_size=32    ,   in_channels=3    ,   out_channels=3    ,   down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')    ,   up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')    ,   )
							return model




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[Any]												= self.dummy_uncond_unet
							UpperCamelCase__   :int												= ScoreSdeVeScheduler()

							UpperCamelCase__   :Tuple												= ScoreSdeVePipeline(unet=a_    ,   scheduler=a_							)
							sde_ve.to(a_							)
							sde_ve.set_progress_bar_config(disable=a_							)

							UpperCamelCase__   :Any												= torch.manual_seed(0							)
							UpperCamelCase__   :Union[str, Any]												= sde_ve(num_inference_steps=2    ,   output_type='''numpy'''    ,   generator=a_							).images

							UpperCamelCase__   :Optional[int]												= torch.manual_seed(0							)
							UpperCamelCase__   :str												= sde_ve(num_inference_steps=2    ,   output_type='''numpy'''    ,   generator=a_    ,   return_dict=a_							)[
							    0
							]

							UpperCamelCase__   :List[str]												= image[0, -3:, -3:, -1]
							UpperCamelCase__   :Union[str, Any]												= image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							UpperCamelCase__   :Tuple												= np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]							)

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



@slow
@require_torch
class    lowercase	(							unittest.TestCase					):

		"""simple docstring"""

		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Tuple												= '''google/ncsnpp-church-256'''
							UpperCamelCase__   :List[Any]												= UNetaDModel.from_pretrained(a_							)

							UpperCamelCase__   :Dict												= ScoreSdeVeScheduler.from_pretrained(a_							)

							UpperCamelCase__   :Any												= ScoreSdeVePipeline(unet=a_    ,   scheduler=a_							)
							sde_ve.to(a_							)
							sde_ve.set_progress_bar_config(disable=a_							)

							UpperCamelCase__   :Dict												= torch.manual_seed(0							)
							UpperCamelCase__   :int												= sde_ve(num_inference_steps=10    ,   output_type='''numpy'''    ,   generator=a_							).images

							UpperCamelCase__   :Union[str, Any]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 256, 256, 3)

							UpperCamelCase__   :Optional[Any]												= np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]							)

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



"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()



@slow
@require_torch_gpu
class 			_UpperCAmelCase (							unittest.TestCase				):






		'''simple docstring'''

		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Tuple						=				StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''					)
			__snake_case	:						List[str]						=				sd_pipe.to(a_					)
			sd_pipe.set_progress_bar_config(disable=a_					)

			sd_pipe.set_scheduler('''sample_euler'''					)

			__snake_case	:						Tuple						=				'''A painting of a squirrel eating a burger'''
			__snake_case	:						str						=				torch.manual_seed(0					)
			__snake_case	:						int						=				sd_pipe([prompt]				,       generator=a_				,       guidance_scale=9.0				,       num_inference_steps=20				,       output_type='''np'''					)

			__snake_case	:						List[str]						=				output.images

			__snake_case	:						Tuple						=				image[0, -3:, -3:, -1]

			assert image.shape == (1, 5_12, 5_12, 3)
			__snake_case	:						Dict						=				np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]					)

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''					)
			__snake_case	:						Union[str, Any]						=				sd_pipe.to(a_					)
			sd_pipe.set_progress_bar_config(disable=a_					)

			sd_pipe.set_scheduler('''sample_euler'''					)

			__snake_case	:						str						=				'''A painting of a squirrel eating a burger'''
			__snake_case	:						List[str]						=				torch.manual_seed(0					)
			__snake_case	:						Dict						=				sd_pipe([prompt]				,       generator=a_				,       guidance_scale=9.0				,       num_inference_steps=20				,       output_type='''np'''					)

			__snake_case	:						Any						=				output.images

			__snake_case	:						Tuple						=				image[0, -3:, -3:, -1]

			assert image.shape == (1, 5_12, 5_12, 3)
			__snake_case	:						Tuple						=				np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]					)

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Any						=				StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''					)
			__snake_case	:						Union[str, Any]						=				sd_pipe.to(a_					)
			sd_pipe.set_progress_bar_config(disable=a_					)

			sd_pipe.set_scheduler('''sample_dpmpp_2m'''					)

			__snake_case	:						int						=				'''A painting of a squirrel eating a burger'''
			__snake_case	:						Optional[Any]						=				torch.manual_seed(0					)
			__snake_case	:						Dict						=				sd_pipe(
			    [prompt]				,       generator=a_				,       guidance_scale=7.5				,       num_inference_steps=15				,       output_type='''np'''				,       use_karras_sigmas=a_				,       )

			__snake_case	:						Optional[int]						=				output.images

			__snake_case	:						Any						=				image[0, -3:, -3:, -1]

			assert image.shape == (1, 5_12, 5_12, 3)
			__snake_case	:						Any						=				np.array(
			    [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]					)

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





