


import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()

class 	__a   (       unittest.TestCase    ):
		@property
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    1
									_UpperCAmelCase					=    3
									_UpperCAmelCase					=    (32, 32)

									_UpperCAmelCase					=    floats_tensor((batch_size, num_channels) + sizes    ,					rng=random.Random(0		)		).to(_SCREAMING_SNAKE_CASE		)
									return image
		@property
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									torch.manual_seed(0		)
									_UpperCAmelCase					=    UNetaDModel(
									    block_out_channels=(32, 64)    ,					layers_per_block=2    ,					sample_size=32    ,					in_channels=6    ,					out_channels=3    ,					down_block_types=('DownBlock2D', 'AttnDownBlock2D')    ,					up_block_types=('AttnUpBlock2D', 'UpBlock2D')    ,					)
									return model
		@property
		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									torch.manual_seed(0		)
									_UpperCAmelCase					=    VQModel(
									    block_out_channels=[32, 64]    ,					in_channels=3    ,					out_channels=3    ,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,					latent_channels=3    ,					)
									return model
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									_UpperCAmelCase					=    'cpu'
									_UpperCAmelCase					=    self.dummy_uncond_unet
									_UpperCAmelCase					=    DDIMScheduler()
									_UpperCAmelCase					=    self.dummy_vq_model

									_UpperCAmelCase					=    LDMSuperResolutionPipeline(unet=_SCREAMING_SNAKE_CASE    ,					vqvae=_SCREAMING_SNAKE_CASE    ,					scheduler=_SCREAMING_SNAKE_CASE		)
									ldm.to(_SCREAMING_SNAKE_CASE		)
									ldm.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    self.dummy_image.to(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
									_UpperCAmelCase					=    ldm(image=_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					num_inference_steps=2    ,					output_type='numpy'		).images

									_UpperCAmelCase					=    image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)
									_UpperCAmelCase					=    np.array([0.8678, 0.8245, 0.6381, 0.6830, 0.4385, 0.5599, 0.4641, 0.6201, 0.5150]		)

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






		@unittest.skipIf(torch_device != 'cuda'    ,					'This test requires a GPU'		)
		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=    self.dummy_uncond_unet
									_UpperCAmelCase					=    DDIMScheduler()
									_UpperCAmelCase					=    self.dummy_vq_model

									# put models in fp16
									_UpperCAmelCase					=    unet.half()
									_UpperCAmelCase					=    vqvae.half()

									_UpperCAmelCase					=    LDMSuperResolutionPipeline(unet=_SCREAMING_SNAKE_CASE    ,					vqvae=_SCREAMING_SNAKE_CASE    ,					scheduler=_SCREAMING_SNAKE_CASE		)
									ldm.to(_SCREAMING_SNAKE_CASE		)
									ldm.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    self.dummy_image.to(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    ldm(_SCREAMING_SNAKE_CASE    ,					num_inference_steps=2    ,					output_type='numpy'		).images

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



@slow
@require_torch
class 	__a   (       unittest.TestCase    ):
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									_UpperCAmelCase					=    load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/vq_diffusion/teddy_bear_pool.png'		)
									_UpperCAmelCase					=    init_image.resize((64, 64)    ,					resample=PIL_INTERPOLATION['lanczos']		)

									_UpperCAmelCase					=    LDMSuperResolutionPipeline.from_pretrained('duongna/ldm-super-resolution'    ,					device_map='auto'		)
									ldm.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    torch.manual_seed(0		)
									_UpperCAmelCase					=    ldm(image=_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					num_inference_steps=20    ,					output_type='numpy'		).images

									_UpperCAmelCase					=    image[0, -3:, -3:, -1]

									assert image.shape == (1, 256, 256, 3)
									_UpperCAmelCase					=    np.array([0.7644, 0.7679, 0.7642, 0.7633, 0.7666, 0.7560, 0.7425, 0.7257, 0.6907]		)

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




import tempfile
import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVeScheduler

class 	__a   (       unittest.TestCase    ):
		# TODO adapt with class SchedulerCommonTest (scheduler needs Numpy Integration)
		_a			:						List[Any]    			=			(ScoreSdeVeScheduler,)
		_a			:						Any    			=			()
		@property
		def 			UpperCAmelCase__       (    self		)				-> List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=    4
									_UpperCAmelCase					=    3
									_UpperCAmelCase					=    8
									_UpperCAmelCase					=    8

									_UpperCAmelCase					=    torch.rand((batch_size, num_channels, height, width)		)

									return sample
		@property
		def 			UpperCAmelCase__       (    self		)				-> Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    4
									_UpperCAmelCase					=    3
									_UpperCAmelCase					=    8
									_UpperCAmelCase					=    8

									_UpperCAmelCase					=    batch_size * num_channels * height * width
									_UpperCAmelCase					=    torch.arange(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    sample.reshape(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    sample / num_elems
									_UpperCAmelCase					=    sample.permute(3    ,					0    ,					1    ,					2		)

									return sample
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									def model(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					*_SCREAMING_SNAKE_CASE		):
																return sample * t / (t + 1)

									return model
		def 			UpperCAmelCase__       (    self    ,					**_SCREAMING_SNAKE_CASE		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    {
									    'num_train_timesteps': 2000,
									    'snr': 0.15,
									    'sigma_min': 0.01,
									    'sigma_max': 1348,
									    'sampling_eps': 1e-5,
									}

									config.update(**_SCREAMING_SNAKE_CASE		)
									return config
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE=0    ,					**_SCREAMING_SNAKE_CASE		)				-> List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																_UpperCAmelCase					=    self.get_scheduler_config(**_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

																_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_correct(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler correction are not identical"
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE=0    ,					**_SCREAMING_SNAKE_CASE		)				-> List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)
									kwargs.update(_SCREAMING_SNAKE_CASE		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																_UpperCAmelCase					=    self.get_scheduler_config()
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

																_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_correct(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler correction are not identical"
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									for timesteps in [10, 100, 1000]:
																self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""simple docstring"""





									for sigma_min, sigma_max in zip([0.0001, 0.001, 0.01]    ,					[1, 100, 1000]		):
																self.check_over_configs(sigma_min=_SCREAMING_SNAKE_CASE    ,					sigma_max=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									for t in [0.1, 0.5, 0.75]:
																self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									_UpperCAmelCase					=    self.scheduler_classes[0]
									_UpperCAmelCase					=    self.get_scheduler_config()
									_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    3

									_UpperCAmelCase					=    self.dummy_model()
									_UpperCAmelCase					=    self.dummy_sample_deter

									scheduler.set_sigmas(_SCREAMING_SNAKE_CASE		)
									scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    torch.manual_seed(0		)

									for i, t in enumerate(scheduler.timesteps		):
																_UpperCAmelCase					=    scheduler.sigmas[i]

																for _ in range(scheduler.config.correct_steps		):
																							with torch.no_grad():
																														_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

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

																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					**_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase	, _UpperCAmelCase					=    output.prev_sample, output.prev_sample_mean

									_UpperCAmelCase					=    torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
									_UpperCAmelCase					=    torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

									assert np.isclose(result_sum.item()    ,					14372758528.0		)
									assert np.isclose(result_mean.item()    ,					18714530.0		)






		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									_UpperCAmelCase					=    kwargs.pop('num_inference_steps'    ,					_SCREAMING_SNAKE_CASE		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.get_scheduler_config()
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE    ,					'set_timesteps'		):
																							scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
																elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE    ,					'set_timesteps'		):
																							_UpperCAmelCase					=    num_inference_steps

																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					0    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					1    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																self.assertEqual(output_a.shape    ,					sample.shape		)
																self.assertEqual(output_a.shape    ,					output_a.shape		)
