main.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import uvicorn
import numpy as np
import tempfile
import os
from typing import Dict, Any
import shutil

# Import resemblyzer for audio embeddings
from resemblyzer import VoiceEncoder, preprocess_wav
from pathlib import Path

app = FastAPI(title="Audio Embedding API", 
              description="API for generating vector embeddings from audio files using resemblyzer")

# Initialize the voice encoder
voice_encoder = VoiceEncoder()

@app.post("/generate-embedding/", response_model=Dict[str, Any])
async def generate_embedding(audio_file: UploadFile = File(...)):
    """
    Generate a vector embedding from an uploaded audio file.

    Args:
        audio_file: The audio file to process (WAV format recommended)

    Returns:
        A JSON object containing the embedding vector and metadata
    """
    # Check if the file is an audio file
    if not audio_file.content_type.startswith("audio/"):
        raise HTTPException(status_code=400, detail="File must be an audio file")

    # Create a temporary file to store the uploaded audio
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
        # Write the uploaded file content to the temporary file
        shutil.copyfileobj(audio_file.file, temp_file)
        temp_file_path = temp_file.name

    try:
        # Preprocess the audio file
        wav = preprocess_wav(Path(temp_file_path))

        # Generate the embedding
        embedding = voice_encoder.embed_utterance(wav)

        # Convert the numpy array to a list for JSON serialization
        embedding_list = embedding.tolist()

        return {
            "filename": audio_file.filename,
            "embedding_dimension": len(embedding_list),
            "embedding": embedding_list
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing audio: {str(e)}")
    finally:
        # Clean up the temporary file
        if os.path.exists(temp_file_path):
            os.remove(temp_file_path)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)