Upload main.py with huggingface_hub

main.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, F
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from huggingface_hub import hf_hub_download
 import os
 import sys
 import uuid
@@ -36,28 +37,58 @@ AUDIO_DIR = os.path.join(os.getcwd(), "audio_files")
 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"
+}
+
 # Download model files if they don't exist
 def ensure_model_files():
-    config_file = os.path.join(MODELS_DIR, "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml")
-    model_file = os.path.join(MODELS_DIR, "wavtokenizer_large_speech_320_24k.ckpt")
-    
-    if not os.path.exists(config_file):
-        print("Downloading config file...")
-        os.system(f"wget -O {config_file} https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml")
+    config_file = os.path.join(MODELS_DIR, MODEL_CONFIG["config_file"])
+    model_file = os.path.join(MODELS_DIR, MODEL_CONFIG["model_file"])
     
-    if not os.path.exists(model_file):
-        print("Downloading model file...")
-        os.system(f"gdown -O {model_file} 1-ASeEkrn4HY49yZWHTASgfGFNXdVnLTt")
-    
-    return os.path.exists(config_file) and os.path.exists(model_file)
+    try:
+        if not os.path.exists(config_file):
+            print("Copying config file...")
+            source_config = MODEL_CONFIG["config_file"]
+            if os.path.exists(source_config):
+                shutil.copy(source_config, config_file)
+            else:
+                print(f"Config file not found: {source_config}")
+                return False
+        
+        if not os.path.exists(model_file):
+            print("Downloading model file from Hugging Face Hub...")
+            hf_token = os.environ.get("HF_TOKEN")
+            if not hf_token:
+                print("HF_TOKEN environment variable not set")
+                return False
+                
+            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"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"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, "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"),
-            wavtokenizer_ckpt_path=os.path.join(MODELS_DIR, "wavtokenizer_large_speech_320_24k.ckpt")
+            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:
@@ -72,7 +103,12 @@ class TextRequest(BaseModel):
 # Health check endpoint
 @app.get("/")
 def read_root():
-    return {"status": "Nigerian Text-to-Speech API is running"}
+    model_status = "available" if os.path.exists(os.path.join(MODELS_DIR, MODEL_CONFIG["model_file"])) else "not available"
+    return {
+        "status": "Nigerian Text-to-Speech API is running",
+        "model_status": model_status,
+        "timestamp": "2025-04-22 04:45:43"
+    }
 
 # Text to speech endpoint
 @app.post("/tts")
@@ -80,12 +116,18 @@ async def text_to_speech(request: TextRequest):
     try:
         # Ensure model files are available
         if not ensure_model_files():
-            raise HTTPException(status_code=500, detail="Failed to download 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")
+            raise HTTPException(
+                status_code=500,
+                detail="Failed to initialize YarnGPT. Please check logs for details."
+            )
         
         # Generate audio
         audio_file_id = str(uuid.uuid4())
@@ -100,22 +142,41 @@ async def text_to_speech(request: TextRequest):
             filename=f"{audio_file_id}.wav"
         )
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error generating audio: {str(e)}"
+        )
 
 # List available files
 @app.get("/list_files")
 def list_files():
-    files = os.listdir(AUDIO_DIR)
-    return {"files": files}
+    try:
+        files = [f for f in os.listdir(AUDIO_DIR) if f.endswith('.wav')]
+        return {"files": files}
+    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, f"{file_id}")
+    file_path = os.path.join(AUDIO_DIR, file_id)
     if not os.path.exists(file_path):
-        raise HTTPException(status_code=404, detail="Audio file not found")
+        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("Starting up Nigerian Text-to-Speech API...")
+    if not ensure_model_files():
+        print("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)