Upload main.py with huggingface_hub

main.py

@@ -1 +1,224 @@
-[Previous main.py content...]
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from huggingface_hub import hf_hub_download
+import os
+import sys
+import uuid
+import torch
+import torchaudio
+import base64
+from io import BytesIO
+import shutil
+import importlib.util
+import subprocess
+from datetime import datetime
+
+# Add YarnGPT to path
+sys.path.append(os.path.join(os.getcwd(), "yarngpt"))
+
+# Initialize FastAPI
+app = FastAPI(
+    title="Nigerian Text-to-Speech API",
+    version="1.0.0",
+    description="A FastAPI service for Nigerian Text-to-Speech generation"
+)
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Models directory
+MODELS_DIR = os.path.join(os.getcwd(), "models")
+AUDIO_DIR = os.path.join(os.getcwd(), "audio_files")
+
+# Ensure directories exist
+os.makedirs(MODELS_DIR, exist_ok=True)
+os.makedirs(AUDIO_DIR, exist_ok=True)
+
+# Model configuration
+MODEL_CONFIG = {
+    "config_file": "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
+    "model_file": "wavtokenizer_large_speech_320_24k.ckpt",
+    "repo_id": "Hameed13/nigerian-tts-model"
+}
+
+def get_current_timestamp():
+    return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
+
+# Download model files if they don't exist
+def ensure_model_files():
+    config_file = os.path.join(MODELS_DIR, MODEL_CONFIG["config_file"])
+    model_file = os.path.join(MODELS_DIR, MODEL_CONFIG["model_file"])
+    
+    try:
+        # First check for HF_TOKEN
+        hf_token = os.environ.get("HF_TOKEN")
+        if not hf_token:
+            print(f"[{get_current_timestamp()}] HF_TOKEN environment variable not set")
+            return False
+        
+        # Check and download config file
+        if not os.path.exists(config_file):
+            print(f"[{get_current_timestamp()}] Downloading config file from Hugging Face Hub...")
+            try:
+                hf_hub_download(
+                    repo_id=MODEL_CONFIG["repo_id"],
+                    filename=MODEL_CONFIG["config_file"],
+                    local_dir=MODELS_DIR,
+                    token=hf_token
+                )
+            except Exception as e:
+                print(f"[{get_current_timestamp()}] Error downloading config file: {e}")
+                return False
+        
+        # Check and download model file
+        if not os.path.exists(model_file):
+            print(f"[{get_current_timestamp()}] Downloading model file from Hugging Face Hub...")
+            try:
+                hf_hub_download(
+                    repo_id=MODEL_CONFIG["repo_id"],
+                    filename=MODEL_CONFIG["model_file"],
+                    local_dir=MODELS_DIR,
+                    token=hf_token
+                )
+            except Exception as e:
+                print(f"[{get_current_timestamp()}] Error downloading model file: {e}")
+                return False
+        
+        return os.path.exists(config_file) and os.path.exists(model_file)
+    except Exception as e:
+        print(f"[{get_current_timestamp()}] Error in ensure_model_files: {e}")
+        return False
+
+# Initialize YarnGPT
+def initialize_yarngpt():
+    try:
+        from yarngpt.generate import TextToSpeech
+        tts = TextToSpeech(
+            wavtokenizer_config_path=os.path.join(MODELS_DIR, MODEL_CONFIG["config_file"]),
+            wavtokenizer_ckpt_path=os.path.join(MODELS_DIR, MODEL_CONFIG["model_file"])
+        )
+        return tts
+    except Exception as e:
+        print(f"[{get_current_timestamp()}] Error initializing YarnGPT: {e}")
+        return None
+
+# Request models
+class TextRequest(BaseModel):
+    text: str
+    accent: str = "nigerian"
+
+# Health check endpoint
+@app.get("/")
+def read_root():
+    model_file = os.path.join(MODELS_DIR, MODEL_CONFIG["model_file"])
+    config_file = os.path.join(MODELS_DIR, MODEL_CONFIG["config_file"])
+    hf_token = os.environ.get("HF_TOKEN")
+    
+    status = {
+        "status": "Nigerian Text-to-Speech API is running",
+        "model_status": {
+            "model_file_exists": os.path.exists(model_file),
+            "config_file_exists": os.path.exists(config_file),
+            "models_dir": MODELS_DIR,
+            "hf_token_set": bool(hf_token),
+            "hf_token_valid": bool(hf_token and len(hf_token) > 0)
+        },
+        "timestamp": "2025-04-22 13:39:07",
+        "version": "1.0.0",
+        "author": "Abdulhameed556"
+    }
+    return JSONResponse(content=status)
+
+# Text to speech endpoint
+@app.post("/tts")
+async def text_to_speech(request: TextRequest):
+    try:
+        # Check HF_TOKEN first
+        if not os.environ.get("HF_TOKEN"):
+            raise HTTPException(
+                status_code=500,
+                detail="HF_TOKEN environment variable not set. Please configure your Hugging Face token."
+            )
+        
+        # Ensure model files are available
+        if not ensure_model_files():
+            raise HTTPException(
+                status_code=500,
+                detail="Failed to download or locate model files. Please check logs for details."
+            )
+            
+        # Initialize YarnGPT
+        tts = initialize_yarngpt()
+        if not tts:
+            raise HTTPException(
+                status_code=500,
+                detail="Failed to initialize YarnGPT. Please check logs for details."
+            )
+        
+        # Generate audio
+        audio_file_id = str(uuid.uuid4())
+        output_path = os.path.join(AUDIO_DIR, f"{audio_file_id}.wav")
+        
+        print(f"[{get_current_timestamp()}] Generating audio for text: {request.text[:50]}...")
+        tts.read_text(request.text, output_path)
+        
+        # Return the audio file
+        return FileResponse(
+            output_path,
+            media_type="audio/wav",
+            filename=f"{audio_file_id}.wav"
+        )
+    except Exception as e:
+        print(f"[{get_current_timestamp()}] Error in text_to_speech: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error generating audio: {str(e)}"
+        )
+
+# List available files
+@app.get("/list_files")
+def list_files():
+    try:
+        files = [f for f in os.listdir(AUDIO_DIR) if f.endswith('.wav')]
+        return {
+            "files": files,
+            "count": len(files),
+            "timestamp": get_current_timestamp()
+        }
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error listing files: {str(e)}"
+        )
+
+# Get audio file by id
+@app.get("/audio/{file_id}")
+def get_audio(file_id: str):
+    file_path = os.path.join(AUDIO_DIR, file_id)
+    if not os.path.exists(file_path):
+        raise HTTPException(
+            status_code=404,
+            detail=f"Audio file not found: {file_id}"
+        )
+    return FileResponse(file_path, media_type="audio/wav")
+
+# Add startup event to ensure model is downloaded when the container starts
+@app.on_event("startup")
+async def startup_event():
+    print(f"[{get_current_timestamp()}] Starting up Nigerian Text-to-Speech API...")
+    if not os.environ.get("HF_TOKEN"):
+        print(f"[{get_current_timestamp()}] Warning: HF_TOKEN environment variable not set")
+    if not ensure_model_files():
+        print(f"[{get_current_timestamp()}] Warning: Failed to initialize model files during startup")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)