Kareem175/DialectLink-API

1import os2os.environ["COQUI_TOS_AGREED"] = "1"3import gc4import tempfile5import uuid6import subprocess7import torch8import torchaudio9import soundfile as sf10import numpy as np11from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Depends12from fastapi.responses import FileResponse13from fastapi.security import APIKeyHeader14from pydantic import BaseModel15from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForSequenceClassification, pipeline16from huggingface_hub import snapshot_download17from fastapi.staticfiles import StaticFiles18from fastapi.middleware.cors import CORSMiddleware19 20# Monkey-patch torchaudio.load because PyTorch 2.6 forces broken torchcodec21def custom_torchaudio_load(filepath, *args, **kwargs):22    data, samplerate = sf.read(filepath, dtype='float32')23    if data.ndim == 1:24        data = data[None, :]25    else:26        data = data.T27    return torch.from_numpy(data), samplerate28torchaudio.load = custom_torchaudio_load29 30# Ensure static directories31os.makedirs(os.path.join("static", "audio"), exist_ok=True)32 33app = FastAPI(title="DialectLink AI Triple-Engine Server")34 35app.mount("/static", StaticFiles(directory="static"), name="static")36 37app.add_middleware(38    CORSMiddleware,39    allow_origins=["*"],40    allow_credentials=True,41    allow_methods=["*"],42    allow_headers=["*"],43)44 45api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True)46 47def get_api_key(api_key: str = Depends(api_key_header)):48    return api_key49 50# --- Configuration: The Models Paths ---51CLASS_MODEL_PATH = "Kareem175/classification-model"52DIALECT_TO_MSA_PATH = "Kareem175/Dialect-To-MSA"53MSA_TO_DIALECT_PATH = "Kareem175/msa-to-dialect"54 55STT_EGY_MAR_PATH = "Kareem175/stt-egy-mar"56STT_ALG_PAL_PATH = "Kareem175/stt-algerian-palestinian"57STT_PAL_PATH = "Kareem175/stt-palestinian"58STT_SAU_TUN_PATH = "Kareem175/stt-saudi-tunisian"59 60XTTS_REPO = "Kareem175/xtts-model"61XTTS_LOCAL_DIR = "./xtts_model_local"62 63# Global pipeline variables64classifier_pipe = None65d2msa_tokenizer = None66d2msa_model = None67msa2d_tokenizer = None68msa2d_model = None69tts_model = None70 71# STT Lazy Loading Globals72current_stt_model_name = None73stt_pipe = None74 75print("Starting DialectLink 3-Stage AI Pipeline...")76 77# --- Loading the NLP and TTS Models ---78try:79    print("1. Loading Classification Model...")80    try:81        class_tokenizer = AutoTokenizer.from_pretrained(CLASS_MODEL_PATH)82    except Exception:83        print(" -> [WARNING] Classification Tokenizer not found in repo. Falling back to UBC-NLP/MARBERT...")84        class_tokenizer = AutoTokenizer.from_pretrained("UBC-NLP/MARBERT")85    class_model = AutoModelForSequenceClassification.from_pretrained(CLASS_MODEL_PATH)86    classifier_pipe = pipeline("text-classification", model=class_model, tokenizer=class_tokenizer)87    print(" -> Classification Model Loaded Successfully!")88 89    print("\n2. Loading Dialect -> MSA Translation Model...")90    d2msa_tokenizer = AutoTokenizer.from_pretrained(DIALECT_TO_MSA_PATH, use_fast=False)91    d2msa_model = AutoModelForSeq2SeqLM.from_pretrained(DIALECT_TO_MSA_PATH)92    print(" -> Dialect->MSA Model Loaded Successfully!")93 94    print("\n3. Loading MSA -> Target Dialect Translation Model...")95    msa2d_tokenizer = AutoTokenizer.from_pretrained(MSA_TO_DIALECT_PATH, use_fast=False)96    msa2d_model = AutoModelForSeq2SeqLM.from_pretrained(MSA_TO_DIALECT_PATH)97    print(" -> MSA->Dialect Model Loaded Successfully!")98 99    print("\n4. STT Models will be lazy-loaded on demand.")100 101    print("\n5. Loading XTTS Model...")102    if not os.path.exists(XTTS_LOCAL_DIR):103        print(" -> Downloading XTTS model from Hugging Face Hub...")104        snapshot_download(repo_id=XTTS_REPO, local_dir=XTTS_LOCAL_DIR)105        106    try:107        from TTS.api import TTS108        original_load = torch.load109        def patched_load(*args, **kwargs):110            kwargs['weights_only'] = False111            return original_load(*args, **kwargs)112        torch.load = patched_load113 114        try:115            device = "cuda" if torch.cuda.is_available() else "cpu"116            tts_model = TTS(model_path=XTTS_LOCAL_DIR, config_path=f"{XTTS_LOCAL_DIR}/config.json").to(device)117            print(" -> XTTS Model Loaded Successfully!")118        finally:119            torch.load = original_load120    except Exception as e:121        print(f" -> [WARNING] Failed to load XTTS: {e}")122 123except Exception as e:124    print(f"\n[ERROR] An error occurred while loading the models: {str(e)}")125 126# --- Lazy Load STT Helper ---127def get_stt_pipeline(source_dialect_code: str):128    global current_stt_model_name, stt_pipe129    130    if source_dialect_code in ["EGY", "MAR"]:131        target_model = STT_EGY_MAR_PATH132    elif source_dialect_code == "PAL":133        target_model = STT_PAL_PATH134    elif source_dialect_code == "ALG":135        target_model = STT_ALG_PAL_PATH136    elif source_dialect_code in ["SAU", "TUN", "KHA"]:137        target_model = STT_SAU_TUN_PATH138    else:139        target_model = STT_EGY_MAR_PATH # Fallback140 141    if current_stt_model_name == target_model and stt_pipe is not None:142        return stt_pipe143 144    print(f"\n[Dynamic Router] Switching STT model to: {target_model}")145    if stt_pipe is not None:146        del stt_pipe147        gc.collect()148        if torch.cuda.is_available():149            torch.cuda.empty_cache()150 151    try:152        stt_pipe = pipeline("automatic-speech-recognition", model=target_model)153    except Exception:154        print(" -> [WARNING] STT Tokenizer mismatch. Using base whisper-small tokenizer...")155        from transformers import WhisperProcessor156        stt_processor = WhisperProcessor.from_pretrained("openai/whisper-small")157        stt_pipe = pipeline("automatic-speech-recognition", model=target_model, tokenizer=stt_processor.tokenizer, feature_extractor=stt_processor.feature_extractor)158        159    current_stt_model_name = target_model160    return stt_pipe161 162 163# --- API Endpoints ---164@app.get("/")165def read_root():166    return {"message": "Welcome to DialectLink AI API. Server is running successfully!"}167 168class TranslationRequest(BaseModel):169    text: str170    target_dialect: str171 172@app.post("/translate")173def translate_text(req: TranslationRequest, api_key: str = Depends(get_api_key)):174    if not all([classifier_pipe, d2msa_model, msa2d_model]):175        return {176            "source_dialect_detected": "mockup",177            "msa_intermediate": f"[Mock MSA] {req.text}",178            "translated_text": f"[Mock Translation to {req.target_dialect}] {req.text}"179        }180    try:181        class_result = classifier_pipe(req.text)182        detected_dialect = class_result[0]['label'] 183        184        inputs1 = d2msa_tokenizer(req.text, return_tensors="pt", max_length=128, truncation=True)185        outputs1 = d2msa_model.generate(**inputs1, max_new_tokens=128)186        intermediate_msa = d2msa_tokenizer.decode(outputs1[0], skip_special_tokens=True)187 188        m2d_prompt = f"translate to {req.target_dialect}: {intermediate_msa}" 189        inputs2 = msa2d_tokenizer(m2d_prompt, return_tensors="pt", max_length=128, truncation=True)190        outputs2 = msa2d_model.generate(**inputs2, max_new_tokens=128)191        final_translation = msa2d_tokenizer.decode(outputs2[0], skip_special_tokens=True)192 193        return {194            "source_dialect_detected": detected_dialect,195            "msa_intermediate": intermediate_msa,196            "translated_text": final_translation197        }198    except Exception as e:199        raise HTTPException(status_code=500, detail=str(e))200 201@app.post("/translate_audio")202async def translate_audio(203    file: UploadFile = File(...), 204    target_dialect: str = Form(...),205    source_dialect: str = Form("EGY"),206    api_key: str = Depends(get_api_key)207):208    try:209        temp_in = tempfile.NamedTemporaryFile(delete=False, suffix=".m4a")210        content = await file.read()211        temp_in.write(content)212        temp_in.close()213 214        temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")215        temp_wav.close()216        217        try:218            subprocess.run(["ffmpeg", "-y", "-i", temp_in.name, "-ar", "16000", "-ac", "1", temp_wav.name], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)219            audio_path_to_use = temp_wav.name220        except Exception:221            import shutil222            shutil.copy(temp_in.name, temp_wav.name)223            audio_path_to_use = temp_wav.name224 225        data, samplerate = sf.read(audio_path_to_use)226        rms = np.sqrt(np.mean(data**2)) if len(data) > 0 else 0227        228        if rms < 0.005:229            transcribed_text = "[صمت/غير مفهوم]"230        else:231            current_stt = get_stt_pipeline(source_dialect)232            stt_result = current_stt(audio_path_to_use, generate_kwargs={"language": "arabic", "task": "transcribe"})233            transcribed_text = stt_result["text"]234 235        inputs1 = d2msa_tokenizer(transcribed_text, return_tensors="pt", max_length=128, truncation=True)236        outputs1 = d2msa_model.generate(**inputs1, max_new_tokens=128)237        intermediate_msa = d2msa_tokenizer.decode(outputs1[0], skip_special_tokens=True)238 239        m2d_prompt = f"translate to {target_dialect}: {intermediate_msa}" 240        inputs2 = msa2d_tokenizer(m2d_prompt, return_tensors="pt", max_length=128, truncation=True)241        outputs2 = msa2d_model.generate(**inputs2, max_new_tokens=128)242        final_translation = msa2d_tokenizer.decode(outputs2[0], skip_special_tokens=True)243 244        filename = f"translated_{uuid.uuid4().hex[:8]}.wav"245        file_path = os.path.join("static", "audio", filename)246 247        if tts_model:248            tts_model.tts_to_file(text=final_translation, file_path=file_path, speaker_wav=audio_path_to_use, language="ar")249        else:250            from gtts import gTTS251            tts = gTTS(text=final_translation, lang='ar')252            tts.save(file_path)253 254        return {255            "audio_url": f"/static/audio/{filename}",256            "translated_text": final_translation,257            "transcribed_text": transcribed_text258        }259 260    except Exception as e:261        raise HTTPException(status_code=500, detail=str(e))262