





"""simple docstring"""

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    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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""")      ,       )
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


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

       __lowercase									=	DDPMScheduler()
       __lowercase									=	AudioDiffusionPipeline(vqvae=_a      ,       unet=self.dummy_unet      ,       mel=_a      ,       scheduler=_a							)
       __lowercase									=	pipe.to(_a							)
       pipe.set_progress_bar_config(disable=_a							)

       __lowercase									=	torch.Generator(device=_a							).manual_seed(4_2							)
       __lowercase									=	pipe(generator=_a      ,       steps=4							)
       __lowercase									=	output.audios[0]
       __lowercase									=	output.images[0]

       __lowercase									=	torch.Generator(device=_a							).manual_seed(4_2							)
       __lowercase									=	pipe(generator=_a      ,       steps=4      ,       return_dict=_a							)
       __lowercase									=	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]
       )
       __lowercase									=	np.frombuffer(image.tobytes()      ,       dtype="""uint8"""							)[:1_0]
       __lowercase									=	np.frombuffer(image_from_tuple.tobytes()      ,       dtype="""uint8"""							)[:1_0]
       __lowercase									=	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

       __lowercase									=	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]      ,       )

       __lowercase									=	DDIMScheduler()
       __lowercase									=	self.dummy_vqvae_and_unet
       __lowercase									=	AudioDiffusionPipeline(
           vqvae=self.dummy_vqvae_and_unet[0]      ,       unet=dummy_vqvae_and_unet[1]      ,       mel=_a      ,       scheduler=_a							)
       __lowercase									=	pipe.to(_a							)
       pipe.set_progress_bar_config(disable=_a							)

       np.random.seed(0							)
       __lowercase									=	np.random.uniform(-1      ,       1      ,       ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)							)
       __lowercase									=	torch.Generator(device=_a							).manual_seed(4_2							)
       __lowercase									=	pipe(raw_audio=_a      ,       generator=_a      ,       start_step=5      ,       steps=1_0							)
       __lowercase									=	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]
       )
       __lowercase									=	np.frombuffer(image.tobytes()      ,       dtype="""uint8"""							)[:1_0]
       __lowercase									=	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

       __lowercase									=	self.dummy_unet_condition
       __lowercase									=	AudioDiffusionPipeline(
           vqvae=self.dummy_vqvae_and_unet[0]      ,       unet=_a      ,       mel=_a      ,       scheduler=_a							)
       __lowercase									=	pipe.to(_a							)
       pipe.set_progress_bar_config(disable=_a							)

       np.random.seed(0							)
       __lowercase									=	torch.rand((1, 1, 1_0)							)
       __lowercase									=	pipe(generator=_a      ,       encoding=_a							)
       __lowercase									=	output.images[0]
       __lowercase									=	np.frombuffer(image.tobytes()      ,       dtype="""uint8"""							)[:1_0]
       __lowercase									=	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    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       __lowercase									=	torch_device

       __lowercase									=	DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""							)
       __lowercase									=	pipe.to(_a							)
       pipe.set_progress_bar_config(disable=_a							)

       __lowercase									=	torch.Generator(device=_a							).manual_seed(4_2							)
       __lowercase									=	pipe(generator=_a							)
       __lowercase									=	output.audios[0]
       __lowercase									=	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]
       __lowercase									=	np.frombuffer(image.tobytes()      ,       dtype="""uint8"""							)[:1_0]
       __lowercase									=	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




"""simple docstring"""

import random


def lowerCAmelCase_      (				UpperCamelCase__	: int				,   UpperCamelCase__	: float				,   UpperCamelCase__	: bool = False ):






 """simple docstring"""







 __lowercase									=	{i: [] for i in range(UpperCamelCase__ )}

 # if probability is greater or equal than 1, then generate a complete graph
 if probability >= 1:
  return complete_graph(UpperCamelCase__ )
 # if probability is lower or equal than 0, then return a graph without edges
 if probability <= 0:
  return graph

 # for each couple of nodes, add an edge from u to v
 # if the number randomly generated is greater than probability probability
 for i in range(UpperCamelCase__ ):
  for j in range(i + 1				,   UpperCamelCase__ ):
   if random.random() < probability:
    graph[i].append(UpperCamelCase__ )
    if not directed:
     # if the graph is undirected, add an edge in from j to i, either
     graph[j].append(UpperCamelCase__ )
 return graph


def lowerCAmelCase_      (				UpperCamelCase__	: int ):






 """simple docstring"""







 return {
     i: [j for j in range(UpperCamelCase__ ) if i != j] for i in range(UpperCamelCase__ )
 }


if __name__ == "__main__":
       import doctest

       doctest.testmod()
