feat(example): add new API (#913)

Add FastAPI-based OpenAI-compatible Text-to-Speech API and Audio Format Conversion Tools

Author: yueguobin

.gitattributes

@@ -1,2 +1,3 @@
 # ignore jupyter notebooks in the language bar on github
 **/*.ipynb linguist-vendored
+*.ipynb

examples/api/README.md

@@ -14,6 +14,11 @@ pip install -r examples/api/requirements.txt
 fastapi dev examples/api/main.py --host 0.0.0.0 --port 8000
 ```
 
+## Run openAI_API server
+
+```
+fastapi dev examples/api/openai_api.py --host 0.0.0.0 --port 8000
+```
 ## Generate audio using requests
 
 ```

examples/api/openai_api.py

@@ -0,0 +1,244 @@
+"""
+openai_api.py
+This module implements a FastAPI-based text-to-speech API compatible with OpenAI's interface specification.
+
+Main features and improvements:
+- Use app.state to manage global state, ensuring thread safety
+- Add exception handling and unified error responses to improve stability
+- Support multiple voice options and audio formats for greater flexibility
+- Add input validation to ensure the validity of request parameters
+- Support additional OpenAI TTS parameters (e.g., speed) for richer functionality
+- Implement health check endpoint for easy service status monitoring
+- Use asyncio.Lock to manage model access, improving concurrency performance
+- Load and manage speaker embedding files to support personalized speech synthesis
+"""
+import io
+import os
+import sys
+import asyncio
+import time 
+from typing import Optional, Dict
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel, Field
+import torch
+
+# Cross-platform compatibility settings
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+# Set working directory and add to system path
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+# Import necessary modules
+import ChatTTS
+from tools.audio import pcm_arr_to_mp3_view, pcm_arr_to_ogg_view, pcm_arr_to_wav_view
+from tools.logger import get_logger
+from tools.normalizer.en import normalizer_en_nemo_text
+from tools.normalizer.zh import normalizer_zh_tn
+
+# Initialize logger
+logger = get_logger("Command")
+
+# Initialize FastAPI application
+app = FastAPI()
+
+# Voice mapping table
+# Download stable voices:
+# ModelScope Community: https://modelscope.cn/studios/ttwwwaa/ChatTTS_Speaker
+# HuggingFace: https://huggingface.co/spaces/taa/ChatTTS_Speaker
+VOICE_MAP = {
+    "default": "1528.pt",
+    "alloy": "1384.pt",
+    "echo": "2443.pt",
+}
+
+# Allowed audio formats
+ALLOWED_FORMATS = {"mp3", "wav", "ogg"}
+
+@app.on_event("startup")
+async def startup_event():
+    """Load ChatTTS model and default speaker embedding when the application starts"""
+    # Initialize ChatTTS and async lock
+    app.state.chat = ChatTTS.Chat(get_logger("ChatTTS"))
+    app.state.model_lock = asyncio.Lock()  # Use async lock instead of thread lock
+    
+    # Register text normalizers
+    app.state.chat.normalizer.register("en", normalizer_en_nemo_text())
+    app.state.chat.normalizer.register("zh", normalizer_zh_tn())
+    
+    logger.info("Initializing ChatTTS...")
+    if app.state.chat.load(source="huggingface"):
+        logger.info("Model loaded successfully.")
+    else:
+        logger.error("Model loading failed, exiting application.")
+        raise RuntimeError("Failed to load ChatTTS model")
+    
+    # Load default speaker embedding
+    # Preload all supported speaker embeddings into memory at startup to avoid repeated loading during runtime
+    app.state.spk_emb_map = {}
+    for voice, spk_path in VOICE_MAP.items():
+        if os.path.exists(spk_path):
+            app.state.spk_emb_map[voice] = torch.load(spk_path, map_location=torch.device("cpu"))
+            logger.info(f"Preloading speaker embedding: {voice} -> {spk_path}")
+        else:
+            logger.warning(f"Speaker embedding not found: {spk_path}, skipping preload")
+    app.state.spk_emb = app.state.spk_emb_map.get("default")  # Default embedding
+
+# Request parameter whitelist
+ALLOWED_PARAMS = {"model", "input", "voice", "response_format", "speed", "stream", "output_format"}
+
+class OpenAITTSRequest(BaseModel):
+    """OpenAI TTS request data model"""
+    model: str = Field(..., description="Speech synthesis model, fixed as 'tts-1'")
+    input: str = Field(..., description="Text content to synthesize", max_length=2048)  # Length limit
+    voice: Optional[str] = Field("default", description="Voice selection, supports: default, alloy, echo")
+    response_format: Optional[str] = Field("mp3", description="Audio format: mp3, wav, ogg")
+    speed: Optional[float] = Field(1.0, ge=0.5, le=2.0, description="Speed, range 0.5-2.0")
+    stream: Optional[bool] = Field(False, description="Whether to stream")
+    output_format: Optional[str] = "mp3"  # Optional formats: mp3, wav, ogg
+    extra_params: Dict[str, Optional[str]] = Field(default_factory=dict, description="Unsupported extra parameters")
+
+    @classmethod
+    def validate_request(cls, request_data: Dict):
+        """Filter unsupported request parameters and unify model value to 'tts-1'"""
+        request_data["model"] = "tts-1"  # Unify model value
+        unsupported_params = set(request_data.keys()) - ALLOWED_PARAMS
+        if unsupported_params:
+            logger.warning(f"Ignoring unsupported parameters: {unsupported_params}")
+        return {key: request_data[key] for key in ALLOWED_PARAMS if key in request_data}
+
+# Unified error response
+@app.exception_handler(Exception)
+async def custom_exception_handler(request, exc):
+    """Custom exception handler"""
+    logger.error(f"Error: {str(exc)}")
+    return JSONResponse(
+        status_code=getattr(exc, "status_code", 500),
+        content={"error": {"message": str(exc), "type": exc.__class__.__name__}}
+    )
+
+@app.post("/v1/audio/speech")
+async def generate_voice(request_data: Dict):
+    """Handle speech synthesis request"""
+    request_data = OpenAITTSRequest.validate_request(request_data)
+    request = OpenAITTSRequest(**request_data)
+    
+    logger.info(f"Received request: text={request.input}..., voice={request.voice}, stream={request.stream}")
+    
+    # Validate audio format
+    if request.response_format not in ALLOWED_FORMATS:
+        raise HTTPException(400, detail=f"Unsupported audio format: {request.response_format}, supported formats: {', '.join(ALLOWED_FORMATS)}")
+
+    # Load speaker embedding for the specified voice
+    spk_emb = app.state.spk_emb_map.get(request.voice, app.state.spk_emb)
+    
+    # Inference parameters
+    params_infer_main = {
+        "text": [request.input],
+        "stream": request.stream,
+        "lang": None,
+        "skip_refine_text": True,  # Do not use text refinement
+        "refine_text_only": False,
+        "use_decoder": True,
+        "audio_seed": 12345678,
+        # "text_seed": 87654321,  # Random seed for text processing, used to control text refinement
+        "do_text_normalization": True,  # Perform text normalization
+        "do_homophone_replacement": True, # Perform homophone replacement
+    }
+     
+    # Inference code parameters
+    params_infer_code = app.state.chat.InferCodeParams(
+        #prompt=f"[speed_{int(request.speed * 10)}]",  # Convert to format supported by ChatTTS
+        prompt="[speed_5]", 
+        top_P=0.5,
+        top_K=10,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        max_new_token=2048,
+        min_new_token=0,
+        show_tqdm=True,
+        ensure_non_empty=True,
+        manual_seed=42,
+        spk_emb=spk_emb,
+        spk_smp=None,
+        txt_smp=None,
+        stream_batch=24,
+        stream_speed=12000,
+        pass_first_n_batches=2
+    )
+
+    try:
+        async with app.state.model_lock:
+            wavs = app.state.chat.infer(
+                text = params_infer_main["text"],
+                stream = params_infer_main["stream"],
+                lang = params_infer_main["lang"],
+                skip_refine_text = params_infer_main["skip_refine_text"],
+                use_decoder = params_infer_main["use_decoder"],
+                do_text_normalization = params_infer_main["do_text_normalization"], 
+                do_homophone_replacement = params_infer_main['do_homophone_replacement'],
+                # params_refine_text = params_refine_text,  
+                params_infer_code=params_infer_code,   
+            )
+    except Exception as e:
+        raise HTTPException(500, detail=f"Speech synthesis failed: {str(e)}")
+
+    def generate_wav_header(sample_rate=24000, bits_per_sample=16, channels=1):
+        """Generate WAV file header (without data length)"""
+        header = bytearray()
+        header.extend(b"RIFF")
+        header.extend(b"\xFF\xFF\xFF\xFF")  # File size unknown
+        header.extend(b"WAVEfmt ")
+        header.extend((16).to_bytes(4, "little"))  # fmt chunk size
+        header.extend((1).to_bytes(2, "little"))  # PCM format
+        header.extend((channels).to_bytes(2, "little"))  # Channels
+        header.extend((sample_rate).to_bytes(4, "little"))  # Sample rate
+        byte_rate = sample_rate * channels * bits_per_sample // 8
+        header.extend((byte_rate).to_bytes(4, "little"))  # Byte rate
+        block_align = channels * bits_per_sample // 8
+        header.extend((block_align).to_bytes(2, "little"))  # Block align
+        header.extend((bits_per_sample).to_bytes(2, "little"))  # Bits per sample
+        header.extend(b"data")
+        header.extend(b"\xFF\xFF\xFF\xFF")  # Data size unknown
+        return bytes(header)
+
+    # Handle audio output format
+    def convert_audio(wav, format):
+        """Convert audio format"""
+        if format == "mp3":
+            return pcm_arr_to_mp3_view(wav)
+        elif format == "wav":
+            return pcm_arr_to_wav_view(wav, include_header=False)  # No header in streaming
+        elif format == "ogg":
+            return pcm_arr_to_ogg_view(wav)
+        return pcm_arr_to_mp3_view(wav) 
+    
+    # Return streaming audio data
+    if request.stream:
+        first_chunk = True
+        async def audio_stream():
+            nonlocal first_chunk
+            for wav in wavs:
+                if request.response_format == "wav" and first_chunk:
+                    yield generate_wav_header()  # Send WAV header
+                    first_chunk = False
+                yield convert_audio(wav, request.response_format)
+        media_type = "audio/wav" if request.response_format == "wav" else "audio/mpeg"
+        return StreamingResponse(audio_stream(), media_type=media_type)
+        
+    # Return audio file directly
+    if request.response_format == 'wav':
+        music_data = pcm_arr_to_wav_view(wavs[0])
+    else:
+        music_data = convert_audio(wavs[0], request.response_format)
+        
+    return StreamingResponse(io.BytesIO(music_data), media_type="audio/mpeg", headers={
+        "Content-Disposition": f"attachment; filename=output.{request.response_format}"
+    })
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "healthy", "model_loaded": bool(app.state.chat)}

openai_api.ipynb

@@ -1,4 +1,4 @@
 from .av import load_audio
-from .pcm import pcm_arr_to_mp3_view
+from .pcm import pcm_arr_to_mp3_view, pcm_arr_to_ogg_view, pcm_arr_to_wav_view
 from .ffmpeg import has_ffmpeg_installed
 from .np import float_to_int16

tools/audio/__init__.py

@@ -1,21 +1,85 @@
 import wave
 from io import BytesIO
-
 import numpy as np
-
 from .np import float_to_int16
 from .av import wav2
 
+def _pcm_to_wav_buffer(wav: np.ndarray, sample_rate: int = 24000) -> BytesIO:
+    """
+    Convert PCM audio data to a WAV format byte stream (internal utility function).
 
-def pcm_arr_to_mp3_view(wav: np.ndarray):
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: WAV format byte stream, stored in a BytesIO object.
+    """
+    # Create an in-memory byte stream buffer
     buf = BytesIO()
+    
+    # Open a WAV file stream in write mode
     with wave.open(buf, "wb") as wf:
-        wf.setnchannels(1)  # Mono channel
-        wf.setsampwidth(2)  # Sample width in bytes
-        wf.setframerate(24000)  # Sample rate in Hz
+        # Set number of channels to 1 (mono)
+        wf.setnchannels(1)
+        # Set sample width to 2 bytes (16-bit)
+        wf.setsampwidth(2)
+        # Set sample rate
+        wf.setframerate(sample_rate)
+        # Convert PCM to 16-bit integer and write
         wf.writeframes(float_to_int16(wav))
+    
+    # Reset buffer pointer to the beginning
     buf.seek(0, 0)
+    return buf
+
+def pcm_arr_to_mp3_view(wav: np.ndarray, sample_rate: int = 24000) -> memoryview:
+    """
+    Convert PCM audio data to MP3 format.
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: MP3 format byte data, returned as a memoryview.
+    """
+    # Get WAV format byte stream
+    buf = _pcm_to_wav_buffer(wav, sample_rate)
+    
+    # Create output buffer
     buf2 = BytesIO()
+    # Convert WAV data to MP3
     wav2(buf, buf2, "mp3")
-    buf.seek(0, 0)
+    # Return MP3 data
     return buf2.getbuffer()
+
+def pcm_arr_to_ogg_view(wav: np.ndarray, sample_rate: int = 24000) -> memoryview:
+    """
+    Convert PCM audio data to OGG format (using Vorbis encoding).
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: OGG format byte data, returned as a memoryview.
+    """
+    # Get WAV format byte stream
+    buf = _pcm_to_wav_buffer(wav, sample_rate)
+    
+    # Create output buffer
+    buf2 = BytesIO()
+    # Convert WAV data to OGG
+    wav2(buf, buf2, "ogg")
+    # Return OGG data
+    return buf2.getbuffer()
+
+def pcm_arr_to_wav_view(wav: np.ndarray, sample_rate: int = 24000, include_header: bool = True) -> memoryview:
+    """
+    Convert PCM audio data to WAV format, with an option to include header.
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :param include_header: Whether to include WAV header, defaults to True.
+    :return: WAV format or raw PCM byte data, returned as a memoryview.
+    """
+    if include_header:
+        # Get complete WAV byte stream
+        buf = _pcm_to_wav_buffer(wav, sample_rate)
+        return buf.getbuffer()
+    else:
+        # Return only converted 16-bit PCM data
+        pcm_data = float_to_int16(wav)
+        return memoryview(pcm_data.tobytes())