


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     __snake_case							(      unittest.TestCase):
					@property
					def   a (	self					:	Tuple     )       ->				Tuple:



											'''simple docstring'''





											__magic_name__ 		=			1
											__magic_name__ 		=			3
											__magic_name__ 		=			(3_2, 3_2)

											__magic_name__ 		=			floats_tensor((batch_size, num_channels) + sizes  ,			rng=random.Random(0     )     ).to(_lowerCAmelCase     )
											return image
					@property
					def   a (	self					:	int     )       ->				Optional[int]:



											'''simple docstring'''





											torch.manual_seed(0     )
											__magic_name__ 		=			UNetaDModel(
											    block_out_channels=(3_2, 6_4)  ,			layers_per_block=2  ,			sample_size=3_2  ,			in_channels=6  ,			out_channels=3  ,			down_block_types=('DownBlock2D', 'AttnDownBlock2D')  ,			up_block_types=('AttnUpBlock2D', 'UpBlock2D')  ,			)
											return model
					@property
					def   a (	self					:	str     )       ->				Dict:



											'''simple docstring'''





											torch.manual_seed(0     )
											__magic_name__ 		=			VQModel(
											    block_out_channels=[3_2, 6_4]  ,			in_channels=3  ,			out_channels=3  ,			down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,			latent_channels=3  ,			)
											return model
					def   a (	self					:	List[Any]     )       ->				Tuple:



											'''simple docstring'''





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

											__magic_name__ 		=			LDMSuperResolutionPipeline(unet=_lowerCAmelCase  ,			vqvae=_lowerCAmelCase  ,			scheduler=_lowerCAmelCase     )
											ldm.to(_lowerCAmelCase     )
											ldm.set_progress_bar_config(disable=_lowerCAmelCase     )

											__magic_name__ 		=			self.dummy_image.to(_lowerCAmelCase     )

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

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

											assert image.shape == (1, 6_4, 6_4, 3)
											__magic_name__ 		=			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   a (	self					:	List[Any]     )       ->				str:



											'''simple docstring'''





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

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

											__magic_name__ 		=			LDMSuperResolutionPipeline(unet=_lowerCAmelCase  ,			vqvae=_lowerCAmelCase  ,			scheduler=_lowerCAmelCase     )
											ldm.to(_lowerCAmelCase     )
											ldm.set_progress_bar_config(disable=_lowerCAmelCase     )

											__magic_name__ 		=			self.dummy_image.to(_lowerCAmelCase     )

											__magic_name__ 		=			ldm(_lowerCAmelCase  ,			num_inference_steps=2  ,			output_type='numpy'     ).images

											assert image.shape == (1, 6_4, 6_4, 3)







@slow
@require_torch
class     __snake_case							(      unittest.TestCase):
					def   a (	self					:	Union[str, Any]     )       ->				Any:



											'''simple docstring'''





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

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

											__magic_name__ 		=			torch.manual_seed(0     )
											__magic_name__ 		=			ldm(image=_lowerCAmelCase  ,			generator=_lowerCAmelCase  ,			num_inference_steps=2_0  ,			output_type='numpy'     ).images

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

											assert image.shape == (1, 2_5_6, 2_5_6, 3)
											__magic_name__ 		=			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





def 							UpperCamelCase_  (				__snake_case			:						int     )->      bool:




						'''simple docstring'''
						if not isinstance(__snake_case	,			__snake_case     ):
												__magic_name__ 		=			F'''Input value of [number={number}] must be an integer'''
												raise TypeError(__snake_case     )
						if number < 0 or number % 2 == 1:
												return False
						__magic_name__ 		=			int(number**0.5     )
						return number == number_sqrt * (number_sqrt + 1)


if __name__ == "__main__":
					import doctest

					doctest.testmod()

