app.py

"""
ClearWave AI Studio

STARTUP FIX:
  Root cause of "stuck at startup":
    - Transcriber() was initializing Groq client at import time
    - Translator() was firing a background thread that downloaded
      Helsinki-NLP models (~600MB) from HuggingFace immediately at startup
    - On HuggingFace Spaces, this causes the container health check to time out

  Fix: All three services are now lazy-initialized — nothing downloads or
  connects until the first actual audio is processed.
"""

import gradio as gr
import os
import tempfile
import subprocess
import threading
import time
import requests as req
import logging
import uuid
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ══════════════════════════════════════════════════════════════════════
# LAZY SERVICE INIT — nothing loads at startup
# ══════════════════════════════════════════════════════════════════════
_denoiser    = None
_transcriber = None
_translator  = None
_init_lock   = threading.Lock()

def get_denoiser():
    global _denoiser
    if _denoiser is None:
        with _init_lock:
            if _denoiser is None:
                from denoiser import Denoiser
                _denoiser = Denoiser()
    return _denoiser

def get_transcriber():
    global _transcriber
    if _transcriber is None:
        with _init_lock:
            if _transcriber is None:
                from transcriber import Transcriber
                _transcriber = Transcriber()
    return _transcriber

def get_translator():
    global _translator
    if _translator is None:
        with _init_lock:
            if _translator is None:
                from translator import Translator
                _translator = Translator()
    return _translator

LANGUAGES = {
    "Auto Detect": "auto",
    "English":  "en", "Telugu":   "te",
    "Hindi":    "hi", "Tamil":    "ta",
    "Kannada":  "kn", "Spanish":  "es",
    "French":   "fr", "German":   "de",
    "Japanese": "ja", "Chinese":  "zh",
    "Arabic":   "ar",
}
OUT_LANGS = {k: v for k, v in LANGUAGES.items() if k != "Auto Detect"}

# ══════════════════════════════════════════════════════════════════════
# JOB STORE
# ══════════════════════════════════════════════════════════════════════
_jobs: dict = {}
_jobs_lock  = threading.Lock()
JOB_TTL     = 3600

def _new_job():
    jid = str(uuid.uuid4())
    with _jobs_lock:
        _jobs[jid] = {"status": "queued", "step": 0,
                      "message": "Queued...", "result": None,
                      "created_at": time.time()}
    return jid

def _update_job(jid, **kw):
    with _jobs_lock:
        if jid in _jobs:
            _jobs[jid].update(kw)

def _get_job(jid):
    with _jobs_lock:
        return dict(_jobs.get(jid, {}))

def _cleanup():
    while True:
        time.sleep(300)
        now = time.time()
        with _jobs_lock:
            exp = [k for k, v in _jobs.items()
                   if now - v.get("created_at", 0) > JOB_TTL]
            for k in exp:
                del _jobs[k]

threading.Thread(target=_cleanup, daemon=True).start()

# ══════════════════════════════════════════════════════════════════════
# FORMAT CONVERTER
# NOTE: This is only used to handle exotic formats (mp4, webm, etc.)
#       that Cleanvoice might not accept. For common formats (mp3, wav,
#       aac, flac) we pass the file directly to Cleanvoice untouched.
# ══════════════════════════════════════════════════════════════════════
# Formats Cleanvoice reliably accepts at its signed upload endpoint.
# .aac, .m4a, .ogg, .opus, .flac all cause HTTP 400 — convert them to WAV.
# denoiser.py handles per-file conversion internally, but this pre-pass
# catches exotic container formats (mp4, webm, mpeg) before the pipeline.
CLEANVOICE_NATIVE = {".mp3", ".wav"}

def maybe_convert_for_cleanvoice(path):
    """Convert any format not natively accepted by Cleanvoice to WAV."""
    if not path:
        return path
    ext = os.path.splitext(path)[1].lower()
    if ext in CLEANVOICE_NATIVE:
        return path  # already safe
    # Convert to PCM WAV — universally accepted by Cleanvoice upload API
    try:
        out = path + "_converted.wav"
        r   = subprocess.run([
            "ffmpeg", "-y", "-i", path,
            "-vn",                   # drop video if present
            "-acodec", "pcm_s16le",  # uncompressed PCM
            "-ar", "16000",          # 16kHz
            "-ac", "1",              # mono
            out
        ], capture_output=True)
        if r.returncode == 0 and os.path.exists(out):
            print(f"[app] Converted {ext} → WAV for Cleanvoice")
            return out
    except Exception as e:
        logger.warning(f"Format conversion failed: {e}")
    return path

# ══════════════════════════════════════════════════════════════════════
# CORE PIPELINE
# ══════════════════════════════════════════════════════════════════════
def run_pipeline(audio_path, src_lang="auto", tgt_lang="te",
                 opt_fillers=True, opt_stutters=True,
                 opt_silences=True, job_id=None):

    def progress(step, msg):
        u = {"status": "processing", "step": step, "message": msg}
        if job_id:
            _update_job(job_id, **u)
        return u

    out_dir = tempfile.mkdtemp()
    try:
        # ── Step 1: Denoise ──────────────────────────────────────────
        yield progress(1, "⏳ Step 1/4 — Cleaning audio (Cleanvoice)...")
        res   = get_denoiser().process(
            audio_path, out_dir,
            remove_fillers=opt_fillers,
            remove_stutters=opt_stutters,
            remove_silences=opt_silences,
        )
        clean = res["audio_path"]
        stats = res.get("stats", {})

        # ── Step 2: Transcribe ───────────────────────────────────────
        yield progress(2, "⏳ Step 2/4 — Transcribing...")
        t_result = get_transcriber().transcribe(clean, src_lang)
        if isinstance(t_result, tuple):
            transcript, detected_lang, t_method = t_result
        elif isinstance(t_result, dict):
            transcript    = t_result.get("text", "")
            detected_lang = t_result.get("language", src_lang)
            t_method      = t_result.get("method", "unknown")
        else:
            transcript    = str(t_result)
            detected_lang = src_lang
            t_method      = "unknown"

        if opt_fillers:
            transcript = get_denoiser().clean_transcript_fillers(transcript)

        # ── Step 3: Translate ────────────────────────────────────────
        translation = transcript
        tl_method   = "same language"

        if tgt_lang != "auto" and detected_lang != tgt_lang:
            yield progress(3, "⏳ Step 3/4 — Translating...")
            t_out = get_translator().translate(transcript, detected_lang, tgt_lang)
            if isinstance(t_out, tuple):
                translation, tl_method = t_out
            else:
                translation = str(t_out)

        # ── Step 4: Export MP3 ───────────────────────────────────────
        yield progress(4, "⏳ Step 4/4 — Exporting...")
        final_mp3 = os.path.join(out_dir, "ClearWave_Final.mp3")
        subprocess.run([
            "ffmpeg", "-y", "-i", clean,
            "-codec:a", "libmp3lame",
            "-b:a", "320k",
            "-ar", "48000",
            "-ac", "2",
            final_mp3
        ], capture_output=True)

        audio_path_out = final_mp3 if os.path.exists(final_mp3) else clean

        result = {
            "status":      "done",
            "step":        4,
            "message":     "✅ Done!",
            "transcript":  transcript,
            "translation": translation,
            "audioPath":   audio_path_out,
            "stats": {
                "language":             detected_lang.upper(),
                "transcription_method": t_method,
                "translation_method":   tl_method,
                "fillers_removed":      stats.get("fillers_removed", "via Cleanvoice"),
                "noise_method":         stats.get("noise_method", "Cleanvoice API"),
            },
        }
        if job_id:
            _update_job(job_id, status="done", step=4,
                        message="✅ Done!", result=result)
        yield result

    except Exception as e:
        logger.error(f"Pipeline failed: {e}", exc_info=True)
        err = {"status": "error", "message": f"❌ Error: {str(e)}"}
        if job_id:
            _update_job(job_id, **err)
        yield err

# ══════════════════════════════════════════════════════════════════════
# BACKGROUND WORKER
# ══════════════════════════════════════════════════════════════════════
def _bg_run(job_id, audio_path, src_lang, tgt_lang,
            opt_fillers, opt_stutters, opt_silences):
    try:
        for _ in run_pipeline(audio_path, src_lang, tgt_lang,
                               opt_fillers, opt_stutters, opt_silences,
                               job_id=job_id):
            pass
    except Exception as e:
        _update_job(job_id, status="error", message=f"❌ {e}")

# ══════════════════════════════════════════════════════════════════════
# GRADIO HANDLER
# ══════════════════════════════════════════════════════════════════════
def process_audio_gradio(upload_file, mic_audio, in_lang_name,
                          out_lang_name, opt_fillers, opt_stutters,
                          opt_silences):

    audio_path = None
    if upload_file is not None:
        if isinstance(upload_file, dict):
            audio_path = upload_file.get("name") or upload_file.get("path", "")
        else:
            audio_path = upload_file
    elif mic_audio is not None:
        if isinstance(mic_audio, str):
            audio_path = mic_audio
        elif isinstance(mic_audio, tuple):
            import soundfile as sf
            sr, data   = mic_audio
            tmp        = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
            sf.write(tmp.name, data, sr)
            tmp.close()
            audio_path = tmp.name
        elif isinstance(mic_audio, dict):
            audio_path = mic_audio.get("name") or mic_audio.get("path", "")

    if not audio_path:
        yield ("❌ Please upload a file or record via microphone.",
               "", "", None, "")
        return

    # Only convert truly exotic formats — leave mp3/wav/aac untouched
    audio_path = maybe_convert_for_cleanvoice(audio_path)

    opt_fillers  = bool(opt_fillers)
    opt_stutters = bool(opt_stutters)
    opt_silences = bool(opt_silences)

    src_lang = LANGUAGES.get(in_lang_name,  "auto")
    tgt_lang = LANGUAGES.get(out_lang_name, "te")

    print(f"[Gradio] src={src_lang} tgt={tgt_lang} "
          f"fillers={opt_fillers} stutters={opt_stutters}")

    job_id = _new_job()
    threading.Thread(
        target=_bg_run,
        args=(job_id, audio_path, src_lang, tgt_lang,
              opt_fillers, opt_stutters, opt_silences),
        daemon=True,
    ).start()

    while True:
        time.sleep(1)
        job    = _get_job(job_id)
        status = job.get("status", "queued")
        msg    = job.get("message", "Processing...")

        if status in ("queued", "downloading", "processing"):
            yield (msg, "", "", None, "")

        elif status == "done":
            r = job.get("result", {})
            s = r.get("stats", {})
            stats_str = "\n".join([
                f"🎙️  Language      : {s.get('language','?')}",
                f"🔊  Noise method  : {s.get('noise_method','?')}",
                f"🗑️  Fillers       : {s.get('fillers_removed','?')}",
                f"📝  Transcription : {s.get('transcription_method','?')}",
                f"🌐  Translation   : {s.get('translation_method','?')}",
            ])
            yield (r.get("message","✅ Done!"),
                   r.get("transcript",""),
                   r.get("translation",""),
                   r.get("audioPath"),
                   stats_str)
            return

        elif status == "error":
            yield (job.get("message","❌ Error"), "", "", None, "")
            return

# ══════════════════════════════════════════════════════════════════════
# GRADIO UI
# ══════════════════════════════════════════════════════════════════════
with gr.Blocks(title="ClearWave AI Studio") as demo:
    gr.Markdown("# 🎙️ ClearWave AI Studio\n### Professional Audio Enhancement")
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Tabs():
                with gr.Tab("📁 Upload File"):
                    audio_in = gr.File(
                        label="Upload Audio (MP3, WAV, MPEG, MP4, AAC, OGG, FLAC...)",
                        file_types=[
                            ".mp3", ".wav", ".mpeg", ".mpg", ".mp4",
                            ".m4a", ".aac", ".ogg", ".flac", ".opus",
                            ".webm", ".amr", ".wma", ".aiff",
                        ],
                    )
                with gr.Tab("🎙️ Record Microphone"):
                    mic_in = gr.Audio(
                        label="Record from Microphone",
                        sources=["microphone"],
                        type="filepath",
                    )

            with gr.Row():
                in_lang  = gr.Dropdown(
                    label="Input Language",
                    choices=list(LANGUAGES.keys()),
                    value="Auto Detect")
                out_lang = gr.Dropdown(
                    label="Output Language",
                    choices=list(OUT_LANGS.keys()),
                    value="Telugu")

            gr.Markdown("### 🛠️ Studio Options")
            with gr.Row():
                opt_f = gr.Checkbox(label="🗑️ Remove Fillers",  value=True)
                opt_s = gr.Checkbox(label="🔁 Remove Stutters", value=True)
            opt_z = gr.Checkbox(label="🤫 Remove Silences", value=True)

            btn        = gr.Button("🚀 Process Audio", variant="primary", size="lg")
            status_out = gr.Textbox(label="Status", interactive=False, lines=2)

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("📝 Text"):
                    with gr.Row():
                        transcript_out  = gr.Textbox(label="Transcript",  lines=12)
                        translation_out = gr.Textbox(label="Translation", lines=12)
                with gr.Tab("🔊 Clean Audio"):
                    audio_out = gr.Audio(label="Enhanced Audio", type="filepath")
                with gr.Tab("📊 Stats"):
                    stats_out = gr.Textbox(label="Stats", lines=10, interactive=False)

    btn.click(
        fn=process_audio_gradio,
        inputs=[audio_in, mic_in, in_lang, out_lang, opt_f, opt_s, opt_z],
        outputs=[status_out, transcript_out, translation_out, audio_out, stats_out],
    )

# ══════════════════════════════════════════════════════════════════════
# API ROUTES
# ══════════════════════════════════════════════════════════════════════
import json as _json
from fastapi import Request as _Req
from fastapi.responses import JSONResponse as _JSON

@demo.app.get("/api/health")
async def api_health():
    return _JSON({"status": "ok", "service": "ClearWave AI Studio",
                  "jobs_active": len(_jobs)})

@demo.app.post("/api/process-url")
async def api_process_url(request: _Req):
    data = await request.json()
    if "data" in data and isinstance(data["data"], dict):
        data = data["data"]

    audio_url    = data.get("audioUrl")
    audio_id     = data.get("audioId",     "")
    src_lang     = data.get("srcLang",     "auto")
    tgt_lang     = data.get("tgtLang",     "te")
    opt_fillers  = bool(data.get("optFillers",  True))
    opt_stutters = bool(data.get("optStutters", True))
    opt_silences = bool(data.get("optSilences", True))

    if not audio_url:
        return _JSON({"error": "audioUrl required"}, status_code=400)

    job_id = _new_job()
    _update_job(job_id, status="downloading", message="Downloading...")

    def _dl_and_run():
        tmp_path = audio_path = None
        try:
            r = req.get(audio_url, timeout=300, stream=True)
            r.raise_for_status()
            url_l  = audio_url.lower()
            suffix = (".wav"  if "wav"  in url_l else
                      ".mpeg" if "mpeg" in url_l else
                      ".mp4"  if "mp4"  in url_l else ".mp3")
            tmp      = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            tmp_path = tmp.name
            total    = int(r.headers.get("content-length", 0))
            dl       = 0
            for chunk in r.iter_content(chunk_size=65536):
                if chunk:
                    tmp.write(chunk)
                    dl += len(chunk)
                    if total:
                        pct = int(dl * 100 / total)
                        _update_job(job_id, status="downloading",
                                    message=f"Downloading... {pct}%")
            tmp.close()
            audio_path = maybe_convert_for_cleanvoice(tmp_path)
            for _ in run_pipeline(audio_path, src_lang, tgt_lang,
                                  opt_fillers, opt_stutters, opt_silences,
                                  job_id=job_id):
                pass
            with _jobs_lock:
                if job_id in _jobs and _jobs[job_id].get("result"):
                    _jobs[job_id]["result"]["audioId"] = audio_id
        except Exception as e:
            logger.error(f"Job {job_id} failed: {e}", exc_info=True)
            _update_job(job_id, status="error", message=f"❌ {e}")
        finally:
            for p in [tmp_path, audio_path]:
                try:
                    if p and p != tmp_path and os.path.exists(p):
                        os.unlink(p)
                except Exception:
                    pass

    threading.Thread(target=_dl_and_run, daemon=True).start()
    return _JSON({"jobId": job_id, "audioId": audio_id,
                  "status": "queued",
                  "pollUrl": f"/api/job/{job_id}",
                  "message": "Job started! Poll pollUrl for progress."})

@demo.app.get("/api/job/{job_id}")
async def api_get_job(job_id: str):
    job = _get_job(job_id)
    if not job:
        return _JSON({"error": "Job not found"}, status_code=404)
    resp = {"jobId": job_id, "status": job.get("status"),
            "step": job.get("step", 0), "message": job.get("message", "")}
    if job.get("status") == "done":
        resp["result"] = job.get("result", {})
    return _JSON(resp)

@demo.app.get("/api/jobs")
async def api_list_jobs():
    with _jobs_lock:
        s = {k: {"status": v["status"], "step": v.get("step", 0)}
             for k, v in _jobs.items()}
    return _JSON({"jobs": s, "total": len(s)})

logger.info("✅ ClearWave AI Studio ready — services lazy-loaded on first use")

if __name__ == "__main__":
    demo.launch(ssr_mode=False)