from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.responses import FileResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import os import uuid import torch import torchaudio import base64 from transformers import AutoModelForCausalLM from yarngpt.audiotokenizer import AudioTokenizerV2 import uvicorn from datetime import datetime, timedelta import asyncio import logging # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI(title="Nigerian TTS API", version="1.0.0") # Add CORS middleware to allow requests from anywhere app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global variables for model components audio_tokenizer = None model = None model_loaded = False loading_error = None # Model configuration - Updated paths for Hugging Face Spaces tokenizer_path = "saheedniyi/YarnGPT2" # These files should be downloaded to /tmp during startup wav_tokenizer_config_path = "/tmp/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml" wav_tokenizer_model_path = "/tmp/wavtokenizer_large_speech_320_24k.ckpt" # Available voices and languages AVAILABLE_VOICES = { "female": ["zainab", "idera", "regina", "chinenye", "joke", "remi"], "male": ["jude", "tayo", "umar", "osagie", "onye", "emma"] } AVAILABLE_LANGUAGES = ["english", "yoruba", "igbo", "hausa"] # Input validation model class TTSRequest(BaseModel): text: str language: str = "english" voice: str = "idera" # Output model with base64-encoded audio class TTSResponse(BaseModel): audio_base64: str audio_url: str text: str voice: str language: str async def download_model_files(): """Download required model files""" global loading_error try: import requests from pathlib import Path logger.info("Starting model file downloads...") # URLs for the model files config_url = "https://huggingface.co/saheedniyi/YarnGPT2/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml" model_url = "https://huggingface.co/saheedniyi/YarnGPT2/resolve/main/wavtokenizer_large_speech_320_24k.ckpt" # Create tmp directory if it doesn't exist Path("/tmp").mkdir(exist_ok=True) # Download config file if not os.path.exists(wav_tokenizer_config_path): logger.info("Downloading tokenizer config...") response = requests.get(config_url, stream=True) response.raise_for_status() with open(wav_tokenizer_config_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) logger.info("Config file downloaded successfully") # Download model file if not os.path.exists(wav_tokenizer_model_path): logger.info("Downloading tokenizer model (this may take a while)...") response = requests.get(model_url, stream=True) response.raise_for_status() total_size = int(response.headers.get('content-length', 0)) downloaded = 0 with open(wav_tokenizer_model_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) downloaded += len(chunk) if total_size > 0: progress = (downloaded / total_size) * 100 if downloaded % (1024 * 1024 * 10) == 0: # Log every 10MB logger.info(f"Download progress: {progress:.1f}%") logger.info("Model file downloaded successfully") logger.info("All model files are ready") except Exception as e: error_msg = f"Error downloading model files: {str(e)}" logger.error(error_msg) loading_error = error_msg raise e async def load_models(): """Load the YarnGPT model and tokenizer""" global audio_tokenizer, model, model_loaded, loading_error try: logger.info("Loading YarnGPT model and tokenizer...") # First download the required files await download_model_files() # Initialize audio tokenizer logger.info("Initializing audio tokenizer...") audio_tokenizer = AudioTokenizerV2( tokenizer_path, wav_tokenizer_model_path, wav_tokenizer_config_path ) # Load the main model logger.info("Loading main model...") model = AutoModelForCausalLM.from_pretrained( tokenizer_path, torch_dtype="auto" ).to(audio_tokenizer.device) model_loaded = True logger.info("Model loaded successfully!") except Exception as e: error_msg = f"Error loading models: {str(e)}" logger.error(error_msg) loading_error = error_msg model_loaded = False @app.on_event("startup") async def startup_event(): """Load models when the API starts""" asyncio.create_task(load_models()) @app.get("/") async def root(): """API health check and info""" return { "status": "ok" if model_loaded else "loading", "message": "Nigerian TTS API is running" if model_loaded else "Models are loading...", "model_loaded": model_loaded, "loading_error": loading_error, "available_languages": AVAILABLE_LANGUAGES, "available_voices": AVAILABLE_VOICES } @app.get("/health") async def health_check(): """Detailed health check""" return { "status": "healthy" if model_loaded else "loading", "model_loaded": model_loaded, "loading_error": loading_error, "timestamp": datetime.now().isoformat() } @app.post("/tts", response_model=TTSResponse) async def text_to_speech(request: TTSRequest, background_tasks: BackgroundTasks): """Convert text to Nigerian-accented speech""" # Check if models are loaded if not model_loaded: if loading_error: raise HTTPException(status_code=503, detail=f"Model loading failed: {loading_error}") else: raise HTTPException(status_code=503, detail="Models are still loading. Please try again in a moment.") # Validate inputs if request.language not in AVAILABLE_LANGUAGES: raise HTTPException(status_code=400, detail=f"Language must be one of {AVAILABLE_LANGUAGES}") all_voices = AVAILABLE_VOICES["female"] + AVAILABLE_VOICES["male"] if request.voice not in all_voices: raise HTTPException(status_code=400, detail=f"Voice must be one of {all_voices}") # Generate unique filename audio_id = str(uuid.uuid4()) output_path = f"audio_files/{audio_id}.wav" os.makedirs("audio_files", exist_ok=True) try: logger.info(f"Generating TTS for text: '{request.text[:50]}...' with voice: {request.voice}") # Create prompt and generate audio prompt = audio_tokenizer.create_prompt( request.text, lang=request.language, speaker_name=request.voice ) input_ids = audio_tokenizer.tokenize_prompt(prompt) output = model.generate( input_ids=input_ids, temperature=0.1, repetition_penalty=1.1, max_length=4000, ) codes = audio_tokenizer.get_codes(output) audio = audio_tokenizer.get_audio(codes) # Save audio file torchaudio.save(output_path, audio, sample_rate=24000) logger.info(f"Audio saved to {output_path}") # Read the file and encode as base64 with open(output_path, "rb") as audio_file: audio_bytes = audio_file.read() audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') # Clean up old files after a while background_tasks.add_task(cleanup_old_files) return TTSResponse( audio_base64=audio_base64, audio_url=f"/audio/{audio_id}.wav", text=request.text, voice=request.voice, language=request.language ) except Exception as e: logger.error(f"Error generating audio: {str(e)}") raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}") @app.get("/audio/{filename}") async def get_audio(filename: str): """Serve audio files""" file_path = f"audio_files/{filename}" if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Audio file not found") return FileResponse(file_path, media_type="audio/wav") def cleanup_old_files(): """Delete audio files older than 6 hours to manage disk space""" try: now = datetime.now() audio_dir = "audio_files" if not os.path.exists(audio_dir): return for filename in os.listdir(audio_dir): if not filename.endswith(".wav"): continue file_path = os.path.join(audio_dir, filename) file_mod_time = datetime.fromtimestamp(os.path.getmtime(file_path)) # Delete files older than 6 hours if now - file_mod_time > timedelta(hours=6): os.remove(file_path) logger.info(f"Deleted old audio file: {filename}") except Exception as e: logger.error(f"Error cleaning up old files: {e}") if __name__ == "__main__": logger.info("Starting Nigerian TTS API server...") uvicorn.run(app, host="0.0.0.0", port=7860)