import gradio as gr
import torch
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import numpy as np
import tempfile
import os

# Load the MusicGen model
model = MusicGen.get_pretrained('small')
model.set_generation_params(duration=30)  # Set maximum duration to 30 seconds

def enhance_audio(audio_file):
    # Load and process the audio file
    waveform = model.compression_model.encode(audio_file)
    
    # Apply AI-based enhancement (this is a simplified example)
    enhanced_waveform = model.compression_model.decode(waveform)
    
    # Convert to numpy array and normalize
    enhanced_audio = enhanced_waveform.squeeze().cpu().numpy()
    enhanced_audio = enhanced_audio / np.max(np.abs(enhanced_audio))
    
    # Save the enhanced audio to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
        audio_write(temp_file.name, enhanced_audio, model.sample_rate, strategy="loudness", loudness_compressor=True)
        output_path = temp_file.name
    
    return output_path

# Create the Gradio interface
iface = gr.Interface(
    fn=enhance_audio,
    inputs=gr.Audio(type="filepath", label="Upload your audio file"),
    outputs=gr.Audio(type="filepath", label="Enhanced Audio"),
    title="AI Music Mastering and Enhancement",
    description="Upload an audio file to apply AI-based mastering and enhancement.",
)

# Launch the app
iface.launch()