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







		"""simple docstring"""






		@property
		def 		UpperCamelCase__							( self					)			->  List[str]:
								_lowerCAmelCase					    =1
								_lowerCAmelCase					    =3
								_lowerCAmelCase					    =(32, 32)

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






		@property
		def 		UpperCamelCase__							( self					)			->  Union[str, Any]:
								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 		UpperCamelCase__							( self					)			->  Dict:
								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 		UpperCamelCase__							( self					)			->  int:
								_lowerCAmelCase					    ='cpu'
								_lowerCAmelCase					    =self.dummy_uncond_unet
								_lowerCAmelCase					    =DDIMScheduler()
								_lowerCAmelCase					    =self.dummy_vq_model

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

								_lowerCAmelCase					    =self.dummy_image.to(__A					)

								_lowerCAmelCase					    =torch.Generator(device=__A					).manual_seed(0					)
								_lowerCAmelCase					    =ldm(image=__A		,    generator=__A		,    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_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					)			->  Any:
								_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=__A		,    vqvae=__A		,    scheduler=__A					)
								ldm.to(__A					)
								ldm.set_progress_bar_config(disable=__A					)

								_lowerCAmelCase					    =self.dummy_image.to(__A					)

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

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

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







		"""simple docstring"""






		def 		UpperCamelCase__							( self					)			->  Union[str, Any]:
								_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=__A					)

								_lowerCAmelCase					    =torch.manual_seed(0					)
								_lowerCAmelCase					    =ldm(image=__A		,    generator=__A		,    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_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'''

from ..utils import DummyObject, requires_backends
class 	SCREAMING_SNAKE_CASE      (    metaclass=__lowercase):







		"""simple docstring"""
		lowercase :   List[Any]							=				['tensorflow_text']






		def __init__( self		,    *__A		,    **__A					)			->  List[Any]:
								requires_backends(self		,    ['tensorflow_text']					)





