


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 lowerCAmelCase_		(					unittest.TestCase			):

 '''simple docstring'''

 @property
 def 		UpperCamelCase__	(      self   ):
       snake_case_			= 1
       snake_case_			= 3
       snake_case_			= (32, 32)

       snake_case_			= floats_tensor((batch_size, num_channels) + sizes						, rng=random.Random(0   )   ).to(_SCREAMING_SNAKE_CASE   )
       return image



 @property
 def 		UpperCamelCase__	(      self   ):
       torch.manual_seed(0   )
       snake_case_			= 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   ):
       torch.manual_seed(0   )
       snake_case_			= 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   ):
       snake_case_			= '''cpu'''
       snake_case_			= self.dummy_uncond_unet
       snake_case_			= DDIMScheduler()
       snake_case_			= self.dummy_vq_model

       snake_case_			= 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   )

       snake_case_			= self.dummy_image.to(_SCREAMING_SNAKE_CASE   )

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

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

       assert image.shape == (1, 64, 64, 3)
       snake_case_			= np.array([0.8_678, 0.8_245, 0.6_381, 0.6_830, 0.4_385, 0.5_599, 0.4_641, 0.6_201, 0.5_150]   )

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





 @unittest.skipIf(torch_device != '''cuda'''						, '''This test requires a GPU'''   )
 def 		UpperCamelCase__	(      self   ):
       snake_case_			= self.dummy_uncond_unet
       snake_case_			= DDIMScheduler()
       snake_case_			= self.dummy_vq_model

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

       snake_case_			= 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   )

       snake_case_			= self.dummy_image.to(_SCREAMING_SNAKE_CASE   )

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

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







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

 '''simple docstring'''

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

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

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

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

       assert image.shape == (1, 2_56, 2_56, 3)
       snake_case_			= np.array([0.7_644, 0.7_679, 0.7_642, 0.7_633, 0.7_666, 0.7_560, 0.7_425, 0.7_257, 0.6_907]   )

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




import tempfile
import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVeScheduler




class       __UpperCAmelCase							(							unittest.TestCase	):





       '''simple docstring'''



       __lowercase   :   str  					=							(ScoreSdeVeScheduler,)
       __lowercase   :   List[str]  					=							()

       @property
       def    __A			( self					)						->	List[str]:
           A_      =  4
           A_      =  3
           A_      =  8
           A_      =  8

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

           return sample

       @property
       def    __A			( self					)						->	Optional[Any]:
           A_      =  4
           A_      =  3
           A_      =  8
           A_      =  8

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

           return sample

       def    __A			( self					)						->	Union[str, Any]:
           def model(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	*_SCREAMING_SNAKE_CASE					):
               return sample * t / (t + 1)

           return model

       def    __A			( self       ,	**_SCREAMING_SNAKE_CASE					)						->	Tuple:
           A_      =  {
               '''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    __A			( self       ,	_SCREAMING_SNAKE_CASE=0       ,	**_SCREAMING_SNAKE_CASE					)						->	Any:
           A_      =  dict(self.forward_default_kwargs					)

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

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

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

               A_      =  scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  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"

               A_      =  scheduler.step_correct(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  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    __A			( self       ,	_SCREAMING_SNAKE_CASE=0       ,	**_SCREAMING_SNAKE_CASE					)						->	Any:
           A_      =  dict(self.forward_default_kwargs					)
           kwargs.update(_SCREAMING_SNAKE_CASE					)

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

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

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

               A_      =  scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  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"

               A_      =  scheduler.step_correct(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  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    __A			( self					)						->	Any:
           for timesteps in [10, 100, 1000]:
               self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	Dict:
           for sigma_min, sigma_max in zip([0.0_001, 0.001, 0.01]       ,	[1, 100, 1000]					):
               self.check_over_configs(sigma_min=_SCREAMING_SNAKE_CASE       ,	sigma_max=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	Optional[int]:
           for t in [0.1, 0.5, 0.75]:
               self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	int:
           A_      =  dict(self.forward_default_kwargs					)

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

           A_      =  3

           A_      =  self.dummy_model()
           A_      =  self.dummy_sample_deter

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

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

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

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

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

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

           assert np.isclose(result_sum.item()       ,	14_372_758_528.0					)
           assert np.isclose(result_mean.item()       ,	18_714_530.0					)

       def    __A			( self					)						->	List[Any]:
           A_      =  dict(self.forward_default_kwargs					)

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

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

               A_      =  self.dummy_sample
               A_      =  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'''					):
                   A_      =  num_inference_steps

               A_      =  scheduler.step_pred(_SCREAMING_SNAKE_CASE       ,	0       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  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					)






