from fastapi import FastAPI
from transformers import AutoProcessor, BarkModel
import torch
import scipy.io.wavfile
import io
from fastapi.responses import StreamingResponse

app = FastAPI()

# Força o uso de CPU e modelo leve para não travar
device = "cpu"
MODEL_ID = "suno/bark-small"

print("Iniciando o download do modelo (pode demorar uns minutos)...")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = BarkModel.from_pretrained(MODEL_ID).to(device)
print("Modelo carregado!")

@app.get("/")
def read_root():
    return {"status": "API Bark Online"}

@app.post("/generate/")
async def generate_audio(text: str):
    try:
        # Prepara o texto
        inputs = processor(text, voice_preset="v2/en_speaker_6").to(device)
        
        # Gera o áudio (limitado para ser mais rápido na CPU)
        audio_array = model.generate(**inputs, coarse_temperature=0.6)
        audio_array = audio_array.cpu().numpy().squeeze()
        
        # Salva em memória
        sample_rate = model.generation_config.sample_rate
        bytes_wav = io.BytesIO()
        scipy.io.wavfile.write(bytes_wav, rate=sample_rate, data=audio_array)
        bytes_wav.seek(0)
        
        return StreamingResponse(bytes_wav, media_type="audio/wav")
    
    except Exception as e:
        return {"error": str(e)}