audio_diffusion_attacks_forhf/tests/test_dac.py

""" test_dac
    Desc: test_dac is another type of encoder, want to try this in addition or instead of AudioLDM encoder
    Req: Run `pip install descript-audio-codec`
            `python3 -m dac download --model_type 16khz`
"""

import numpy as np 
import pandas as pd 
import matplotlib.pyplot as plt 
import seaborn as sns 
import IPython

import scipy
import torch
import torchaudio
import os
import ast

import dac
from audiotools import AudioSignal


if __name__ == '__main__':
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    data_loc = '/data/robbizorg/music_datasets/fma/'
    example_audio_loc = os.path.join(data_loc, 'data/fma_large/000/000400.mp3')

    # Download a model
    model_path = dac.utils.download(model_type="16khz")
    model = dac.DAC.load(model_path)

    model = model.to(device)

    # Lets try again with their tools
    # Load audio signal file
    signal = AudioSignal(example_audio_loc, sample_rate = 16000)
    signal.resample(16000)
    
    # signal.to(model.device)

    # resamp_audio = torchaudio.functional.resample(signal.audio_data, sr, 16000)

    # x = model.preprocess(signal.audio_data, signal.sample_rate).mean(axis = 1).unsqueeze(0)

    # x = model.preprocess(resamp_audio, 16000)

    # Runs out of space for large file
    # z, codes, latents, _, _ = model.encode(x)

    # Decode audio signal
    # y = model.decode(z)
    signal.write('./assets/audios/example_track.wav')

    x = model.compress(signal)

    y = model.decompress(x)
    y.write('./assets/audios/example_track_decompressed.wav')