

import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class        a__       (   unittest.TestCase       ):







		"""simple docstring"""




		def  _lowercase      (			self    : List[Any]		)					->Optional[Any]:
								"""simple docstring"""







								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


		@property
		def  _lowercase      (			self    : Optional[Any]		)					->Tuple:
								"""simple docstring"""







								torch.manual_seed(0		)
								SCREAMING_SNAKE_CASE :      List[str]       =   UNetaDModel(
								    sample_size=(3_2, 6_4)				,				in_channels=1				,				out_channels=1				,				layers_per_block=2				,				block_out_channels=(1_2_8, 1_2_8)				,				down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")				,				up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")				,				)
								return model


		@property
		def  _lowercase      (			self    : Tuple		)					->List[Any]:
								"""simple docstring"""







								torch.manual_seed(0		)
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   UNetaDConditionModel(
								    sample_size=(6_4, 3_2)				,				in_channels=1				,				out_channels=1				,				layers_per_block=2				,				block_out_channels=(1_2_8, 1_2_8)				,				down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")				,				up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")				,				cross_attention_dim=1_0				,				)
								return model


		@property
		def  _lowercase      (			self    : Optional[Any]		)					->Tuple:
								"""simple docstring"""







								torch.manual_seed(0		)
								SCREAMING_SNAKE_CASE :      Dict       =   AutoencoderKL(
								    sample_size=(1_2_8, 6_4)				,				in_channels=1				,				out_channels=1				,				latent_channels=1				,				layers_per_block=2				,				block_out_channels=(1_2_8, 1_2_8)				,				down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")				,				up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")				,				)
								SCREAMING_SNAKE_CASE :      int       =   UNetaDModel(
								    sample_size=(6_4, 3_2)				,				in_channels=1				,				out_channels=1				,				layers_per_block=2				,				block_out_channels=(1_2_8, 1_2_8)				,				down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")				,				up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")				,				)
								return vqvae, unet




		@slow
		def  _lowercase      (			self    : Optional[int]		)					->int:
								"""simple docstring"""







								SCREAMING_SNAKE_CASE :      Tuple       =   """cpu"""  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE :      List[str]       =   Mel(
								    x_res=self.dummy_unet.config.sample_size[1]				,				y_res=self.dummy_unet.config.sample_size[0]				,				)

								SCREAMING_SNAKE_CASE :      List[str]       =   DDPMScheduler()
								SCREAMING_SNAKE_CASE :      Dict       =   AudioDiffusionPipeline(vqvae=_a				,				unet=self.dummy_unet				,				mel=_a				,				scheduler=_a		)
								SCREAMING_SNAKE_CASE :      Optional[int]       =   pipe.to(_a		)
								pipe.set_progress_bar_config(disable=_a		)

								SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.Generator(device=_a		).manual_seed(4_2		)
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   pipe(generator=_a				,				steps=4		)
								SCREAMING_SNAKE_CASE :      int       =   output.audios[0]
								SCREAMING_SNAKE_CASE :      Dict       =   output.images[0]

								SCREAMING_SNAKE_CASE :      Any       =   torch.Generator(device=_a		).manual_seed(4_2		)
								SCREAMING_SNAKE_CASE :      Tuple       =   pipe(generator=_a				,				steps=4				,				return_dict=_a		)
								SCREAMING_SNAKE_CASE :      List[Any]       =   output[0][0]

								assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
								assert (
								    image.height == self.dummy_unet.config.sample_size[0]
								    and image.width == self.dummy_unet.config.sample_size[1]
								)
								SCREAMING_SNAKE_CASE :      List[str]       =   np.frombuffer(image.tobytes()				,				dtype="""uint8"""		)[:1_0]
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   np.frombuffer(image_from_tuple.tobytes()				,				dtype="""uint8"""		)[:1_0]
								SCREAMING_SNAKE_CASE :      List[Any]       =   np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]		)

								assert np.abs(image_slice.flatten() - expected_slice		).max() == 0
								assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() == 0

								SCREAMING_SNAKE_CASE :      str       =   Mel(
								    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]				,				y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]				,				)

								SCREAMING_SNAKE_CASE :      Any       =   DDIMScheduler()
								SCREAMING_SNAKE_CASE :      Tuple       =   self.dummy_vqvae_and_unet
								SCREAMING_SNAKE_CASE :      List[str]       =   AudioDiffusionPipeline(
								    vqvae=self.dummy_vqvae_and_unet[0]				,				unet=dummy_vqvae_and_unet[1]				,				mel=_a				,				scheduler=_a		)
								SCREAMING_SNAKE_CASE :      List[Any]       =   pipe.to(_a		)
								pipe.set_progress_bar_config(disable=_a		)

								np.random.seed(0		)
								SCREAMING_SNAKE_CASE :      Dict       =   np.random.uniform(-1				,				1				,				((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
								SCREAMING_SNAKE_CASE :      int       =   torch.Generator(device=_a		).manual_seed(4_2		)
								SCREAMING_SNAKE_CASE :      int       =   pipe(raw_audio=_a				,				generator=_a				,				start_step=5				,				steps=1_0		)
								SCREAMING_SNAKE_CASE :      Dict       =   output.images[0]

								assert (
								    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
								    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
								)
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   np.frombuffer(image.tobytes()				,				dtype="""uint8"""		)[:1_0]
								SCREAMING_SNAKE_CASE :      List[Any]       =   np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]		)

								assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

								SCREAMING_SNAKE_CASE :      List[Any]       =   self.dummy_unet_condition
								SCREAMING_SNAKE_CASE :      List[Any]       =   AudioDiffusionPipeline(
								    vqvae=self.dummy_vqvae_and_unet[0]				,				unet=_a				,				mel=_a				,				scheduler=_a		)
								SCREAMING_SNAKE_CASE :      List[str]       =   pipe.to(_a		)
								pipe.set_progress_bar_config(disable=_a		)

								np.random.seed(0		)
								SCREAMING_SNAKE_CASE :      List[str]       =   torch.rand((1, 1, 1_0)		)
								SCREAMING_SNAKE_CASE :      Tuple       =   pipe(generator=_a				,				encoding=_a		)
								SCREAMING_SNAKE_CASE :      Tuple       =   output.images[0]
								SCREAMING_SNAKE_CASE :      Optional[int]       =   np.frombuffer(image.tobytes()				,				dtype="""uint8"""		)[:1_0]
								SCREAMING_SNAKE_CASE :      int       =   np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]		)

								assert np.abs(image_slice.flatten() - expected_slice		).max() == 0







@slow
@require_torch_gpu
class        a__       (   unittest.TestCase       ):







		"""simple docstring"""




		def  _lowercase      (			self    : Dict		)					->Dict:
								"""simple docstring"""







								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




		def  _lowercase      (			self    : Union[str, Any]		)					->Optional[int]:
								"""simple docstring"""







								SCREAMING_SNAKE_CASE :      Tuple       =   torch_device

								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
								SCREAMING_SNAKE_CASE :      Optional[int]       =   pipe.to(_a		)
								pipe.set_progress_bar_config(disable=_a		)

								SCREAMING_SNAKE_CASE :      List[str]       =   torch.Generator(device=_a		).manual_seed(4_2		)
								SCREAMING_SNAKE_CASE :      Dict       =   pipe(generator=_a		)
								SCREAMING_SNAKE_CASE :      Optional[int]       =   output.audios[0]
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   output.images[0]

								assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
								assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   np.frombuffer(image.tobytes()				,				dtype="""uint8"""		)[:1_0]
								SCREAMING_SNAKE_CASE :      Any       =   np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]		)

								assert np.abs(image_slice.flatten() - expected_slice		).max() == 0







from typing import TYPE_CHECKING

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


UpperCAmelCase__ :     List[Any]				    =    {"""configuration_vit""": ["""VIT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViTConfig""", """ViTOnnxConfig"""]}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Union[str, Any]				    =    ["""ViTFeatureExtractor"""]
  UpperCAmelCase__ :     List[Any]				    =    ["""ViTImageProcessor"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     List[Any]				    =    [
      """VIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """ViTForImageClassification""",
      """ViTForMaskedImageModeling""",
      """ViTModel""",
      """ViTPreTrainedModel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Tuple				    =    [
      """TFViTForImageClassification""",
      """TFViTModel""",
      """TFViTPreTrainedModel""",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     str				    =    [
      """FlaxViTForImageClassification""",
      """FlaxViTModel""",
      """FlaxViTPreTrainedModel""",
  ]

if TYPE_CHECKING:
  from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_vit import ViTFeatureExtractor
    from .image_processing_vit import ViTImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_vit import (
        VIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        ViTForImageClassification,
        ViTForMaskedImageModeling,
        ViTModel,
        ViTPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_vit import TFViTForImageClassification, TFViTModel, TFViTPreTrainedModel

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel, FlaxViTPreTrainedModel


else:
  import sys

  UpperCAmelCase__ :     Optional[Any]				    =    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




