387 lines
14 KiB
Python
387 lines
14 KiB
Python
"""Orchestrate convert -> separate -> transcribe.
|
|
|
|
Two modes:
|
|
- sequential: for each file, run all enabled stages before moving to the next.
|
|
- staged: run stage 1 over all files, then stage 2 over all files, then stage 3.
|
|
|
|
Intermediate artifacts (convert.wav, vocals.wav) are written to a temp dir and
|
|
NEVER persisted to the output directory — unless `keep_intermediate=True`.
|
|
Only the artifact of the LAST enabled stage is moved into `<output>/<stem>/`.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import tempfile
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Optional
|
|
|
|
from .config import PipelineConfig
|
|
from .convert import convert_to_wav
|
|
from .separate import separate_vocals
|
|
from .transcribe import transcribe
|
|
|
|
|
|
AUDIO_EXTENSIONS = {
|
|
".mp3", ".wav", ".flac", ".ogg", ".m4a", ".aac", ".opus", ".wma",
|
|
}
|
|
|
|
LogFn = Callable[[str], None]
|
|
ProgressFn = Callable[[int, int, str], None] # done, total, stage
|
|
|
|
|
|
@dataclass
|
|
class FileResult:
|
|
input_path: Path
|
|
text: Optional[str]
|
|
output_dir: Optional[Path]
|
|
error: Optional[str] = None
|
|
|
|
|
|
class Cancelled(Exception):
|
|
"""Raised when the stop event is set."""
|
|
|
|
|
|
# ---------------------------------------------------------------- helpers
|
|
|
|
_ARTIFACT_NAMES = {
|
|
"convert": "convert.wav",
|
|
"separate": "vocals.wav",
|
|
"transcribe": "text.txt",
|
|
}
|
|
|
|
|
|
def collect_audio_files(input_dir: Path, recursive: bool) -> list[Path]:
|
|
it = input_dir.rglob("*") if recursive else input_dir.iterdir()
|
|
return sorted(p for p in it if p.is_file() and p.suffix.lower() in AUDIO_EXTENSIONS)
|
|
|
|
|
|
def _final_stage(config: PipelineConfig) -> Optional[str]:
|
|
"""The stage whose artifact is the final result (single element)."""
|
|
if config.transcribe.enabled:
|
|
return "transcribe"
|
|
if config.separate.enabled:
|
|
return "separate"
|
|
if config.convert.enabled:
|
|
return "convert"
|
|
return None
|
|
|
|
|
|
def _check_cancel(stop_event: Optional[threading.Event]) -> None:
|
|
if stop_event is not None and stop_event.is_set():
|
|
raise Cancelled()
|
|
|
|
|
|
def _persist_final(work_dir: Path, persist_dir: Path, final_stage: str, log: LogFn) -> None:
|
|
"""Move the final-stage artifact from work_dir into persist_dir."""
|
|
name = _ARTIFACT_NAMES[final_stage]
|
|
src = work_dir / name
|
|
if not src.exists():
|
|
raise FileNotFoundError(f"Final artifact missing: {src}")
|
|
persist_dir.mkdir(parents=True, exist_ok=True)
|
|
dst = persist_dir / name
|
|
if dst.exists():
|
|
dst.unlink()
|
|
shutil.move(str(src), str(dst))
|
|
log(f" [persist] {final_stage} -> {persist_dir.name}/{name}")
|
|
|
|
|
|
def _work_dir_for(
|
|
config: PipelineConfig,
|
|
stem: str,
|
|
output_root: Path,
|
|
tmp_root: Optional[Path],
|
|
) -> tuple[Path, Path]:
|
|
"""Return (work_dir, persist_dir).
|
|
|
|
- keep_intermediate=True -> work_dir == persist_dir == output_root/stem
|
|
- keep_intermediate=False -> work_dir = tmp_root/stem, persist_dir = output_root/stem
|
|
"""
|
|
persist_dir = output_root / stem
|
|
if config.keep_intermediate:
|
|
persist_dir.mkdir(parents=True, exist_ok=True)
|
|
return persist_dir, persist_dir
|
|
if tmp_root is None:
|
|
raise ValueError("tmp_root required when keep_intermediate=False")
|
|
work_dir = tmp_root / stem
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
return work_dir, persist_dir
|
|
|
|
|
|
# ---------------------------------------------------------------- entry
|
|
|
|
def run_pipeline(
|
|
config: PipelineConfig,
|
|
log: LogFn = print,
|
|
on_progress: Optional[ProgressFn] = None,
|
|
stop_event: Optional[threading.Event] = None,
|
|
) -> list[FileResult]:
|
|
input_dir = Path(config.input_dir)
|
|
if not input_dir.is_dir():
|
|
raise FileNotFoundError(f"Input directory not found: {input_dir}")
|
|
|
|
output_dir = Path(config.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
files = collect_audio_files(input_dir, config.recursive)
|
|
if not files:
|
|
log("Аудиофайлы не найдены.")
|
|
return []
|
|
|
|
mode = (config.mode or "sequential").lower()
|
|
log(f"Найдено треков: {len(files)} | режим: {mode} | "
|
|
f"keep_intermediate={config.keep_intermediate}")
|
|
if not config.keep_intermediate:
|
|
log(" Промежуточные файлы будут во временной папке и удалены после прогона.")
|
|
|
|
if mode == "staged":
|
|
return _run_staged(config, files, log, on_progress, stop_event)
|
|
return _run_sequential(config, files, log, on_progress, stop_event)
|
|
|
|
|
|
# ---------------------------------------------------------------- sequential
|
|
|
|
def _run_sequential(
|
|
config: PipelineConfig,
|
|
files: list[Path],
|
|
log: LogFn,
|
|
on_progress: Optional[ProgressFn],
|
|
stop_event: Optional[threading.Event],
|
|
) -> list[FileResult]:
|
|
total = len(files)
|
|
final_stage = _final_stage(config)
|
|
results: list[FileResult] = []
|
|
output_root = Path(config.output_dir)
|
|
|
|
for idx, f in enumerate(files, 1):
|
|
if stop_event is not None and stop_event.is_set():
|
|
log("⏹ Отменено")
|
|
break
|
|
if on_progress:
|
|
on_progress(idx - 1, total, "file")
|
|
log(f"\n▶ [{idx}/{total}] {f.name}")
|
|
|
|
tmp_root = Path(tempfile.mkdtemp(prefix=f"at_{f.stem}_"))
|
|
work_dir, persist_dir = _work_dir_for(config, f.stem, output_root, tmp_root)
|
|
|
|
audio_path: Path = f
|
|
text: Optional[str] = None
|
|
error: Optional[str] = None
|
|
processed_ok = False
|
|
|
|
try:
|
|
if config.convert.enabled:
|
|
_check_cancel(stop_event)
|
|
log(" [1] Конвертация в WAV...")
|
|
target = work_dir / "convert.wav"
|
|
if config.overwrite or not target.exists():
|
|
convert_to_wav(
|
|
f, target,
|
|
sample_rate=config.convert.sample_rate,
|
|
channels=config.convert.channels,
|
|
ffmpeg_path=config.convert.ffmpeg_path,
|
|
)
|
|
audio_path = target
|
|
log(" ✓ convert.wav")
|
|
|
|
if config.separate.enabled:
|
|
_check_cancel(stop_event)
|
|
log(" [2] Отделение вокала...")
|
|
target = work_dir / "vocals.wav"
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
raw = separate_vocals(
|
|
audio_path, Path(tmp),
|
|
model=config.separate.model,
|
|
isolate_stem=config.separate.isolate_stem,
|
|
unblend_path=config.separate.unblend_path,
|
|
custom_kernels=config.separate.custom_kernels,
|
|
)
|
|
shutil.move(str(raw), str(target))
|
|
audio_path = target
|
|
log(" ✓ vocals.wav")
|
|
|
|
if config.transcribe.enabled:
|
|
_check_cancel(stop_event)
|
|
log(" [3] Транскрибация...")
|
|
text = transcribe(
|
|
audio_path,
|
|
engine=config.transcribe.engine,
|
|
model=config.transcribe.model,
|
|
language=config.transcribe.language,
|
|
device=config.transcribe.device,
|
|
compute_type=config.transcribe.compute_type,
|
|
batch_size=config.transcribe.batch_size,
|
|
gigaam_provider=config.transcribe.gigaam_provider,
|
|
local_files_only=config.transcribe.local_files_only,
|
|
vad_filter=config.transcribe.vad_filter,
|
|
)
|
|
(work_dir / "text.txt").write_text(text, encoding="utf-8")
|
|
log(f" ✓ text.txt ({len(text)} символов)")
|
|
|
|
processed_ok = True
|
|
|
|
except Cancelled:
|
|
log("⏹ Отменено")
|
|
break
|
|
except Exception as e: # noqa: BLE001
|
|
error = f"{type(e).__name__}: {e}"
|
|
log(f" ✗ Ошибка: {error}")
|
|
finally:
|
|
if processed_ok and not config.keep_intermediate and final_stage is not None:
|
|
try:
|
|
_persist_final(work_dir, persist_dir, final_stage, log)
|
|
except Exception as e: # noqa: BLE001
|
|
error = f"persist failed: {e}"
|
|
log(f" ✗ {error}")
|
|
shutil.rmtree(tmp_root, ignore_errors=True)
|
|
if on_progress:
|
|
on_progress(idx, total, "file")
|
|
results.append(FileResult(f, text, persist_dir if not error else None, error))
|
|
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------- staged
|
|
|
|
def _run_staged(
|
|
config: PipelineConfig,
|
|
files: list[Path],
|
|
log: LogFn,
|
|
on_progress: Optional[ProgressFn],
|
|
stop_event: Optional[threading.Event],
|
|
) -> list[FileResult]:
|
|
output_root = Path(config.output_dir)
|
|
total = len(files)
|
|
final_stage = _final_stage(config)
|
|
|
|
# One tmp_root for the whole batch: stage 2 needs convert.wav of every file,
|
|
# stage 3 needs vocals.wav of every file. All of it is removed at the end.
|
|
tmp_root_ctx = tempfile.mkdtemp(prefix="at_staged_")
|
|
tmp_root = Path(tmp_root_ctx)
|
|
|
|
state: dict[Path, dict] = {}
|
|
for f in files:
|
|
work_dir, persist_dir = _work_dir_for(config, f.stem, output_root, tmp_root)
|
|
state[f] = {
|
|
"audio_path": f,
|
|
"work_dir": work_dir,
|
|
"persist_dir": persist_dir,
|
|
"text": None,
|
|
"error": None,
|
|
}
|
|
|
|
def active() -> list[Path]:
|
|
return [f for f in files if state[f]["error"] is None]
|
|
|
|
def cancelled() -> bool:
|
|
return stop_event is not None and stop_event.is_set()
|
|
|
|
try:
|
|
# ---- Stage 1: convert ----
|
|
if config.convert.enabled:
|
|
log(f"\n══ Этап 1/3 · Конвертация ({total} файлов) ══")
|
|
for i, f in enumerate(files, 1):
|
|
if cancelled():
|
|
log("⏹ Отменено")
|
|
break
|
|
if on_progress:
|
|
on_progress(i, total, "convert")
|
|
try:
|
|
target = state[f]["work_dir"] / "convert.wav"
|
|
if config.overwrite or not target.exists():
|
|
convert_to_wav(
|
|
f, target,
|
|
sample_rate=config.convert.sample_rate,
|
|
channels=config.convert.channels,
|
|
ffmpeg_path=config.convert.ffmpeg_path,
|
|
)
|
|
state[f]["audio_path"] = target
|
|
log(f" ✓ {f.name}")
|
|
except Exception as e: # noqa: BLE001
|
|
state[f]["error"] = f"{type(e).__name__}: {e}"
|
|
log(f" ✗ {f.name}: {e}")
|
|
|
|
# ---- Stage 2: separate ----
|
|
if config.separate.enabled and not cancelled():
|
|
todo = active()
|
|
log(f"\n══ Этап 2/3 · Отделение вокала ({len(todo)} файлов) ══")
|
|
for i, f in enumerate(todo, 1):
|
|
if cancelled():
|
|
log("⏹ Отменено")
|
|
break
|
|
if on_progress:
|
|
on_progress(i, len(todo), "separate")
|
|
try:
|
|
target = state[f]["work_dir"] / "vocals.wav"
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
raw = separate_vocals(
|
|
state[f]["audio_path"], Path(tmp),
|
|
model=config.separate.model,
|
|
isolate_stem=config.separate.isolate_stem,
|
|
unblend_path=config.separate.unblend_path,
|
|
custom_kernels=config.separate.custom_kernels,
|
|
)
|
|
shutil.move(str(raw), str(target))
|
|
state[f]["audio_path"] = target
|
|
log(f" ✓ {f.name}")
|
|
except Exception as e: # noqa: BLE001
|
|
state[f]["error"] = f"{type(e).__name__}: {e}"
|
|
log(f" ✗ {f.name}: {e}")
|
|
|
|
# ---- Stage 3: transcribe ----
|
|
if config.transcribe.enabled and not cancelled():
|
|
todo = active()
|
|
log(f"\n══ Этап 3/3 · Транскрибация ({len(todo)} файлов) ══")
|
|
for i, f in enumerate(todo, 1):
|
|
if cancelled():
|
|
log("⏹ Отменено")
|
|
break
|
|
if on_progress:
|
|
on_progress(i, len(todo), "transcribe")
|
|
try:
|
|
text = transcribe(
|
|
state[f]["audio_path"],
|
|
engine=config.transcribe.engine,
|
|
model=config.transcribe.model,
|
|
language=config.transcribe.language,
|
|
device=config.transcribe.device,
|
|
compute_type=config.transcribe.compute_type,
|
|
batch_size=config.transcribe.batch_size,
|
|
gigaam_provider=config.transcribe.gigaam_provider,
|
|
local_files_only=config.transcribe.local_files_only,
|
|
vad_filter=config.transcribe.vad_filter,
|
|
)
|
|
(state[f]["work_dir"] / "text.txt").write_text(text, encoding="utf-8")
|
|
state[f]["text"] = text
|
|
log(f" ✓ {f.name} ({len(text)} символов)")
|
|
except Exception as e: # noqa: BLE001
|
|
state[f]["error"] = f"{type(e).__name__}: {e}"
|
|
log(f" ✗ {f.name}: {e}")
|
|
|
|
# ---- Persist final artifacts ----
|
|
if not config.keep_intermediate and final_stage is not None:
|
|
log(f"\n══ Перенос финальных файлов ({final_stage}) ══")
|
|
for f in files:
|
|
s = state[f]
|
|
if s["error"] is not None:
|
|
continue
|
|
try:
|
|
_persist_final(s["work_dir"], s["persist_dir"], final_stage, log)
|
|
except Exception as e: # noqa: BLE001
|
|
s["error"] = f"persist failed: {e}"
|
|
log(f" ✗ {f.name}: {e}")
|
|
|
|
finally:
|
|
shutil.rmtree(tmp_root, ignore_errors=True)
|
|
|
|
return [
|
|
FileResult(
|
|
f,
|
|
state[f]["text"],
|
|
state[f]["persist_dir"] if state[f]["error"] is None else None,
|
|
state[f]["error"],
|
|
)
|
|
for f in files
|
|
]
|