api.py

from fastapi import FastAPI, HTTPException, Response, Depends, File, Form
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import soundfile as sf
import numpy as np
from voxcpm import VoxCPM
from pydantic import BaseModel
import os
import requests
import zipfile
from utils import *
import uuid
import queue
import threading
import asyncio
import time


security = HTTPBearer()
app = FastAPI()

# 验证函数
def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    # 从环境变量获取token
    expected_token = os.getenv("API_TOKEN", "my_secret_token")
    if credentials.credentials != expected_token:
        raise HTTPException(status_code=401, detail="Invalid or missing token")
    return credentials.credentials

class GenerateRequest(BaseModel):
    text: str
    voice: str
    cfg_value: float = 2.0
    inference_timesteps: int = 10
    do_normalize: bool = True
    denoise: bool = True

def download_voices(bForce=False):
    # 检查 /workspace/voices/ 目录中是否有 .pmt 文件
    voices_dir = "/workspace/voices"
    if not os.path.exists(voices_dir):
        os.makedirs(voices_dir)
    
    # 改进：递归查找所有子目录中的.pmt文件，找到一个就返回
    pmt_file_found = None
    for root, dirs, files in os.walk(voices_dir):
        for file in files:
            if file.endswith(".pmt"):
                pmt_file_found = file
                break
        if pmt_file_found:
            break
    
    # 如果没有找到.pmt文件且需要强制下载
    if bForce or pmt_file_found is None:
        # 如果没有 .pmt 文件，尝试从远程下载
        voice_download_url = os.getenv("VOICE_DOWNLOAD_URL")
        
        if voice_download_url:
            try:
                response = requests.get(voice_download_url)
                response.raise_for_status()
                
                # 保存下载的zip文件
                zip_path = f"{voices_dir}/voices.zip"
                with open(zip_path, "wb") as f:
                    f.write(response.content)
                
                # 解压zip文件
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall(voices_dir)
                
                # 删除临时zip文件
                os.remove(zip_path)
                
            except Exception as e:
                print_with_time(f"Failed to download and extract voices: {e}")
                raise HTTPException(status_code=500, detail="Failed to download voice files")

# 队列相关变量
task_queue = queue.Queue()
output_dir = "./output"
max_output_files = 10

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

cleanup_interval = 24 * 60 * 60  # 24小时，以秒为单位
def cleanup_thread():
    while True:
        try:
            current_time = time.time()
            cutoff_time = current_time - cleanup_interval
            deleted_count = 0
            
            # 先获取所有wav文件列表，避免在遍历中删除文件导致问题
            wav_files = []
            for filename in os.listdir(output_dir):
                wav_files.append(filename)
            
            for filename in wav_files:
                filepath = os.path.join(output_dir, filename)
                try:
                    ctime = os.path.getctime(filepath)
                    if ctime < cutoff_time:
                        os.remove(filepath)
                        deleted_count += 1
                except Exception as e:
                    print_with_time(f"[Error] deleting file {filename}: {e}")
            if deleted_count > 0:
                print_with_time(f"--Cleaned up {deleted_count} old files--")
        except Exception as e:
            print_with_time(f"[Error] cleaning up old files: {e}")
        finally:
            time.sleep(cleanup_interval)

async def process_queue():
    print_with_time("Loading VoxCPM model...")
    model = VoxCPM.from_pretrained("openbmb/VoxCPM-0.5B")
    print_with_time("VoxCPM model loaded.")

    while True:
        try:
            task_data = task_queue.get_nowait()
            request = task_data["request"]
            text = (request.text or "").strip()
            if len(text) == 0:
                continue

            if model is None:
                raise RuntimeError("Failed to initialize model")
            
            download_voices()
            print_with_time(f"Generating audio for : '{text[:60]}...'")
            with open(f"./voices/{request.voice}.pmt", 'r', encoding='utf-8') as f:
                wav = model.generate(
                    text=text,
                    prompt_wav_path=f"./voices/{request.voice}.wav",
                    prompt_text=f.read(),
                    cfg_value=request.cfg_value,
                    inference_timesteps=request.inference_timesteps,
                    normalize=request.do_normalize,
                    denoise=request.denoise
                )
                sf.write(os.path.join(output_dir, f"{task_data['task_id']}.wav"), wav, 16000)

            task_queue.task_done()
            print_with_time("audio generated.")
            await asyncio.sleep(0.6)
        except queue.Empty:
            await asyncio.sleep(0.6)
        except Exception as e:
            print_with_time(f"Error processing queue item: {e}")
            await asyncio.sleep(0.6)


@app.post("/generate")
async def generate_tts_async(request: GenerateRequest, token: str = Depends(verify_token)):
    task_id = str(uuid.uuid4())
    
    # 将任务添加到队列
    task_data = {"task_id": task_id, "request": request}
    task_queue.put(task_data)
    
    return {"task_id": task_id}

@app.get("/tts")
async def get_generate_result(task_id: str, token: str = Depends(verify_token)):
    filepath = os.path.join(output_dir, f"{task_id}.wav")
    
    if not os.path.exists(filepath):
        raise HTTPException(status_code=404, detail="Result file not found")
    try:
        with open(filepath, 'rb') as f:
            content = f.read()
        return Response(content=content, media_type="audio/wav")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to read result file: {str(e)}")

@app.post("/upload_voice")
def upload_voice(name: str = Form(...), wav: bytes = File(...), prompt: str = Form(...), token: str = Depends(verify_token)):
    # 保存wav文件
    with open(f"/workspace/voices/{name}.wav", 'wb') as f:
        f.write(wav)
    
    # 保存pmt文件
    with open(f"/workspace/voices/{name}.pmt", 'w', encoding='utf-8') as f:
        f.write(prompt)
    
    return {"status": "success"}

@app.delete("/delete_voice")
def delete_voice(name: str, token: str = Depends(verify_token)):
    wav_file = f"/workspace/voices/{name}.wav"
    pmt_file = f"/workspace/voices/{name}.pmt"
    
    # 检查文件是否存在
    if os.path.exists(wav_file):
        os.remove(wav_file)
        if os.path.exists(pmt_file):
            os.remove(pmt_file)
        return {"status": "success"}
    else:
        return {"status": "不存在"}

@app.get("/voices")
def get_voices(token: str = Depends(verify_token)):
    download_voices()
    # 获取所有 .pmt 文件（递归搜索子目录）
    pmt_files = []
    for root, dirs, files in os.walk("./voices"):
        for file in files:
            if file.endswith(".pmt"):
                # 获取文件相对于当前工作目录的完整相对路径
                full_path = os.path.join(root, file)
                relative_path = os.path.relpath(full_path, "./voices")
                pmt_files.append(relative_path[:-4])
    
    # 确保对应的 .wav 文件也存在
    valid_voices = []
    for voice in pmt_files:
        if os.path.exists(f"./voices/{voice}.wav"):
            valid_voices.append(voice)
    return {"voices": valid_voices}

@app.post("/re_download_voices")
def download_voices_file(token: str = Depends(verify_token)):
    download_voices(True)
    return {"status": "success"}


# ↓↓↓↓↓↓↓↓↓无需验证↓↓↓↓↓↓↓↓
@app.get("/")
@app.get("/health")
def health_check():
    return {"status": "health"}

def start_api_server():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)


threading.Thread(target=cleanup_thread, daemon=True).start()
threading.Thread(target=start_api_server, daemon=True).start()
asyncio.run(process_queue())