Ari123er/Ai_Speech_Synthesis

1"""
2Caption Generator — Vosk Speech-to-Text
3========================================
4Generates word-level captions (JSON) from audio chunks using Vosk.
5Optimized for HuggingFace Free Tier (16GB RAM, 2 vCPU).
6
7Usage (imported by app.py):
8    from caption_generator import generate_captions
9    captions = generate_captions(trimmed_chunks, sr)
10"""
11
12import os
13import gc
14import json
15import wave
16import struct
17import tempfile
18import numpy as np
19
20# ═══════════════════════════════════════════════════════════════════════════
21# CONFIGURATION
22# ═══════════════════════════════════════════════════════════════════════════
23
24# Vosk model — smallest English model (~40MB download, ~300MB RAM)
25VOSK_MODEL_NAME = "vosk-model-small-en-us-0.15"
26VOSK_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vosk-models")
27VOSK_SR = 16000  # Vosk requires 16kHz mono 16-bit PCM
28
29# Singleton model instance
30_vosk_model = None
31
32
33# ═══════════════════════════════════════════════════════════════════════════
34# MODEL LOADING
35# ═══════════════════════════════════════════════════════════════════════════
36
37def _load_vosk_model():
38    """
39    Lazy-load Vosk model. Downloads on first call, caches on disk.
40    Returns the Model instance.
41    """
42    global _vosk_model
43    if _vosk_model is not None:
44        return _vosk_model
45
46    from vosk import Model, SetLogLevel
47
48    # Suppress Vosk's verbose logging
49    SetLogLevel(-1)
50
51    model_path = os.path.join(VOSK_MODEL_DIR, VOSK_MODEL_NAME)
52
53    if not os.path.isdir(model_path):
54        print(f"[VOSK] Downloading model '{VOSK_MODEL_NAME}'...")
55        os.makedirs(VOSK_MODEL_DIR, exist_ok=True)
56
57        # vosk.Model auto-downloads if given just the model name
58        _vosk_model = Model(model_name=VOSK_MODEL_NAME)
59        print(f"[VOSK] Model downloaded and loaded.")
60    else:
61        print(f"[VOSK] Loading cached model from {model_path}")
62        _vosk_model = Model(model_path=model_path)
63        print(f"[VOSK] Model loaded.")
64
65    return _vosk_model
66
67
68# ═══════════════════════════════════════════════════════════════════════════
69# AUDIO CONVERSION
70# ═══════════════════════════════════════════════════════════════════════════
71
72def _audio_to_pcm16(audio: np.ndarray, sr: int) -> bytes:
73    """
74    Convert float32 numpy audio to 16kHz mono 16-bit PCM bytes.
75    Resamples if sr != 16000. This is what Vosk expects.
76    """
77    # Resample to 16kHz if needed
78    if sr != VOSK_SR:
79        import librosa
80        audio = librosa.resample(audio, orig_sr=sr, target_sr=VOSK_SR)
81
82    # Ensure mono
83    if audio.ndim > 1:
84        audio = np.mean(audio, axis=1)
85
86    # Clip and convert float32 → int16
87    audio = np.clip(audio, -1.0, 1.0)
88    pcm_data = (audio * 32767).astype(np.int16)
89
90    return pcm_data.tobytes()
91
92
93# ═══════════════════════════════════════════════════════════════════════════
94# CAPTION GENERATION
95# ═══════════════════════════════════════════════════════════════════════════
96
97def generate_captions(trimmed_chunks: list, sr: int) -> list:
98    """
99    Generate word-level captions from trimmed audio chunks using Vosk.
100
101    Args:
102        trimmed_chunks: List of dicts with keys:
103            - 'audio': np.ndarray (trimmed float32 audio at original sr)
104            - 'text':  str (original script text for this chunk)
105            - 'tag':   str (style tag, e.g. 'SHOCK', 'FAST')
106        sr: Sample rate of the audio (e.g. 24000)
107
108    Returns:
109        Flat list of word dicts:
110        [
111            {"word": "Dragon", "start": 0.0, "end": 0.32, "conf": 0.98},
112            {"word": "Ball", "start": 0.32, "end": 0.55, "conf": 0.99},
113            ...
114        ]
115    """
116    from vosk import KaldiRecognizer
117
118    model = _load_vosk_model()
119
120    print(f"\n{'='*60}")
121    print(f"  CAPTION GENERATOR — Vosk STT")
122    print(f"{'='*60}")
123
124    all_words = []
125    time_offset = 0.0  # Running offset for multi-chunk positioning
126
127    for i, chunk in enumerate(trimmed_chunks):
128        audio = chunk['audio']
129        script_text = chunk['text']
130        tag = chunk['tag']
131
132        chunk_duration = len(audio) / sr
133        print(f"  ▸ [{i+1}/{len(trimmed_chunks)}] [{tag}] {chunk_duration:.2f}s — \"{script_text[:50]}...\"")
134
135        # Convert to 16kHz PCM
136        pcm_bytes = _audio_to_pcm16(audio, sr)
137
138        # Create recognizer with word-level timestamps
139        rec = KaldiRecognizer(model, VOSK_SR)
140        rec.SetWords(True)
141
142        # Feed audio in small chunks (4000 bytes ≈ 125ms at 16kHz 16-bit)
143        FEED_SIZE = 4000
144        for pos in range(0, len(pcm_bytes), FEED_SIZE):
145            rec.AcceptWaveform(pcm_bytes[pos:pos + FEED_SIZE])
146
147        # Get final result
148        result = json.loads(rec.FinalResult())
149
150        # Extract words with offset-adjusted timestamps
151        if 'result' in result:
152            for w in result['result']:
153                all_words.append({
154                    "word": w['word'],
155                    "start": round(w['start'] + time_offset, 3),
156                    "end": round(w['end'] + time_offset, 3),
157                    "conf": round(w.get('conf', 0.0), 3)
158                })
159
160        time_offset += chunk_duration
161
162        # Free recognizer memory
163        del rec, pcm_bytes
164        gc.collect()
165
166        print(f"    → {len(all_words)} total words so far")
167
168    print(f"{'='*60}")
169    print(f"  DONE — {len(all_words)} total words across {len(trimmed_chunks)} chunk(s)")
170    print(f"{'='*60}\n")
171
172    return all_words
173"""
174Standalone module for Vosk STT captions.
175Uses the smallest English model for efficient CPU inference.
176"""