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







   @property
   def    UpperCAmelCase__	( self)   ->Optional[Any]:
      a_						   =     1
      a_						   =     3
      a_						   =     (32, 32)

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







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

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

      a_						   =     self.dummy_image.to(__UpperCAmelCase)

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

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

      assert image.shape == (1, 64, 64, 3)
      a_						   =     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)   ->List[str]:
      a_						   =     self.dummy_uncond_unet
      a_						   =     DDIMScheduler()
      a_						   =     self.dummy_vq_model

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

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

      a_						   =     self.dummy_image.to(__UpperCAmelCase)

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

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






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







   def    UpperCAmelCase__	( self)   ->Dict:
      a_						   =     load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/vq_diffusion/teddy_bear_pool.png")
      a_						   =     init_image.resize((64, 64)       ,							resample=PIL_INTERPOLATION["lanczos"])

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

      a_						   =     torch.manual_seed(0)
      a_						   =     ldm(image=__UpperCAmelCase       ,							generator=__UpperCAmelCase       ,							num_inference_steps=20       ,							output_type="numpy").images

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

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






class 		_lowerCAmelCase:

  """simple docstring"""







  def __init__(   self      )->   None:
     # numbermap keys are the number and its values are lists of indexes sorted
     # in ascending order
     __A    =  {}
     # indexmap keys are an index and it's values are the number at that index
     __A    =  {}

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   list[int]:
     if isinstance(UpperCAmelCase    ,							(range, str)      ):
        __A    =  list(UpperCAmelCase      )
     elif not isinstance(UpperCAmelCase    ,							UpperCAmelCase      ):
        raise TypeError(
            '''binary_search_delete() only accepts either a list, range or str'''      )

     __A    =  0
     __A    =  len(UpperCAmelCase      ) - 1

     while low <= high:
        __A    =  (low + high) // 2
        if array[mid] == item:
           array.pop(UpperCAmelCase      )
           return array
        elif array[mid] < item:
           __A    =  mid + 1
        else:
           __A    =  mid - 1
     raise ValueError(
         '''Either the item is not in the array or the array was unsorted'''      )

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   list[int]:
     if isinstance(UpperCAmelCase    ,							(range, str)      ):
        __A    =  list(UpperCAmelCase      )
     elif not isinstance(UpperCAmelCase    ,							UpperCAmelCase      ):
        raise TypeError(
            '''binary_search_insert() only accepts either a list, range or str'''      )

     __A    =  0
     __A    =  len(UpperCAmelCase      ) - 1

     while low <= high:
        __A    =  (low + high) // 2
        if array[mid] == index:
           # If the item already exists in the array,
           # insert it after the existing item
           array.insert(mid + 1    ,							UpperCAmelCase      )
           return array
        elif array[mid] < index:
           __A    =  mid + 1
        else:
           __A    =  mid - 1

        # If the item doesn't exist in the array, insert it at the appropriate position
     array.insert(UpperCAmelCase    ,							UpperCAmelCase      )
     return array

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   None:
     # Remove previous index
     if index in self.indexmap:
        __A    =  self.indexmap[index]
        if len(self.numbermap[n]      ) == 1:
           del self.numbermap[n]
        else:
           __A    =  self.binary_search_delete(self.numbermap[n]    ,							UpperCAmelCase      )

        # Set new index
     __A    =  number

     # Number not seen before or empty so insert number value
     if number not in self.numbermap:
        __A    =  [index]

     # Here we need to perform a binary search insertion in order to insert
     # The item in the correct place
     else:
        __A    =  self.binary_search_insert(
            self.numbermap[number]    ,							UpperCAmelCase      )





  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase      )->   int:
     return self.numbermap.get(UpperCAmelCase    ,							[-1]      )[0]


if __name__ == "__main__":
    import doctest

    doctest.testmod()

