OnTheWay/ChickenRice-Transcribe
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add ASMR VAD inference support and model assets

Browse source
This commit is contained in:
xuwei 2026-03-13 22:16:57 +08:00
parent 26ea9aba51
commit d605aadef3
9 changed files with 1259 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

300
infer.py
View file

@ -1,12 +1,93 @@
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Iterable
import librosa
import numpy as np
from faster_whisper import WhisperModel
class AsmrVadModel:
def __init__(
self,
model_path: Path,
force_cpu: bool = False,
num_threads: int = 1,
) -> None:
try:
import onnxruntime as ort
except ImportError as exc:
raise RuntimeError("onnxruntime is required for --vad-mode asmr") from exc
try:
from transformers import WhisperFeatureExtractor
except ImportError as exc:
raise RuntimeError("transformers is required for --vad-mode asmr") from exc
metadata_path = model_path.with_name("model_metadata.json")
metadata = {
"whisper_model_name": "openai/whisper-base",
"frame_duration_ms": 20,
"total_duration_ms": 30000,
}
if metadata_path.exists():
metadata.update(json.loads(metadata_path.read_text(encoding="utf-8")))
self.sample_rate = 16000
self.frame_duration_ms = int(metadata.get("frame_duration_ms", 20))
self.chunk_duration_ms = int(metadata.get("total_duration_ms", 30000))
self.chunk_samples = int(self.chunk_duration_ms * self.sample_rate / 1000)
opts = ort.SessionOptions()
opts.inter_op_num_threads = num_threads
opts.intra_op_num_threads = num_threads
providers = ["CPUExecutionProvider"]
if not force_cpu and "CUDAExecutionProvider" in ort.get_available_providers():
providers.insert(0, "CUDAExecutionProvider")
whisper_model_name = metadata.get("whisper_model_name", "openai/whisper-base")
local_whisper_path = Path(whisper_model_name)
if local_whisper_path.exists():
feature_extractor_source = str(local_whisper_path)
elif Path("model").exists():
feature_extractor_source = "model"
else:
feature_extractor_source = whisper_model_name
self.feature_extractor = WhisperFeatureExtractor.from_pretrained(feature_extractor_source)
self.session = ort.InferenceSession(str(model_path), providers=providers, sess_options=opts)
self.input_name = self.session.get_inputs()[0].name
self.output_names = [output.name for output in self.session.get_outputs()]
self.providers = self.session.get_providers()
def load_audio(self, audio_path: Path) -> np.ndarray:
audio, _ = librosa.load(str(audio_path), sr=self.sample_rate, mono=True)
return audio.astype(np.float32, copy=False)
def predict_probabilities(self, audio: np.ndarray) -> np.ndarray:
probabilities: list[np.ndarray] = []
for start in range(0, len(audio), self.chunk_samples):
chunk = audio[start : start + self.chunk_samples]
if len(chunk) < self.chunk_samples:
chunk = np.pad(chunk, (0, self.chunk_samples - len(chunk)), mode="constant")
features = self.feature_extractor(
chunk,
sampling_rate=self.sample_rate,
return_tensors="np",
).input_features
logits = self.session.run(self.output_names, {self.input_name: features})[0][0]
probabilities.append(1.0 / (1.0 + np.exp(-logits)))
if not probabilities:
return np.array([], dtype=np.float32)
return np.concatenate(probabilities, axis=0)
def load_model(
device: str,
compute_type: str,
@ -40,6 +121,79 @@ def format_timestamp(seconds: float) -> str:
return f"[{minutes:02d}:{secs:02d}.{centiseconds:02d}]"
def detect_asmr_speech_segments(
audio: np.ndarray,
vad_model: AsmrVadModel,
threshold: float,
neg_threshold: float | None,
min_speech_duration_ms: int,
min_silence_duration_ms: int,
speech_pad_ms: int,
) -> list[dict[str, float]]:
speech_probs = vad_model.predict_probabilities(audio)
if speech_probs.size == 0:
return []
frame_ms = vad_model.frame_duration_ms
min_speech_frames = max(1, int(round(min_speech_duration_ms / frame_ms)))
min_silence_frames = max(1, int(round(min_silence_duration_ms / frame_ms)))
speech_pad_frames = max(0, int(round(speech_pad_ms / frame_ms)))
neg_threshold = max(threshold - 0.15, 0.01) if neg_threshold is None else neg_threshold
raw_segments: list[tuple[int, int]] = []
triggered = False
current_start = 0
temp_end: int | None = None
for frame_idx, speech_prob in enumerate(speech_probs):
if speech_prob >= threshold and not triggered:
triggered = True
current_start = frame_idx
temp_end = None
continue
if not triggered:
continue
if speech_prob < neg_threshold:
if temp_end is None:
temp_end = frame_idx
elif frame_idx - temp_end >= min_silence_frames:
if temp_end - current_start >= min_speech_frames:
raw_segments.append((current_start, temp_end))
triggered = False
temp_end = None
elif temp_end is not None:
temp_end = None
if triggered:
end_frame = temp_end if temp_end is not None else len(speech_probs)
if end_frame - current_start >= min_speech_frames:
raw_segments.append((current_start, end_frame))
segments: list[dict[str, float]] = []
for idx, (start_frame, end_frame) in enumerate(raw_segments):
prev_end = raw_segments[idx - 1][1] if idx > 0 else 0
next_start = raw_segments[idx + 1][0] if idx + 1 < len(raw_segments) else len(speech_probs)
padded_start = max(prev_end, start_frame - speech_pad_frames)
padded_end = min(next_start, end_frame + speech_pad_frames)
segments.append(
{
"start": padded_start * frame_ms / 1000,
"end": padded_end * frame_ms / 1000,
}
)
return segments
def build_clip_timestamps(segments: list[dict[str, float]]) -> list[float]:
clip_timestamps: list[float] = []
for segment in segments:
clip_timestamps.extend([segment["start"], segment["end"]])
return clip_timestamps
def write_lrc(segments: Iterable, output_path: Path) -> None:
lines = []
for seg in segments:
@ -57,17 +211,21 @@ def transcribe_file(
model: WhisperModel,
audio_path: Path,
beam_size: int,
language: str,
task: str,
vad: bool,
vad_parameters: dict | None,
clip_timestamps: str | list[float] = "0",
extra_generation_args: dict[str, Any] | None = None,
) -> tuple[list, str, float]:
segments_iter, info = model.transcribe(
str(audio_path),
task="translate",
task=task,
beam_size=beam_size,
vad_filter=vad,
vad_parameters=vad_parameters if vad else None,
language="ja",
clip_timestamps=clip_timestamps,
language=language,
**(extra_generation_args or {}),
)
print(f"[{audio_path.name}] Detected language: {info.language} (prob={info.language_probability:.2f})")
@ -98,7 +256,24 @@ def parse_args() -> argparse.Namespace:
help="Path to an audio file or directory (default: ./mp3).",
)
parser.add_argument("--beam-size", type=int, default=5, help="Beam size for decoding.")
parser.add_argument("--no-vad", action="store_true", help="Disable VAD (voice activity detection) filtering.")
parser.add_argument(
"--language",
default="ja",
help="Input language code passed to Whisper (default: ja). Use 'auto' for auto-detection.",
)
parser.add_argument(
"--task",
default="translate",
choices=["transcribe", "translate"],
help="Whisper task mode (default: translate).",
)
parser.add_argument(
"--vad-mode",
default="asmr",
choices=["asmr", "builtin", "none"],
help="VAD mode: asmr ONNX model, faster-whisper builtin VAD, or none (default: asmr).",
)
parser.add_argument("--no-vad", action="store_true", help=argparse.SUPPRESS)
parser.add_argument(
"--vad-threshold",
type=float,
@ -111,17 +286,51 @@ def parse_args() -> argparse.Namespace:
default=None,
help="Optional silence probability threshold for VAD (useful to smooth speech end).",
)
parser.add_argument(
"--vad-min-speech-ms",
type=int,
default=300,
help="Minimum speech duration (ms) kept by VAD (default: 300).",
)
parser.add_argument(
"--vad-min-silence-ms",
type=int,
default=400,
help="Minimum silence (ms) to cut a speech chunk when VAD is enabled (default: 400).",
default=100,
help="Minimum silence (ms) to cut a speech chunk when VAD is enabled (default: 100).",
)
parser.add_argument(
"--vad-pad-ms",
type=int,
default=500,
help="Padding (ms) added before/after each detected speech chunk (default: 500).",
default=200,
help="Padding (ms) added before/after each detected speech chunk (default: 200).",
)
parser.add_argument(
"--vad-model-path",
default="vad_models/Whisper-Vad-EncDec-ASMR-onnx/model.onnx",
help="Path to the external ASMR VAD ONNX model.",
)
parser.add_argument(
"--vad-force-cpu",
action="store_true",
help="Force the external ASMR VAD to run on CPU.",
)
parser.add_argument(
"--vad-num-threads",
type=int,
default=1,
help="CPU thread count for the external ASMR VAD (default: 1).",
)
parser.add_argument(
"--max-initial-timestamp",
type=float,
default=30.0,
help="Maximum initial timestamp passed to Whisper decoding (default: 30).",
)
parser.add_argument(
"--repetition-penalty",
type=float,
default=1.1,
help="Repetition penalty passed to Whisper decoding (default: 1.1).",
)
parser.add_argument(
"--device",
@ -148,6 +357,8 @@ def parse_args() -> argparse.Namespace:
def main() -> None:
args = parse_args()
if args.no_vad:
args.vad_mode = "none"
target = Path(args.audio)
audio_files = collect_audio_files(target)
@ -166,37 +377,86 @@ def main() -> None:
else:
print(f"Using device={device_used}, compute_type={compute_used}")
vad_parameters = {
builtin_vad_parameters = {
"threshold": args.vad_threshold,
"neg_threshold": args.vad_neg_threshold,
"min_speech_duration_ms": args.vad_min_speech_ms,
"min_silence_duration_ms": args.vad_min_silence_ms,
"speech_pad_ms": args.vad_pad_ms,
}
if args.no_vad:
vad_parameters = None
else:
asmr_vad: AsmrVadModel | None = None
if args.vad_mode == "builtin":
print(
"VAD enabled: "
f"threshold={vad_parameters['threshold']}, "
f"neg_threshold={vad_parameters['neg_threshold']}, "
f"min_silence_ms={vad_parameters['min_silence_duration_ms']}, "
f"pad_ms={vad_parameters['speech_pad_ms']}"
"Built-in VAD enabled: "
f"threshold={builtin_vad_parameters['threshold']}, "
f"neg_threshold={builtin_vad_parameters['neg_threshold']}, "
f"min_speech_ms={builtin_vad_parameters['min_speech_duration_ms']}, "
f"min_silence_ms={builtin_vad_parameters['min_silence_duration_ms']}, "
f"pad_ms={builtin_vad_parameters['speech_pad_ms']}"
)
elif args.vad_mode == "asmr":
vad_model_path = Path(args.vad_model_path)
if not vad_model_path.exists():
raise FileNotFoundError(f"ASMR VAD model not found: {vad_model_path}")
asmr_vad = AsmrVadModel(
model_path=vad_model_path,
force_cpu=args.vad_force_cpu,
num_threads=args.vad_num_threads,
)
print(
"ASMR VAD enabled: "
f"model={vad_model_path}, "
f"providers={asmr_vad.providers}, "
f"threshold={args.vad_threshold}, "
f"neg_threshold={args.vad_neg_threshold}, "
f"min_speech_ms={args.vad_min_speech_ms}, "
f"min_silence_ms={args.vad_min_silence_ms}, "
f"pad_ms={args.vad_pad_ms}"
)
else:
print("VAD disabled")
generation_args: dict[str, Any] = {
"max_initial_timestamp": 30,
"repetition_penalty": 1.1,
"max_initial_timestamp": args.max_initial_timestamp,
"repetition_penalty": args.repetition_penalty,
}
for idx, audio_path in enumerate(audio_files, start=1):
print(f"\n[{idx}/{len(audio_files)}] Processing {audio_path}")
try:
use_builtin_vad = args.vad_mode == "builtin"
clip_timestamps: str | list[float] = "0"
if args.vad_mode == "asmr" and asmr_vad is not None:
audio = asmr_vad.load_audio(audio_path)
speech_segments = detect_asmr_speech_segments(
audio=audio,
vad_model=asmr_vad,
threshold=args.vad_threshold,
neg_threshold=args.vad_neg_threshold,
min_speech_duration_ms=args.vad_min_speech_ms,
min_silence_duration_ms=args.vad_min_silence_ms,
speech_pad_ms=args.vad_pad_ms,
)
if speech_segments:
kept_duration = sum(segment["end"] - segment["start"] for segment in speech_segments)
print(
"ASMR VAD kept "
f"{len(speech_segments)} segments "
f"({kept_duration:.2f}s speech)"
)
clip_timestamps = build_clip_timestamps(speech_segments)
else:
print("ASMR VAD found no speech segments; falling back to full-audio decoding.")
segments, _, _ = transcribe_file(
model=model,
audio_path=audio_path,
beam_size=args.beam_size,
vad=not args.no_vad,
vad_parameters=vad_parameters,
language=args.language,
task=args.task,
vad=use_builtin_vad,
vad_parameters=builtin_vad_parameters if use_builtin_vad else None,
clip_timestamps=clip_timestamps,
extra_generation_args=generation_args,
)
write_lrc(segments, audio_path.with_suffix(".lrc"))