uxoxo/eb2ab

1"""2Background worker for TTS processing.3"""4 5import os6import sys7import threading8import traceback9import uuid10from datetime import datetime11from pathlib import Path12from queue import Queue, Empty13from typing import Dict, Optional14 15# Add parent directory to path to import lib modules16sys.path.insert(0, str(Path(__file__).parent.parent.parent))17 18from lib.functions import SessionContext, convert_ebook19from lib import FULL_DOCKER20from api.models import JobStatus21from api.storage import (22    save_text_input,23    save_job_metadata,24    get_audio_file_path,25    get_job_directory26)27 28 29# Job state storage30_jobs: Dict[str, dict] = {}31_jobs_lock = threading.Lock()32 33# Job queue34_job_queue: Queue = Queue(maxsize=50)35 36# Worker thread37_worker_thread: Optional[threading.Thread] = None38 39 40def create_job(41    text: str,42    language: str,43    voice: Optional[str],44    engine: str,45    output_format: str,46    speed: Optional[float] = None,47    temperature: Optional[float] = None48) -> str:49    """50    Create a new TTS job and add it to the queue.51 52    Args:53        text: Text to convert54        language: Language code55        voice: Voice file path or None56        engine: TTS engine name57        output_format: Output audio format58        speed: Speech speed (optional)59        temperature: Temperature (optional)60 61    Returns:62        Job ID63 64    Raises:65        Exception: If queue is full66    """67    job_id = str(uuid.uuid4())68    created_at = datetime.now()69 70    # Save text input71    text_file = save_text_input(job_id, text)72 73    # Save metadata74    save_job_metadata(job_id, created_at)75 76    # Create job record77    job = {78        "job_id": job_id,79        "status": JobStatus.QUEUED,80        "progress": 0,81        "error": None,82        "text_file": str(text_file),83        "language": language,84        "voice": voice,85        "engine": engine,86        "output_format": output_format,87        "speed": speed,88        "temperature": temperature,89        "created_at": created_at,90        "audio_file": None91    }92 93    # Store job94    with _jobs_lock:95        _jobs[job_id] = job96 97    # Add to queue98    try:99        _job_queue.put_nowait(job_id)100    except Exception as e:101        # Queue is full102        with _jobs_lock:103            _jobs[job_id]["status"] = JobStatus.FAILED104            _jobs[job_id]["error"] = "Job queue is full. Please try again later."105        raise Exception("Job queue is full")106 107    return job_id108 109 110def get_job_status(job_id: str) -> Optional[dict]:111    """112    Get current job status.113 114    Args:115        job_id: Job identifier116 117    Returns:118        Job status dict or None if not found119    """120    with _jobs_lock:121        return _jobs.get(job_id)122 123 124def _process_job(job_id: str):125    """126    Process a TTS job.127 128    Args:129        job_id: Job identifier130    """131    with _jobs_lock:132        job = _jobs.get(job_id)133        if not job:134            return135 136        job["status"] = JobStatus.PROCESSING137        job["progress"] = 10138 139    try:140        # Create session context141        ctx = SessionContext()142        session_id = str(uuid.uuid4())143 144        # Get voice file path if voice name is provided145        voice_path = None146        if job["voice"]:147            from lib import voices_dir148            # Search for voice file in voices directory149            voices_path = Path(voices_dir)150            for voice_file in voices_path.rglob("*.wav"):151                # Match voice name (case-insensitive, check if voice name is in filename)152                if job["voice"].lower().replace(" ", "_") in voice_file.stem.lower():153                    voice_path = str(voice_file)154                    break155 156        # Prepare conversion arguments157        args = {158            "script_mode": FULL_DOCKER,159            "session": session_id,160            "headless": True,161            "is_gui_process": False,162            "ebook": job["text_file"],163            "language": job["language"],164            "voice": voice_path,165            "device": os.environ.get("DEFAULT_DEVICE", "cpu"),166            "tts_engine": job["engine"],167            "custom_model": None,168            "fine_tuned": None,169            "output_format": job["output_format"],170            "temperature": job["temperature"],171            "length_penalty": None,172            "num_beams": None,173            "repetition_penalty": None,174            "top_k": None,175            "top_p": None,176            "speed": job["speed"],177            "enable_text_splitting": False,178            "text_temp": None,179            "waveform_temp": None,180            "audiobooks_dir": str(get_job_directory(job_id)),181            "output_split": False,182            "output_split_hours": "6",183            "ebook_list": None184        }185 186        # Update progress187        with _jobs_lock:188            _jobs[job_id]["progress"] = 20189 190        # Run conversion191        progress_status, passed = convert_ebook(args, ctx)192 193        if not passed:194            raise Exception(f"Conversion failed: {progress_status}")195 196        # Find generated audio file197        job_dir = get_job_directory(job_id)198        audio_files = list(job_dir.glob(f"*.{job['output_format']}"))199 200        if not audio_files:201            raise Exception(f"No audio file generated with format {job['output_format']}")202 203        audio_file = audio_files[0]204 205        # Update job status206        with _jobs_lock:207            _jobs[job_id]["status"] = JobStatus.COMPLETED208            _jobs[job_id]["progress"] = 100209            _jobs[job_id]["audio_file"] = str(audio_file)210 211        print(f"Job {job_id} completed successfully")212 213    except Exception as e:214        error_msg = str(e)215        traceback.print_exc()216 217        with _jobs_lock:218            _jobs[job_id]["status"] = JobStatus.FAILED219            _jobs[job_id]["error"] = error_msg220            _jobs[job_id]["progress"] = 0221 222        print(f"Job {job_id} failed: {error_msg}")223 224 225def _worker_loop():226    """Worker thread main loop."""227    print("TTS worker thread started")228 229    while True:230        try:231            # Wait for job (blocking)232            job_id = _job_queue.get(timeout=1)233 234            print(f"Processing job: {job_id}")235            _process_job(job_id)236 237            _job_queue.task_done()238 239        except Empty:240            # No jobs in queue, continue waiting241            continue242        except Exception as e:243            print(f"Worker error: {e}")244            traceback.print_exc()245 246 247def start_worker():248    """Start background worker thread."""249    global _worker_thread250 251    if _worker_thread is not None and _worker_thread.is_alive():252        print("Worker thread already running")253        return254 255    _worker_thread = threading.Thread(256        target=_worker_loop,257        daemon=True,258        name="TTSWorker"259    )260    _worker_thread.start()261    print("TTS worker thread initialized")262 263 264def get_queue_size() -> int:265    """Get current queue size."""266    return _job_queue.qsize()267 268 269def get_total_jobs() -> int:270    """Get total number of jobs."""271    with _jobs_lock:272        return len(_jobs)273