mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
424 lines
12 KiB
Python
424 lines
12 KiB
Python
"""
|
|
FastAPI application for voicebox backend.
|
|
|
|
Handles voice cloning, generation history, and server mode.
|
|
"""
|
|
|
|
from fastapi import FastAPI, Depends, UploadFile, File, Form, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import FileResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from sqlalchemy.orm import Session
|
|
from typing import List, Optional
|
|
import uvicorn
|
|
import argparse
|
|
import torch
|
|
import tempfile
|
|
from pathlib import Path
|
|
import uuid
|
|
|
|
from . import database, models, profiles, history, tts, transcribe
|
|
from .database import get_db, init_db
|
|
|
|
# Initialize database
|
|
init_db()
|
|
|
|
app = FastAPI(
|
|
title="voicebox API",
|
|
description="Production-quality Qwen3-TTS voice cloning API",
|
|
version="0.1.0",
|
|
)
|
|
|
|
# CORS middleware
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"], # Configure appropriately for production
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
# ============================================
|
|
# ROOT & HEALTH ENDPOINTS
|
|
# ============================================
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
"""Root endpoint."""
|
|
return {"message": "voicebox API", "version": "0.1.0"}
|
|
|
|
|
|
@app.get("/health", response_model=models.HealthResponse)
|
|
async def health():
|
|
"""Health check endpoint."""
|
|
tts_model = tts.get_tts_model()
|
|
gpu_available = torch.cuda.is_available()
|
|
|
|
vram_used = None
|
|
if gpu_available:
|
|
vram_used = torch.cuda.memory_allocated() / 1024 / 1024 # MB
|
|
|
|
return models.HealthResponse(
|
|
status="healthy",
|
|
model_loaded=tts_model.is_loaded(),
|
|
gpu_available=gpu_available,
|
|
vram_used_mb=vram_used,
|
|
)
|
|
|
|
|
|
# ============================================
|
|
# VOICE PROFILE ENDPOINTS
|
|
# ============================================
|
|
|
|
@app.post("/profiles", response_model=models.VoiceProfileResponse)
|
|
async def create_profile(
|
|
data: models.VoiceProfileCreate,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Create a new voice profile."""
|
|
try:
|
|
return await profiles.create_profile(data, db)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
|
|
|
|
@app.get("/profiles", response_model=List[models.VoiceProfileResponse])
|
|
async def list_profiles(db: Session = Depends(get_db)):
|
|
"""List all voice profiles."""
|
|
return await profiles.list_profiles(db)
|
|
|
|
|
|
@app.get("/profiles/{profile_id}", response_model=models.VoiceProfileResponse)
|
|
async def get_profile(
|
|
profile_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Get a voice profile by ID."""
|
|
profile = await profiles.get_profile(profile_id, db)
|
|
if not profile:
|
|
raise HTTPException(status_code=404, detail="Profile not found")
|
|
return profile
|
|
|
|
|
|
@app.put("/profiles/{profile_id}", response_model=models.VoiceProfileResponse)
|
|
async def update_profile(
|
|
profile_id: str,
|
|
data: models.VoiceProfileCreate,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Update a voice profile."""
|
|
profile = await profiles.update_profile(profile_id, data, db)
|
|
if not profile:
|
|
raise HTTPException(status_code=404, detail="Profile not found")
|
|
return profile
|
|
|
|
|
|
@app.delete("/profiles/{profile_id}")
|
|
async def delete_profile(
|
|
profile_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Delete a voice profile."""
|
|
success = await profiles.delete_profile(profile_id, db)
|
|
if not success:
|
|
raise HTTPException(status_code=404, detail="Profile not found")
|
|
return {"message": "Profile deleted successfully"}
|
|
|
|
|
|
@app.post("/profiles/{profile_id}/samples", response_model=models.ProfileSampleResponse)
|
|
async def add_profile_sample(
|
|
profile_id: str,
|
|
file: UploadFile = File(...),
|
|
reference_text: str = Form(...),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Add a sample to a voice profile."""
|
|
# Save uploaded file to temporary location
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
|
content = await file.read()
|
|
tmp.write(content)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
sample = await profiles.add_profile_sample(
|
|
profile_id,
|
|
tmp_path,
|
|
reference_text,
|
|
db,
|
|
)
|
|
return sample
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
finally:
|
|
# Clean up temp file
|
|
Path(tmp_path).unlink(missing_ok=True)
|
|
|
|
|
|
@app.get("/profiles/{profile_id}/samples", response_model=List[models.ProfileSampleResponse])
|
|
async def get_profile_samples(
|
|
profile_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Get all samples for a profile."""
|
|
return await profiles.get_profile_samples(profile_id, db)
|
|
|
|
|
|
@app.delete("/profiles/samples/{sample_id}")
|
|
async def delete_profile_sample(
|
|
sample_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Delete a profile sample."""
|
|
success = await profiles.delete_profile_sample(sample_id, db)
|
|
if not success:
|
|
raise HTTPException(status_code=404, detail="Sample not found")
|
|
return {"message": "Sample deleted successfully"}
|
|
|
|
|
|
# ============================================
|
|
# GENERATION ENDPOINTS
|
|
# ============================================
|
|
|
|
@app.post("/generate", response_model=models.GenerationResponse)
|
|
async def generate_speech(
|
|
data: models.GenerationRequest,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Generate speech from text using a voice profile."""
|
|
try:
|
|
# Get profile
|
|
profile = await profiles.get_profile(data.profile_id, db)
|
|
if not profile:
|
|
raise HTTPException(status_code=404, detail="Profile not found")
|
|
|
|
# Create voice prompt from profile
|
|
voice_prompt = await profiles.create_voice_prompt_for_profile(
|
|
data.profile_id,
|
|
db,
|
|
)
|
|
|
|
# Generate audio
|
|
tts_model = tts.get_tts_model()
|
|
audio, sample_rate = await tts_model.generate(
|
|
data.text,
|
|
voice_prompt,
|
|
data.language,
|
|
data.seed,
|
|
)
|
|
|
|
# Calculate duration
|
|
duration = len(audio) / sample_rate
|
|
|
|
# Save audio
|
|
generation_id = str(uuid.uuid4())
|
|
audio_path = history.GENERATIONS_DIR / f"{generation_id}.wav"
|
|
|
|
from .utils.audio import save_audio
|
|
save_audio(audio, str(audio_path), sample_rate)
|
|
|
|
# Create history entry
|
|
generation = await history.create_generation(
|
|
profile_id=data.profile_id,
|
|
text=data.text,
|
|
language=data.language,
|
|
audio_path=str(audio_path),
|
|
duration=duration,
|
|
seed=data.seed,
|
|
db=db,
|
|
)
|
|
|
|
return generation
|
|
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ============================================
|
|
# HISTORY ENDPOINTS
|
|
# ============================================
|
|
|
|
@app.get("/history", response_model=List[models.GenerationResponse])
|
|
async def list_history(
|
|
profile_id: Optional[str] = None,
|
|
search: Optional[str] = None,
|
|
limit: int = 50,
|
|
offset: int = 0,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""List generation history with optional filters."""
|
|
query = models.HistoryQuery(
|
|
profile_id=profile_id,
|
|
search=search,
|
|
limit=limit,
|
|
offset=offset,
|
|
)
|
|
generations, total = await history.list_generations(query, db)
|
|
return generations
|
|
|
|
|
|
@app.get("/history/{generation_id}", response_model=models.GenerationResponse)
|
|
async def get_generation(
|
|
generation_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Get a generation by ID."""
|
|
generation = await history.get_generation(generation_id, db)
|
|
if not generation:
|
|
raise HTTPException(status_code=404, detail="Generation not found")
|
|
return generation
|
|
|
|
|
|
@app.delete("/history/{generation_id}")
|
|
async def delete_generation(
|
|
generation_id: str,
|
|
db: Session = Depends(get_db),
|
|
):
|
|
"""Delete a generation."""
|
|
success = await history.delete_generation(generation_id, db)
|
|
if not success:
|
|
raise HTTPException(status_code=404, detail="Generation not found")
|
|
return {"message": "Generation deleted successfully"}
|
|
|
|
|
|
@app.get("/history/stats")
|
|
async def get_stats(db: Session = Depends(get_db)):
|
|
"""Get generation statistics."""
|
|
return await history.get_generation_stats(db)
|
|
|
|
|
|
# ============================================
|
|
# TRANSCRIPTION ENDPOINTS
|
|
# ============================================
|
|
|
|
@app.post("/transcribe", response_model=models.TranscriptionResponse)
|
|
async def transcribe_audio(
|
|
file: UploadFile = File(...),
|
|
language: Optional[str] = Form(None),
|
|
):
|
|
"""Transcribe audio file to text."""
|
|
# Save uploaded file to temporary location
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
|
content = await file.read()
|
|
tmp.write(content)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
# Get audio duration
|
|
from .utils.audio import load_audio
|
|
audio, sr = load_audio(tmp_path)
|
|
duration = len(audio) / sr
|
|
|
|
# Transcribe
|
|
whisper_model = transcribe.get_whisper_model()
|
|
text = await whisper_model.transcribe(tmp_path, language)
|
|
|
|
return models.TranscriptionResponse(
|
|
text=text,
|
|
duration=duration,
|
|
)
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
finally:
|
|
# Clean up temp file
|
|
Path(tmp_path).unlink(missing_ok=True)
|
|
|
|
|
|
# ============================================
|
|
# FILE SERVING
|
|
# ============================================
|
|
|
|
@app.get("/audio/{generation_id}")
|
|
async def get_audio(generation_id: str, db: Session = Depends(get_db)):
|
|
"""Serve generated audio file."""
|
|
generation = await history.get_generation(generation_id, db)
|
|
if not generation:
|
|
raise HTTPException(status_code=404, detail="Generation not found")
|
|
|
|
audio_path = Path(generation.audio_path)
|
|
if not audio_path.exists():
|
|
raise HTTPException(status_code=404, detail="Audio file not found")
|
|
|
|
return FileResponse(
|
|
audio_path,
|
|
media_type="audio/wav",
|
|
filename=f"generation_{generation_id}.wav",
|
|
)
|
|
|
|
|
|
# ============================================
|
|
# MODEL MANAGEMENT
|
|
# ============================================
|
|
|
|
@app.post("/models/load")
|
|
async def load_model(model_size: str = "1.7B"):
|
|
"""Manually load TTS model."""
|
|
try:
|
|
tts_model = tts.get_tts_model()
|
|
tts_model.load_model(model_size)
|
|
return {"message": f"Model {model_size} loaded successfully"}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/models/unload")
|
|
async def unload_model():
|
|
"""Unload TTS model to free memory."""
|
|
try:
|
|
tts.unload_tts_model()
|
|
return {"message": "Model unloaded successfully"}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ============================================
|
|
# STARTUP & SHUTDOWN
|
|
# ============================================
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""Run on application startup."""
|
|
print("voicebox API starting up...")
|
|
print(f"Database initialized at {database._db_path}")
|
|
print(f"GPU available: {torch.cuda.is_available()}")
|
|
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
"""Run on application shutdown."""
|
|
print("voicebox API shutting down...")
|
|
# Unload models to free memory
|
|
tts.unload_tts_model()
|
|
transcribe.unload_whisper_model()
|
|
|
|
|
|
# ============================================
|
|
# MAIN
|
|
# ============================================
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="voicebox backend server")
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
default="127.0.0.1",
|
|
help="Host to bind to (use 0.0.0.0 for remote access)",
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=8000,
|
|
help="Port to bind to",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
uvicorn.run(
|
|
"main:app",
|
|
host=args.host,
|
|
port=args.port,
|
|
reload=False, # Disable reload in production
|
|
)
|