import gradio as gr
from transformers import pipeline
import torch
import numpy as np
import scipy.io.wavfile
import os

# --- AI Model Setup ---
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(f"Device set to use {device}")
pipe = pipeline("text-to-audio", "facebook/musicgen-small", device=device)

# --- Music Generation Function ---
def generate_music(prompt, duration):
    try:
        max_new_tokens = int(duration * 50)
        music = pipe(prompt, forward_params={"max_new_tokens": max_new_tokens})
        
        sampling_rate = music["sampling_rate"]
        audio_data = music["audio"][0].T
        
        output_filename = "music_output.wav"
        audio_int16 = np.int16(audio_data * 32767)
        scipy.io.wavfile.write(output_filename, rate=sampling_rate, data=audio_int16)
        
        return output_filename
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

# --- Gradio Interface ---
iface = gr.Interface(
    fn=generate_music,
    inputs=[
        gr.Textbox(lines=2, label="Music Description", placeholder="e.g., A futuristic synthwave track..."),
        gr.Slider(minimum=2, maximum=60, value=15, label="Duration (seconds)")
    ],
    outputs=gr.Audio(type="filepath", label="Generated Music", downloadable=True),
    title="MusicGen AI",
    description="Transform your ideas into music with AI. Describe a sound and set the duration to get started.",
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)