generation_utilities.py

import torch
from IPython.display import Audio
from audiodiffusion import AudioDiffusion, AudioDiffusionPipeline
from audiodiffusion.audio_encoder import AudioEncoder
import librosa
import librosa.display
import numpy as np
import random 

device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = ["SAint7579/orpheus_ldm_model_v1-0", "teticio/audio-diffusion-ddim-256"]

audio_diffusion_v0 = AudioDiffusionPipeline.from_pretrained("teticio/latent-audio-diffusion-256").to(device)
audio_diffusion_v1 = AudioDiffusionPipeline.from_pretrained("SAint7579/orpheus_ldm_model_v1-0").to(device)
audio_diffusion_v2 = AudioDiffusionPipeline.from_pretrained("SAint7579/orpheus_ldm_model_v2-0").to(device)
ddim = AudioDiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256").to(device)

### Add numpy docstring to generate_from_music
def generate_from_music(song_array, diffuser, start_step, total_steps=100, device=device):
    """
    Generates audio from a given song array using a given diffuser.
    Parameters
    ----------
    song_array : numpy.ndarray
        The song array to use as the raw audio.
    diffuser : AudioDiffusionPipeline
        The diffuser to use to generate the audio.
    start_step : int
        The step to start generating from.
    total_steps : int
        The total number of steps to generate.
    device : str
        The device to use for generation.
    Returns
    -------
    numpy.ndarray
        The generated audio.
    """
    generator = torch.Generator(device=device)
    generator.seed()
    output = diffuser(raw_audio=song_array, generator = generator, start_step=start_step, steps=total_steps)
    return output.images[0], output.audios[0, 0]

def generate_from_music_long(song_array, diffuser, start_step, total_steps=100, device=device):
    """
    Generates a 10 second audio from a given song array using a given diffuser.
    Parameters
    ----------
    song_array : numpy.ndarray
        The song array to use as the raw audio.
    diffuser : AudioDiffusionPipeline
        The diffuser to use to generate the audio.
    start_step : int
        The step to start generating from.
    total_steps : int
        The total number of steps to generate.
    device : str
        The device to use for generation.
    Returns
    -------
    numpy.ndarray
        The generated audio.
    """
    generator = torch.Generator(device=device)
    generator.seed()
    output = diffuser(raw_audio=song_array, generator = generator, start_step=start_step, steps=total_steps)

    # Get the track and use the diffuser again to create the continuation
    track = output.audios[0, 0]
    sample_rate = diffuser.mel.get_sample_rate()
    overlap_secs = 2
    overlap_samples = overlap_secs * sample_rate

    continue_output = diffuser(raw_audio=track[-overlap_samples:],
                               generator=generator,
                               start_step=start_step,
                               mask_start_secs=overlap_secs)
    # image2 = output.images[0]
    audio2 = continue_output.audios[0, 0]
    track = np.concatenate([track, audio2[overlap_samples:]])

    return output.images[0], track


## Add docstring to iterative_slerp function in numpy format
def iterative_slerp(song_arrays, ddim, steps=10):
    """Iterative slerp function.

    Parameters
    ----------
    song_arrays : list
        List of song arrays to slerp.
    ddim : AudioDiffusion ddim model
        AudioDiffusion object.

    Returns
    -------
    slerp : torch.Tensor
        Slerped tensor.
    """
    noise = []
    for arr in song_arrays:
        ddim.mel.audio = arr
        noise.append(ddim.encode([ddim.mel.audio_slice_to_image(0)], steps=steps))

    slerp = noise[0]
    for i in range(1, len(noise)):
        slerp = ddim.slerp(slerp, noise[i], 0.5)

    return slerp

def merge_songs(song_arrays, ddim, slerp_steps=10, diffusion_steps=100, device=device):
    """Merge songs.

    Parameters
    ----------
    song_arrays : list
        List of song arrays to merge.
    ddim : AudioDiffusion ddim model
        AudioDiffusion object.

    Returns
    -------
    spectrogram : np.ndarray
        Merged spectrogram.
    audio : np.ndarray
        Merged audio.
    """
    generator = torch.Generator(device=device)
    generator.manual_seed(7579)
    slerp = iterative_slerp(song_arrays, ddim, slerp_steps)
    merged = ddim(noise=slerp, generator=generator, steps=diffusion_steps)
    return merged.images[0], merged.audios[0, 0]

## Write generate songs function with numpy docstring
def generate_songs(conditioning_songs, similarity=0.9, quality=500, merging_quality=100, device=device):
    """Generate songs.

    Parameters
    ----------
    conditioning_songs : list
        List of conditioning songs.
    similarity : float
        Similarity between conditioning songs.
    quality : int
        Quality of generated song.

    Returns
    -------
    spec_generated : np.ndarray
        Spectrogram of generated song.
    generated : np.ndarray
        Generated song.
    """
    ## Merging songs
    print("Merging songs...")
    if len(conditioning_songs)>1:
        # print(conditioning_songs)
        # for c in conditioning_songs:
        #     print(len(c))
        spec_merged, merged = merge_songs(conditioning_songs, ddim, slerp_steps=merging_quality, diffusion_steps=merging_quality, device=device)
    else:
        merged = conditioning_songs[0]

    ## Take a random 10 second slice from the merged song
    # sample_rate = ddim.mel.get_sample_rate()
    # start = np.random.randint(0, len(merged) - 5 * sample_rate)
    # merged = merged[start:start + 5 * sample_rate]



    diffuser, model_name = random.choice([
        (audio_diffusion_v0, "v0"),
        (audio_diffusion_v1, "v1"),
        (audio_diffusion_v2, "v2")
    ])

    print("Generating song...")
    ## quality = X - similarity*X
    total_steps = min([1000, int(quality/(1-similarity))])
    start_step = int(total_steps*similarity)
    spec_generated, generated = generate_from_music(merged, diffuser, start_step=start_step, total_steps=total_steps, device=device)

    return spec_generated, generated, model_name

# if __name__ == '__main__':
#     song1 = "D:/Projects/Orpheus_ai/DataSet/Sep_Dataset/accompaniment/000010.mp3"
#     song2 = "D:/Projects/Orpheus_ai/DataSet/Sep_Dataset/accompaniment/000002.mp3"
#     song3 = "D:/Projects/Orpheus_ai/DataSet/Sep_Dataset/accompaniment/000003.mp3"

#     song_array_1, sr = librosa.load(song1, sr=22050)
#     song_array_1 = song_array_1[:sr*5]

#     song_array_2, sr = librosa.load(song2, sr=22050)
#     song_array_2 = song_array_2[:sr*10]

#     song_array_3, sr = librosa.load(song3, sr=22050)
#     song_array_3 = song_array_3[:sr*10]

#     generator = torch.Generator(device=device)

#     # seed = 239150437427
#     image, audio = generate_songs([song_array_2, song_array_3], similarity=0.5, quality=10, device=device)
#     Audio(audio, rate=sr)