


"""simple docstring"""



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     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					lowerCAmelCase_			       = 1
					lowerCAmelCase_			       = 3
					lowerCAmelCase_			       = (32, 32)

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







				@property
				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					torch.manual_seed(0      )
					lowerCAmelCase_			       = 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     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					torch.manual_seed(0      )
					lowerCAmelCase_			       = 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     lowerCamelCase_    (   self      ):

					'''simple docstring'''






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

					lowerCAmelCase_			       = LDMSuperResolutionPipeline(unet=UpperCAmelCase			,  vqvae=UpperCAmelCase			,  scheduler=UpperCAmelCase      )
					ldm.to(UpperCAmelCase      )
					ldm.set_progress_bar_config(disable=UpperCAmelCase      )

					lowerCAmelCase_			       = self.dummy_image.to(UpperCAmelCase      )

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

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

					assert image.shape == (1, 64, 64, 3)
					lowerCAmelCase_			       = np.array([0.8_6_7_8, 0.8_2_4_5, 0.6_3_8_1, 0.6_8_3_0, 0.4_3_8_5, 0.5_5_9_9, 0.4_6_4_1, 0.6_2_0_1, 0.5_1_5_0]      )

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







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

					'''simple docstring'''






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

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

					lowerCAmelCase_			       = LDMSuperResolutionPipeline(unet=UpperCAmelCase			,  vqvae=UpperCAmelCase			,  scheduler=UpperCAmelCase      )
					ldm.to(UpperCAmelCase      )
					ldm.set_progress_bar_config(disable=UpperCAmelCase      )

					lowerCAmelCase_			       = self.dummy_image.to(UpperCAmelCase      )

					lowerCAmelCase_			       = ldm(UpperCAmelCase			,  num_inference_steps=2			,  output_type='''numpy'''      ).images

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







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







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






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

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

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

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

					assert image.shape == (1, 256, 256, 3)
					lowerCAmelCase_			       = np.array([0.7_6_4_4, 0.7_6_7_9, 0.7_6_4_2, 0.7_6_3_3, 0.7_6_6_6, 0.7_5_6_0, 0.7_4_2_5, 0.7_2_5_7, 0.6_9_0_7]      )

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



"""simple docstring"""



import torch

from diffusers import LMSDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class       __a	(   __snake_case       ):
				lowerCamelCase    :					Union[str, Any] 					=(LMSDiscreteScheduler,)
				lowerCamelCase    :					int 					=10







				def     lowerCamelCase_    (   self			,  **UpperCAmelCase      ):

					'''simple docstring'''






					lowerCAmelCase_			       = {
					    '''num_train_timesteps''': 1100,
					    '''beta_start''': 0.0_0_0_1,
					    '''beta_end''': 0.0_2,
					    '''beta_schedule''': '''linear''',
					}

					config.update(**UpperCAmelCase      )
					return config







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					for timesteps in [10, 50, 100, 1000]:
						self.check_over_configs(num_train_timesteps=UpperCAmelCase      )







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]			,  [0.0_0_0_2, 0.0_0_2, 0.0_2]      ):
						self.check_over_configs(beta_start=UpperCAmelCase			,  beta_end=UpperCAmelCase      )







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					for schedule in ["linear", "scaled_linear"]:
						self.check_over_configs(beta_schedule=UpperCAmelCase      )







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					for prediction_type in ["epsilon", "v_prediction"]:
						self.check_over_configs(prediction_type=UpperCAmelCase      )







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					for t in [0, 500, 800]:
						self.check_over_forward(time_step=UpperCAmelCase      )







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					lowerCAmelCase_			       = self.scheduler_classes[0]
					lowerCAmelCase_			       = self.get_scheduler_config()
					lowerCAmelCase_			       = scheduler_class(**UpperCAmelCase      )

					scheduler.set_timesteps(self.num_inference_steps      )

					lowerCAmelCase_			       = self.dummy_model()
					lowerCAmelCase_			       = self.dummy_sample_deter * scheduler.init_noise_sigma

					for i, t in enumerate(scheduler.timesteps      ):
						lowerCAmelCase_			       = scheduler.scale_model_input(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = model(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = scheduler.step(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase      )
						lowerCAmelCase_			       = output.prev_sample

					lowerCAmelCase_			       = torch.sum(torch.abs(UpperCAmelCase      )      )
					lowerCAmelCase_			       = torch.mean(torch.abs(UpperCAmelCase      )      )

					assert abs(result_sum.item() - 1_0_0_6.3_8_8      ) < 1E-2
					assert abs(result_mean.item() - 1.3_1      ) < 1E-3







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					lowerCAmelCase_			       = self.scheduler_classes[0]
					lowerCAmelCase_			       = self.get_scheduler_config(prediction_type='''v_prediction'''      )
					lowerCAmelCase_			       = scheduler_class(**UpperCAmelCase      )

					scheduler.set_timesteps(self.num_inference_steps      )

					lowerCAmelCase_			       = self.dummy_model()
					lowerCAmelCase_			       = self.dummy_sample_deter * scheduler.init_noise_sigma

					for i, t in enumerate(scheduler.timesteps      ):
						lowerCAmelCase_			       = scheduler.scale_model_input(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = model(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = scheduler.step(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase      )
						lowerCAmelCase_			       = output.prev_sample

					lowerCAmelCase_			       = torch.sum(torch.abs(UpperCAmelCase      )      )
					lowerCAmelCase_			       = torch.mean(torch.abs(UpperCAmelCase      )      )

					assert abs(result_sum.item() - 0.0_0_1_7      ) < 1E-2
					assert abs(result_mean.item() - 2.2_6_7_6E-0_6      ) < 1E-3







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					lowerCAmelCase_			       = self.scheduler_classes[0]
					lowerCAmelCase_			       = self.get_scheduler_config()
					lowerCAmelCase_			       = scheduler_class(**UpperCAmelCase      )

					scheduler.set_timesteps(self.num_inference_steps			,  device=UpperCAmelCase      )

					lowerCAmelCase_			       = self.dummy_model()
					lowerCAmelCase_			       = self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
					lowerCAmelCase_			       = sample.to(UpperCAmelCase      )

					for i, t in enumerate(scheduler.timesteps      ):
						lowerCAmelCase_			       = scheduler.scale_model_input(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = model(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = scheduler.step(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase      )
						lowerCAmelCase_			       = output.prev_sample

					lowerCAmelCase_			       = torch.sum(torch.abs(UpperCAmelCase      )      )
					lowerCAmelCase_			       = torch.mean(torch.abs(UpperCAmelCase      )      )

					assert abs(result_sum.item() - 1_0_0_6.3_8_8      ) < 1E-2
					assert abs(result_mean.item() - 1.3_1      ) < 1E-3







				def     lowerCamelCase_    (   self      ):

					'''simple docstring'''






					lowerCAmelCase_			       = self.scheduler_classes[0]
					lowerCAmelCase_			       = self.get_scheduler_config()
					lowerCAmelCase_			       = scheduler_class(**UpperCAmelCase			,  use_karras_sigmas=UpperCAmelCase      )

					scheduler.set_timesteps(self.num_inference_steps			,  device=UpperCAmelCase      )

					lowerCAmelCase_			       = self.dummy_model()
					lowerCAmelCase_			       = self.dummy_sample_deter.to(UpperCAmelCase      ) * scheduler.init_noise_sigma
					lowerCAmelCase_			       = sample.to(UpperCAmelCase      )

					for t in scheduler.timesteps:
						lowerCAmelCase_			       = scheduler.scale_model_input(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = model(UpperCAmelCase			,  UpperCAmelCase      )

						lowerCAmelCase_			       = scheduler.step(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase      )
						lowerCAmelCase_			       = output.prev_sample

					lowerCAmelCase_			       = torch.sum(torch.abs(UpperCAmelCase      )      )
					lowerCAmelCase_			       = torch.mean(torch.abs(UpperCAmelCase      )      )

					assert abs(result_sum.item() - 3_8_1_2.9_9_2_7      ) < 2E-2
					assert abs(result_mean.item() - 4.9_6_4_8      ) < 1E-3