






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


      def  __UpperCAmelCase		(	self	)->	Dict:
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


      @property
      def  __UpperCAmelCase		(	self	)->	List[str]:
           torch.manual_seed(0	)
           __lowerCAmelCase								=					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  __UpperCAmelCase		(	self	)->	List[str]:
           torch.manual_seed(0	)
           __lowerCAmelCase								=					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  __UpperCAmelCase		(	self	)->	Optional[int]:
           torch.manual_seed(0	)
           __lowerCAmelCase								=					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")					,						)
           __lowerCAmelCase								=					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  __UpperCAmelCase		(	self	)->	Optional[int]:
           __lowerCAmelCase								=					"cpu"  # ensure determinism for the device-dependent torch.Generator
           __lowerCAmelCase								=					Mel(
               x_res=self.dummy_unet.config.sample_size[1]					,						y_res=self.dummy_unet.config.sample_size[0]					,						)

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

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

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

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

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

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

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

           np.random.seed(0	)
           __lowerCAmelCase								=					torch.rand((1, 1, 1_0)	)
           __lowerCAmelCase								=					pipe(generator=__UpperCamelCase					,						encoding=__UpperCamelCase	)
           __lowerCAmelCase								=					output.images[0]
           __lowerCAmelCase								=					np.frombuffer(image.tobytes()					,						dtype="uint8"	)[:1_0]
           __lowerCAmelCase								=					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 							_UpperCamelCase			(unittest.TestCase			):


      def  __UpperCAmelCase		(	self	)->	Any:
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


      def  __UpperCAmelCase		(	self	)->	Dict:
           __lowerCAmelCase								=					torch_device

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

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





def       __lowerCAmelCase      (						__snake_case					, __snake_case					, __snake_case       ):
     __lowerCAmelCase								=					list(range(len(__snake_case       )       )       )
     __lowerCAmelCase								=					[v / w for v, w in zip(__snake_case					, __snake_case       )]
     index.sort(key=lambda __snake_case       : ratio[i]					, reverse=__snake_case       )

     __lowerCAmelCase								=					0
     __lowerCAmelCase								=					[0] * len(__snake_case       )
     for i in index:
          if weight[i] <= capacity:
               __lowerCAmelCase								=					1
               max_value += value[i]
               capacity -= weight[i]
          else:
               __lowerCAmelCase								=					capacity / weight[i]
               max_value += value[i] * capacity / weight[i]
               break

     return max_value, fractions


if __name__ == "__main__":
 import doctest

 doctest.testmod()






