nite
a565382a52
A conversational video-editing planning assistant. Provides tools for: - Audio transcription via bundled funasr-script (Fun-ASR-Nano / SenseVoice) - On-demand clip extraction and frame sampling (ffmpeg, hardware-accelerated) - SQLite index to track all artifacts and avoid duplicate processing - Vision analysis guidance with binary-search-style frame sampling - Iterative Markdown-table edit plan output Agent-agnostic: no platform-specific tool names, works with any agent runtime. Dependencies checked at guidance level with OS package manager install hints.
366 lines
12 KiB
Python
366 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
REGULAR_MODEL_CONFIGS: dict[str, dict[str, Any]] = {
|
|
"sensevoice": {
|
|
"model": "iic/SenseVoiceSmall",
|
|
"vad_model": "fsmn-vad",
|
|
"vad_kwargs": {"max_single_segment_time": 30000},
|
|
},
|
|
"paraformer": {
|
|
"model": "paraformer-zh",
|
|
"vad_model": "fsmn-vad",
|
|
"punc_model": "ct-punc",
|
|
},
|
|
"paraformer-en": {
|
|
"model": "paraformer-en",
|
|
"vad_model": "fsmn-vad",
|
|
},
|
|
}
|
|
|
|
NANO_MODEL_CONFIG: dict[str, Any] = {
|
|
"model": "FunAudioLLM/Fun-ASR-Nano-2512",
|
|
"vad_model": "fsmn-vad",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class AudioTrack:
|
|
index: int
|
|
codec: str | None
|
|
channels: int | None
|
|
channel_layout: str | None
|
|
language: str | None
|
|
title: str | None
|
|
|
|
|
|
def require_command(name: str) -> None:
|
|
if shutil.which(name) is None:
|
|
raise SystemExit(f"缺少命令: {name}")
|
|
|
|
|
|
def run_command(cmd: list[str], *, capture: bool = True) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
cmd,
|
|
check=True,
|
|
text=True,
|
|
stdout=subprocess.PIPE if capture else None,
|
|
stderr=subprocess.PIPE if capture else None,
|
|
)
|
|
|
|
|
|
def list_audio_tracks(media_path: Path) -> list[AudioTrack]:
|
|
require_command("ffprobe")
|
|
proc = run_command(
|
|
[
|
|
"ffprobe",
|
|
"-v",
|
|
"error",
|
|
"-select_streams",
|
|
"a",
|
|
"-show_entries",
|
|
"stream=index,codec_name,channels,channel_layout:stream_tags=language,title",
|
|
"-of",
|
|
"json",
|
|
str(media_path),
|
|
]
|
|
)
|
|
data = json.loads(proc.stdout)
|
|
tracks: list[AudioTrack] = []
|
|
for stream in data.get("streams", []):
|
|
tags = stream.get("tags") or {}
|
|
tracks.append(
|
|
AudioTrack(
|
|
index=int(stream["index"]),
|
|
codec=stream.get("codec_name"),
|
|
channels=stream.get("channels"),
|
|
channel_layout=stream.get("channel_layout"),
|
|
language=tags.get("language"),
|
|
title=tags.get("title"),
|
|
)
|
|
)
|
|
return tracks
|
|
|
|
|
|
def print_audio_tracks(media_path: Path) -> None:
|
|
tracks = list_audio_tracks(media_path)
|
|
if not tracks:
|
|
print("未找到音频轨")
|
|
return
|
|
for i, track in enumerate(tracks, 1):
|
|
parts = [
|
|
f"#{i}",
|
|
f"stream_index={track.index}",
|
|
f"codec={track.codec or '?'}",
|
|
f"channels={track.channels or '?'}",
|
|
]
|
|
if track.channel_layout:
|
|
parts.append(f"layout={track.channel_layout}")
|
|
if track.language:
|
|
parts.append(f"lang={track.language}")
|
|
if track.title:
|
|
parts.append(f"title={track.title}")
|
|
print(" ".join(parts))
|
|
|
|
|
|
def resolve_track(media_path: Path, requested: int | None) -> int:
|
|
tracks = list_audio_tracks(media_path)
|
|
if not tracks:
|
|
raise SystemExit(f"未找到音频轨: {media_path}")
|
|
if requested is None:
|
|
return tracks[0].index
|
|
valid_indexes = {track.index for track in tracks}
|
|
if requested in valid_indexes:
|
|
return requested
|
|
raise SystemExit(
|
|
f"音轨 stream index 不存在: {requested}\n"
|
|
f"可用音轨: {', '.join(str(t.index) for t in tracks)}\n"
|
|
"提示: --track 使用 ffprobe/ffmpeg 的 stream index，不是第几条音轨的序号。"
|
|
)
|
|
|
|
|
|
def extract_audio(media_path: Path, wav_path: Path, track_index: int) -> None:
|
|
require_command("ffmpeg")
|
|
wav_path.parent.mkdir(parents=True, exist_ok=True)
|
|
cmd = [
|
|
"ffmpeg",
|
|
"-hide_banner",
|
|
"-y",
|
|
"-i",
|
|
str(media_path),
|
|
"-map",
|
|
f"0:{track_index}",
|
|
"-vn",
|
|
"-ar",
|
|
"16000",
|
|
"-ac",
|
|
"1",
|
|
"-c:a",
|
|
"pcm_s16le",
|
|
str(wav_path),
|
|
]
|
|
subprocess.run(cmd, check=True)
|
|
|
|
|
|
def clean_text(text: str, *, sensevoice: bool) -> str:
|
|
if sensevoice:
|
|
try:
|
|
from funasr.utils.postprocess_utils import rich_transcription_postprocess
|
|
|
|
text = rich_transcription_postprocess(text)
|
|
except Exception:
|
|
# Keep transcription usable even if FunASR changes this helper.
|
|
pass
|
|
return re.sub(r"<\|[^|]*\|>", "", text).strip()
|
|
|
|
|
|
def get_audio_duration(audio_path: Path) -> float | None:
|
|
try:
|
|
import soundfile as sf
|
|
|
|
return round(float(sf.info(str(audio_path)).duration), 3)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def normalize_result(
|
|
raw_result: list[dict[str, Any]],
|
|
*,
|
|
audio_path: Path,
|
|
source_media: Path,
|
|
track_index: int,
|
|
model_name: str,
|
|
model_id: str,
|
|
language: str | None,
|
|
diarize: bool,
|
|
elapsed: float,
|
|
sensevoice: bool,
|
|
) -> dict[str, Any]:
|
|
first = raw_result[0] if raw_result else {}
|
|
text = clean_text(str(first.get("text", "")), sensevoice=sensevoice)
|
|
segments: list[dict[str, Any]] = []
|
|
for seg in first.get("sentence_info") or []:
|
|
item: dict[str, Any] = {
|
|
"start": seg.get("start", 0),
|
|
"end": seg.get("end", 0),
|
|
"text": clean_text(str(seg.get("sentence") or seg.get("text") or ""), sensevoice=sensevoice),
|
|
}
|
|
if diarize and "spk" in seg:
|
|
item["speaker"] = seg["spk"]
|
|
segments.append(item)
|
|
|
|
output: dict[str, Any] = {
|
|
"text": text,
|
|
"segments": segments,
|
|
"file": source_media.name,
|
|
"audio_file": audio_path.name,
|
|
"track": track_index,
|
|
"model": model_name,
|
|
"model_id": model_id,
|
|
"language": language or "auto",
|
|
"audio_duration_s": get_audio_duration(audio_path),
|
|
"processing_s": round(elapsed, 3),
|
|
}
|
|
for key in ("timestamps", "timestamp"):
|
|
if key in first:
|
|
output[key] = first[key]
|
|
return output
|
|
|
|
|
|
def write_json(path: Path, data: dict[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def make_output_path(output_dir: Path, media_path: Path, track_index: int, suffix: str) -> Path:
|
|
return output_dir / f"{media_path.stem}_track{track_index}_{suffix}.json"
|
|
|
|
|
|
def resolve_output_dir(output_dir: Path | None, media_path: Path) -> Path:
|
|
if output_dir is None:
|
|
return media_path.parent
|
|
return output_dir.expanduser().resolve()
|
|
|
|
|
|
def add_common_args(parser: argparse.ArgumentParser) -> None:
|
|
parser.add_argument("media", type=Path, help="视频/音频文件路径")
|
|
parser.add_argument("--track", type=int, default=None, help="ffmpeg stream index；默认第一条音频轨")
|
|
parser.add_argument("--list-tracks", action="store_true", help="列出音轨后退出")
|
|
parser.add_argument("--output-dir", "-o", type=Path, default=None, help="输出目录；默认源视频/音频文件所在目录")
|
|
parser.add_argument("--language", "-l", default="auto", help="语言，如 auto/zh/en/ja/ko/yue；默认 auto")
|
|
parser.add_argument("--device", default=None, help="设备，如 cuda:0/cpu；默认自动")
|
|
parser.add_argument("--no-diarize", action="store_true", help="禁用 cam++ 说话人分离")
|
|
parser.add_argument("--keep-wav", action="store_true", help="保留提取出的 16k wav 文件")
|
|
parser.add_argument("--wav-dir", type=Path, default=None, help="临时 wav 输出目录；默认系统临时目录")
|
|
parser.add_argument("--verbose", "-v", action="store_true", help="输出更多进度信息")
|
|
|
|
|
|
def prepare_audio(args: argparse.Namespace) -> tuple[Path, int, tempfile.TemporaryDirectory[str] | None]:
|
|
media_path: Path = args.media.expanduser().resolve()
|
|
if not media_path.exists():
|
|
raise SystemExit(f"文件不存在: {media_path}")
|
|
|
|
if args.list_tracks:
|
|
print_audio_tracks(media_path)
|
|
raise SystemExit(0)
|
|
|
|
track_index = resolve_track(media_path, args.track)
|
|
temp_dir: tempfile.TemporaryDirectory[str] | None = None
|
|
if args.keep_wav:
|
|
wav_dir = args.wav_dir.expanduser().resolve() if args.wav_dir else resolve_output_dir(args.output_dir, media_path)
|
|
wav_dir.mkdir(parents=True, exist_ok=True)
|
|
elif args.wav_dir is not None:
|
|
wav_dir = args.wav_dir.expanduser().resolve()
|
|
wav_dir.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
temp_dir = tempfile.TemporaryDirectory(prefix="funasr-script-")
|
|
wav_dir = Path(temp_dir.name)
|
|
|
|
wav_path = wav_dir / f"{media_path.stem}_track{track_index}.wav"
|
|
print(f"[1/2] 提取音轨 stream_index={track_index} -> {wav_path}", file=sys.stderr)
|
|
extract_audio(media_path, wav_path, track_index)
|
|
return wav_path, track_index, temp_dir
|
|
|
|
|
|
def auto_device(device: str | None) -> str:
|
|
if device:
|
|
return device
|
|
import torch
|
|
|
|
return "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def transcribe_with_config(
|
|
*,
|
|
wav_path: Path,
|
|
media_path: Path,
|
|
track_index: int,
|
|
output_path: Path,
|
|
model_name: str,
|
|
config: dict[str, Any],
|
|
language: str | None,
|
|
device: str | None,
|
|
diarize: bool,
|
|
use_itn: bool,
|
|
sensevoice: bool,
|
|
verbose: bool,
|
|
) -> dict[str, Any]:
|
|
from funasr import AutoModel
|
|
|
|
config = config.copy()
|
|
if diarize and "spk_model" not in config:
|
|
config["spk_model"] = "cam++"
|
|
resolved_device = auto_device(device)
|
|
|
|
print(f"[2/2] 加载模型 {model_name} ({config['model']}) on {resolved_device}", file=sys.stderr)
|
|
load_start = time.time()
|
|
model = AutoModel(device=resolved_device, disable_update=True, **config)
|
|
if verbose:
|
|
print(f"模型加载耗时: {time.time() - load_start:.1f}s", file=sys.stderr)
|
|
|
|
gen_kw: dict[str, Any] = {"input": str(wav_path), "batch_size": 1}
|
|
if language and language != "auto":
|
|
gen_kw["language"] = language
|
|
if use_itn:
|
|
gen_kw["use_itn"] = True
|
|
|
|
start = time.time()
|
|
raw_result = model.generate(**gen_kw)
|
|
elapsed = time.time() - start
|
|
output = normalize_result(
|
|
raw_result,
|
|
audio_path=wav_path,
|
|
source_media=media_path,
|
|
track_index=track_index,
|
|
model_name=model_name,
|
|
model_id=str(config["model"]),
|
|
language=language,
|
|
diarize=diarize,
|
|
elapsed=elapsed,
|
|
sensevoice=sensevoice,
|
|
)
|
|
write_json(output_path, output)
|
|
print(f"完成: {output_path}", file=sys.stderr)
|
|
print(f"转录耗时: {elapsed:.2f}s", file=sys.stderr)
|
|
return output
|
|
|
|
|
|
def run_pipeline(args: argparse.Namespace, *, model_name: str, config: dict[str, Any], suffix: str, use_itn: bool, sensevoice: bool) -> Path:
|
|
media_path = args.media.expanduser().resolve()
|
|
temp_dir: tempfile.TemporaryDirectory[str] | None = None
|
|
try:
|
|
wav_path, track_index, temp_dir = prepare_audio(args)
|
|
output_path = make_output_path(resolve_output_dir(args.output_dir, media_path), media_path, track_index, suffix)
|
|
transcribe_with_config(
|
|
wav_path=wav_path,
|
|
media_path=media_path,
|
|
track_index=track_index,
|
|
output_path=output_path,
|
|
model_name=model_name,
|
|
config=config,
|
|
language=args.language,
|
|
device=args.device,
|
|
diarize=not args.no_diarize,
|
|
use_itn=use_itn,
|
|
sensevoice=sensevoice,
|
|
verbose=args.verbose,
|
|
)
|
|
print(output_path)
|
|
return output_path
|
|
finally:
|
|
if temp_dir is not None:
|
|
temp_dir.cleanup()
|