

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

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= 1
						SCREAMING_SNAKE_CASE__									= 3
						SCREAMING_SNAKE_CASE__									= (3_2, 3_2)

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

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Union[str, Any]:
						torch.manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= 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    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Optional[Any]:
						torch.manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= 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    SCREAMING_SNAKE_CASE			(   self			:  Dict    )				->      Union[str, Any]:
						SCREAMING_SNAKE_CASE__									= """cpu"""
						SCREAMING_SNAKE_CASE__									= self.dummy_uncond_unet
						SCREAMING_SNAKE_CASE__									= DDIMScheduler()
						SCREAMING_SNAKE_CASE__									= self.dummy_vq_model

						SCREAMING_SNAKE_CASE__									= LDMSuperResolutionPipeline(unet=__UpperCAmelCase      ,			vqvae=__UpperCAmelCase      ,			scheduler=__UpperCAmelCase    )
						ldm.to(__UpperCAmelCase    )
						ldm.set_progress_bar_config(disable=__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= self.dummy_image.to(__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= ldm(image=__UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			num_inference_steps=2      ,			output_type="""numpy"""    ).images

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

						assert image.shape == (1, 6_4, 6_4, 3)
						SCREAMING_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    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= self.dummy_uncond_unet
						SCREAMING_SNAKE_CASE__									= DDIMScheduler()
						SCREAMING_SNAKE_CASE__									= self.dummy_vq_model

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

						SCREAMING_SNAKE_CASE__									= LDMSuperResolutionPipeline(unet=__UpperCAmelCase      ,			vqvae=__UpperCAmelCase      ,			scheduler=__UpperCAmelCase    )
						ldm.to(__UpperCAmelCase    )
						ldm.set_progress_bar_config(disable=__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= self.dummy_image.to(__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= ldm(__UpperCAmelCase      ,			num_inference_steps=2      ,			output_type="""numpy"""    ).images

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





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

			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      List[str]:
						SCREAMING_SNAKE_CASE__									= load_image(
						    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
						    """/vq_diffusion/teddy_bear_pool.png"""    )
						SCREAMING_SNAKE_CASE__									= init_image.resize((6_4, 6_4)      ,			resample=PIL_INTERPOLATION["""lanczos"""]    )

						SCREAMING_SNAKE_CASE__									= LDMSuperResolutionPipeline.from_pretrained("""duongna/ldm-super-resolution"""      ,			device_map="""auto"""    )
						ldm.set_progress_bar_config(disable=__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= torch.manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= ldm(image=__UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			num_inference_steps=2_0      ,			output_type="""numpy"""    ).images

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

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




def 		A    (       snake_case__						,			snake_case__      ):




			'''simple docstring'''


			return a + b


if __name__ == "__main__":
						A_							:     List[str]     =	5
						A_							:     Optional[int]     =	6
						print(F'The sum of {a} + {b} is {add(a, b)}')


