from flask import jsonify, send_file, request
from main import *
import torch, soundfile as sf, numpy as np, io, base64

def generate_music(prompt, output_path="output_music.wav"):
    if musicgen_model is None: return {"error": "Music generation model not initialized."}
    attributes = [prompt]; sample_rate = 32000; duration = 8
    audio_values = musicgen_model.sample(attributes=attributes, sample_rate=sample_rate, duration=duration); output_audio = audio_values.cpu().numpy().squeeze()
    sf.write(output_path, output_audio, sample_rate); return output_path

def musicgen_api(prompt):
    output_file = generate_music(prompt)
    if isinstance(output_file, dict) and "error" in output_file: return {"error": output_file["error"]}
    with open(output_file, 'rb') as f: audio_content = f.read()
    audio_base64 = base64.b64encode(audio_content).decode('utf-8'); os.remove(output_file); return {"audio_base64": audio_base64, "mimetype": "audio/wav"}