102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""
|
|
Subtitle Generation Module using Faster-Whisper on IntelLLM.
|
|
Generates word-level synchronized .srt subtitle tracks in English and Spanish.
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
import logging
|
|
from typing import Dict, Any
|
|
from .remote import RemotePC
|
|
|
|
logger = logging.getLogger("YTFactory.Subs")
|
|
|
|
WHISPER_WORKER_SCRIPT = r'''
|
|
import sys
|
|
from faster_whisper import WhisperModel
|
|
|
|
audio_path = sys.argv[1]
|
|
lang = sys.argv[2]
|
|
out_srt = sys.argv[3]
|
|
|
|
# Load faster-whisper on CPU using int8 quantization
|
|
model = WhisperModel("base", device="cpu", compute_type="int8")
|
|
segments, _ = model.transcribe(audio_path, language=lang, word_timestamps=True)
|
|
|
|
def format_timestamp(seconds):
|
|
ms = int(round(seconds * 1000))
|
|
hours, ms = divmod(ms, 3600000)
|
|
minutes, ms = divmod(ms, 60000)
|
|
secs, ms = divmod(ms, 1000)
|
|
return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"
|
|
|
|
words = []
|
|
for seg in segments:
|
|
words.extend(seg.words or [])
|
|
|
|
cues = []
|
|
current_words = []
|
|
start_time = None
|
|
last_end_time = None
|
|
|
|
for w in words:
|
|
if current_words and (len(current_words) >= 6 or (w.start - start_time) > 3.2):
|
|
cues.append((start_time, last_end_time, " ".join(x.word.strip() for x in current_words)))
|
|
current_words = []
|
|
|
|
if not current_words:
|
|
start_time = w.start
|
|
|
|
current_words.append(w)
|
|
last_end_time = w.end
|
|
|
|
if current_words:
|
|
cues.append((start_time, last_end_time, " ".join(x.word.strip() for x in current_words)))
|
|
|
|
with open(out_srt, "w", encoding="utf-8") as f:
|
|
for i, (s, e, text) in enumerate(cues, 1):
|
|
f.write(f"{i}\n{format_timestamp(s)} --> {format_timestamp(e)}\n{text}\n\n")
|
|
|
|
print(f"GENERATED_{len(cues)}_CUES")
|
|
'''
|
|
|
|
class SubGen:
|
|
def __init__(self, pcs: Dict[str, RemotePC]):
|
|
self.pcs = pcs
|
|
|
|
def generate(self, audio_path: str, job_dir: str, lang: str = "en") -> str:
|
|
"""Generate accurate .srt subtitle file on IntelLLM."""
|
|
logger.info(f"[SUBS:{lang.upper()}] Transcribing audio with Faster-Whisper on IntelLLM...")
|
|
pc = self.pcs.get("intelllm")
|
|
if not pc:
|
|
raise RuntimeError("IntelLLM is not available for Whisper transcription")
|
|
|
|
ts = int(time.time())
|
|
remote_audio = f"/tmp/whisper_in_{lang}_{ts}.wav"
|
|
remote_srt = f"/tmp/whisper_out_{lang}_{ts}.srt"
|
|
remote_script = f"/tmp/run_whisper_{lang}_{ts}.py"
|
|
|
|
pc.scp_to(audio_path, remote_audio)
|
|
|
|
# Upload whisper worker script
|
|
rc, _, err = pc.run(f"cat > {remote_script} << 'PYEOF'\n{WHISPER_WORKER_SCRIPT.strip()}\nPYEOF")
|
|
if rc != 0:
|
|
raise RuntimeError(f"Failed to deploy whisper worker script: {err}")
|
|
|
|
# Execute in intel venv
|
|
cmd = (
|
|
f"source ~/venvs/youtube-factory-intel/bin/activate 2>/dev/null || true; "
|
|
f"python3 {remote_script} {remote_audio} {lang} {remote_srt}"
|
|
)
|
|
rc, stdout, stderr = pc.run(cmd, timeout=900)
|
|
if rc != 0:
|
|
logger.error(f"[SUBS:{lang.upper()}] Whisper transcription failed: {stderr}")
|
|
raise RuntimeError(f"Whisper failed: {stderr}")
|
|
|
|
local_srt = os.path.join(job_dir, f"subtitles_{lang}.srt")
|
|
pc.scp_from(remote_srt, local_srt)
|
|
|
|
# Cleanup remote tmp files
|
|
pc.run(f"rm -f {remote_audio} {remote_srt} {remote_script}")
|
|
logger.info(f"[SUBS:{lang.upper()}] Subtitles generated: {local_srt} ({stdout.strip()})")
|
|
return local_srt
|