app.py

import gradio as gr
import torch
import soundfile as sf
import numpy as np
import yaml
from inference import MasteringStyleTransfer
from utils import download_youtube_audio
from config import args
import pyloudnorm as pyln
import tempfile
import os
import matplotlib.pyplot as plt
import io

mastering_transfer = MasteringStyleTransfer(args)

def denormalize_audio(audio, dtype=np.int16):
    """
    Denormalize the audio from the range [-1, 1] to the full range of the specified dtype.
    """
    if dtype == np.int16:
        audio = np.clip(audio, -1, 1)  # Ensure the input is in the range [-1, 1]
        return (audio * 32767).astype(np.int16)
    elif dtype == np.float32:
        return audio.astype(np.float32)
    else:
        raise ValueError("Unsupported dtype. Use np.int16 or np.float32.")

def loudness_normalize(audio, sample_rate, target_loudness=-12.0):
    # Ensure audio is float32
    if audio.dtype != np.float32:
        audio = audio.astype(np.float32)
    
    # If audio is mono, reshape to (samples, 1)
    if audio.ndim == 1:
        audio = audio.reshape(-1, 1)
    
    meter = pyln.Meter(sample_rate)  # create BS.1770 meter
    loudness = meter.integrated_loudness(audio)
    
    loudness_normalized_audio = pyln.normalize.loudness(audio, loudness, target_loudness)
    return loudness_normalized_audio

def process_audio(input_audio, reference_audio):
    output_audio, predicted_params, _, _, _, sr = mastering_transfer.process_audio(
        input_audio, reference_audio, reference_audio, {}, False
    )
    
    param_output = mastering_transfer.get_param_output_string(predicted_params)
    
    # Convert output_audio to numpy array if it's a tensor
    if isinstance(output_audio, torch.Tensor):
        output_audio = output_audio.cpu().numpy()

    # # Normalize output audio
    # output_audio = loudness_normalize(output_audio, sr)
    
    # Denormalize the audio to int16
    output_audio = denormalize_audio(output_audio, dtype=np.int16)
    
    if output_audio.ndim == 1:
        output_audio = output_audio.reshape(-1, 1)
    elif output_audio.ndim > 2:
        output_audio = output_audio.squeeze()

    # Ensure the audio is in the correct shape (samples, channels)
    if output_audio.shape[1] > output_audio.shape[0]:
        output_audio = output_audio.transpose(1,0)

    print(output_audio.shape)
    print(param_output)
    
    return (sr, output_audio), param_output

def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
    if ito_reference_audio is None:
        ito_reference_audio = reference_audio

    ito_config = {
        'optimizer': optimizer,
        'learning_rate': learning_rate,
        'num_steps': num_steps,
        'af_weights': af_weights,
        'sample_rate': args.sample_rate
    }

    input_tensor = mastering_transfer.preprocess_audio(input_audio, args.sample_rate)
    reference_tensor = mastering_transfer.preprocess_audio(reference_audio, args.sample_rate)
    ito_reference_tensor = mastering_transfer.preprocess_audio(ito_reference_audio, args.sample_rate)

    initial_reference_feature = mastering_transfer.get_reference_embedding(reference_tensor)

    ito_log = ""
    loss_values = []
    for log_entry, current_output, current_params, step, loss in mastering_transfer.inference_time_optimization(
        input_tensor, ito_reference_tensor, ito_config, initial_reference_feature
    ):
        ito_log += log_entry
        ito_param_output = mastering_transfer.get_param_output_string(current_params)
        loss_values.append(loss)
        
        # Convert current_output to numpy array if it's a tensor
        if isinstance(current_output, torch.Tensor):
            current_output = current_output.cpu().numpy()
        
        # Normalize output audio
        current_output = loudness_normalize(current_output, args.sample_rate)
        
        # Denormalize the audio to int16
        current_output = denormalize_audio(current_output, dtype=np.int16)
        
        # Ensure the audio is in the correct shape (samples, channels)
        if current_output.ndim == 1:
            current_output = current_output.reshape(-1, 1)
        elif current_output.ndim > 2:
            current_output = current_output.squeeze()
        
        yield (args.sample_rate, current_output), ito_param_output, step, ito_log, loss_values

def plot_loss_curve(loss_values):
    plt.figure(figsize=(10, 6))
    plt.plot(loss_values)
    plt.title('ITO Loss Curve')
    plt.xlabel('Step')
    plt.ylabel('Loss')
    plt.grid(True)
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    return buf

""" APP display """
with gr.Blocks() as demo:
    gr.Markdown("# Mastering Style Transfer Demo")

    with gr.Tab("Upload Audio"):
        with gr.Row():
            input_audio = gr.Audio(label="Input Audio")
            reference_audio = gr.Audio(label="Reference Audio")
        
        process_button = gr.Button("Process Mastering Style Transfer")
        
        with gr.Row():
            output_audio = gr.Audio(label="Output Audio", type='numpy')
            param_output = gr.Textbox(label="Predicted Parameters", lines=5)

        process_button.click(
            process_audio, 
            inputs=[input_audio, reference_audio], 
            outputs=[output_audio, param_output]
        )

    gr.Markdown("## Inference Time Optimization (ITO)")
    
    with gr.Row():
        ito_reference_audio = gr.Audio(label="ITO Reference Audio (optional)")
        with gr.Column():
            num_steps = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of Steps")
            optimizer = gr.Dropdown(["Adam", "RAdam", "SGD"], value="RAdam", label="Optimizer")
            learning_rate = gr.Slider(minimum=0.0001, maximum=0.1, value=0.001, step=0.0001, label="Learning Rate")
            af_weights = gr.Textbox(label="AudioFeatureLoss Weights (comma-separated)", value="0.1,0.001,1.0,1.0,0.1")
        
    ito_button = gr.Button("Perform ITO")

    with gr.Row():
        with gr.Column():
            ito_output_audio = gr.Audio(label="ITO Output Audio")
            ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=15)
        with gr.Column():
            ito_steps_taken = gr.Number(label="ITO Steps Taken")
            ito_loss_plot = gr.Image(label="ITO Loss Curve")
            ito_log = gr.Textbox(label="ITO Log", lines=10)

    def run_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
        af_weights = [float(w.strip()) for w in af_weights.split(',')]
        ito_generator = perform_ito(
            input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
        )
        
        # Initialize variables to store the final results
        final_audio = None
        final_params = None
        final_steps = 0
        final_log = ""
        
        # Iterate through the generator to get the final results
        for audio, params, steps, log, losses in ito_generator:
            final_audio = audio
            final_params = params
            final_steps = steps
            final_log = log
            loss_values = losses

        loss_plot = plot_loss_curve(loss_values)
        
        return final_audio, final_params, final_steps, final_log, loss_plot

    ito_button.click(
        run_ito,
        inputs=[input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights],
        outputs=[ito_output_audio, ito_param_output, ito_steps_taken, ito_log, ito_loss_plot]
    )

demo.launch()