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







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       1
          _UpperCAmelCase 						=       3
          _UpperCAmelCase 						=       (32, 32)

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







       @property
       def  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)
          _UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

          torch.manual_seed(0		)
          _UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase 						=       "cpu"
          _UpperCAmelCase 						=       self.dummy_uncond_unet
          _UpperCAmelCase 						=       DDIMScheduler()
          _UpperCAmelCase 						=       self.dummy_vq_model

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

          _UpperCAmelCase 						=       self.dummy_image.to(a_		)

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

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

          assert image.shape == (1, 64, 64, 3)
          _UpperCAmelCase 						=       np.array([0.8_6_7_8, 0.8_2_4_5, 0.6_3_8_1, 0.6_8_3_0, 0.4_3_8_5, 0.5_5_9_9, 0.4_6_4_1, 0.6_2_0_1, 0.5_1_5_0]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2







       @unittest.skipIf(torch_device != "cuda"				,"This test requires a GPU"		)
       def  _snake_case      (       self		):



          '''simple docstring'''

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

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

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

          _UpperCAmelCase 						=       self.dummy_image.to(a_		)

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

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






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







       def  _snake_case      (       self		):



          '''simple docstring'''

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

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

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

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

          assert image.shape == (1, 256, 256, 3)
          _UpperCAmelCase 						=       np.array([0.7_6_4_4, 0.7_6_7_9, 0.7_6_4_2, 0.7_6_3_3, 0.7_6_6_6, 0.7_5_6_0, 0.7_4_2_5, 0.7_2_5_7, 0.6_9_0_7]		)

          assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2









"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase											=     {"""configuration_pegasus""": ["""PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PegasusConfig"""]}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizer"""]

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """PegasusForCausalLM""",
         """PegasusForConditionalGeneration""",
         """PegasusModel""",
         """PegasusPreTrainedModel""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """TFPegasusForConditionalGeneration""",
         """TFPegasusModel""",
         """TFPegasusPreTrainedModel""",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """FlaxPegasusForConditionalGeneration""",
         """FlaxPegasusModel""",
         """FlaxPegasusPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus import PegasusTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus_fast import PegasusTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_pegasus import (
              PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
              PegasusForCausalLM,
              PegasusForConditionalGeneration,
              PegasusModel,
              PegasusPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_pegasus import (
              FlaxPegasusForConditionalGeneration,
              FlaxPegasusModel,
              FlaxPegasusPreTrainedModel,
          )

else:
     import sys

     UpperCamelCase											=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


