capcut-agent/capcut_agent/transcribe.py

"""faster-whisper 기반 전체 대본 전사 (세그먼트 + 단어 타임스탬프).

★ 핵심 원칙(참고 영상 08:11): 전체 대본을 먼저 뽑아 맥락을 잡는다.
   → 이 대본이 NG 탐지 · 자막 · 숏폼 하이라이트 선택의 공통 재료.

함정:
- ASR(numba)은 동시 호출 시 segfault → 호출부(pipeline)에서 asyncio.Lock 직렬화.
- 캐시는 content hash 기준(mtime 아님) → 같은 영상 재업로드 시 hit.
"""
from __future__ import annotations

import hashlib
import json
import os
from dataclasses import asdict, dataclass
from typing import List, Optional

CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".cache")
os.makedirs(CACHE_DIR, exist_ok=True)

# 모델 인스턴스 재사용(로딩 비쌈). (model_size, compute_type) → WhisperModel
_MODEL_CACHE: dict = {}


@dataclass
class Word:
    start: float
    end: float
    text: str


@dataclass
class Seg:
    start: float
    end: float
    text: str
    words: List[Word]


@dataclass
class Transcript:
    language: str
    text: str            # 전체 대본 (한 덩어리)
    segments: List[Seg]

    def to_json(self) -> dict:
        return {
            "language": self.language,
            "text": self.text,
            "segments": [
                {"start": s.start, "end": s.end, "text": s.text,
                 "words": [asdict(w) for w in s.words]}
                for s in self.segments
            ],
        }

    @staticmethod
    def from_json(d: dict) -> "Transcript":
        return Transcript(
            language=d["language"], text=d["text"],
            segments=[
                Seg(s["start"], s["end"], s["text"],
                    [Word(**w) for w in s["words"]])
                for s in d["segments"]
            ],
        )


def _file_hash(path: str) -> str:
    h = hashlib.sha1()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()[:12]


def _get_model(model_size: str, compute_type: str):
    key = (model_size, compute_type)
    if key not in _MODEL_CACHE:
        from faster_whisper import WhisperModel  # 지연 import (무거움)
        _MODEL_CACHE[key] = WhisperModel(model_size, device="cpu", compute_type=compute_type)
    return _MODEL_CACHE[key]


def transcribe(
    video_path: str,
    *,
    model_size: str = "medium",
    language: Optional[str] = "ko",
    compute_type: str = "int8",
    use_cache: bool = True,
    vad_filter: bool = True,
) -> Transcript:
    """영상 전사 → Transcript. content-hash 캐시 적용.

    vad_filter=True: 무음 제거(환각↓) 하지만 짧은 외침을 놓칠 수 있음.
    vad_filter=False: 모든 발화 포착(자막 커버리지↑) — 컷이 오디오 기준일 때 자막용으로 적합.
    동기 함수. 호출부에서 ASR_LOCK 으로 직렬화하고 thread 로 돌릴 것.
    """
    h = _file_hash(video_path)
    tag = ("vad" if vad_filter else "novad") + "2"  # 2 = 환각억제(no_repeat) 설정
    cache_path = os.path.join(CACHE_DIR, f"asr_{model_size}_{tag}_{h}.json")
    if use_cache and os.path.exists(cache_path):
        with open(cache_path, encoding="utf-8") as f:
            return Transcript.from_json(json.load(f))

    model = _get_model(model_size, compute_type)
    seg_iter, info = model.transcribe(
        video_path,
        language=language,
        word_timestamps=True,
        vad_filter=vad_filter,
        vad_parameters={"min_silence_duration_ms": 400} if vad_filter else None,
        condition_on_previous_text=False,   # 세그먼트 간 환각 드리프트 억제
        no_repeat_ngram_size=3,             # 긴 어구 환각 루프(외국어 등) 차단. 단어 반복은 허용
    )

    segments: List[Seg] = []
    for s in seg_iter:            # 제너레이터 → 여기서 실제 추론 진행
        words = [Word(w.start, w.end, w.word) for w in (s.words or [])]
        segments.append(Seg(s.start, s.end, s.text.strip(), words))

    full_text = " ".join(s.text for s in segments).strip()
    tr = Transcript(language=info.language, text=full_text, segments=segments)

    if use_cache:
        with open(cache_path, "w", encoding="utf-8") as f:
            json.dump(tr.to_json(), f, ensure_ascii=False, indent=1)
    return tr