




'''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       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            __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       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




            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       lowerCAmelCase      (   self  )      ->  Optional[int]:



            '''simple docstring'''




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



            '''simple docstring'''




            __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       lowerCAmelCase      (   self  )      ->  str:



            '''simple docstring'''




            __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):


       def       lowerCAmelCase      (   self  )      ->  List[Any]:



            '''simple docstring'''




            __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, 256, 256, 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'''



from transformers import ProcessorMixin





class     lowerCAmelCase   ( __lowerCAmelCase):
       __lowercase  :      Union[str, Any]        =			'''AutoFeatureExtractor'''
       __lowercase  :      List[Any]        =			'''AutoTokenizer'''






