migrate to mikroservices architecture, and add more features
This commit is contained in:
jadssss
208
core/pipeline.py
Normal file
208
core/pipeline.py
Normal file
@@ -0,0 +1,208 @@
|
||||
"""Orchestrate convert -> separate -> transcribe over a directory of audio files."""
|
||||
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]
|
||||
StageFn = Callable[[str, str], None]
|
||||
ProgressFn = Callable[[int, int], None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileResult:
|
||||
input_path: Path
|
||||
text: Optional[str]
|
||||
output_dir: Optional[Path]
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class Cancelled(Exception):
|
||||
"""Raised internally when the stop event is set."""
|
||||
|
||||
|
||||
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 _check_cancel(stop_event: Optional[threading.Event]) -> None:
|
||||
if stop_event is not None and stop_event.is_set():
|
||||
raise Cancelled()
|
||||
|
||||
|
||||
def _final_stages(config: PipelineConfig) -> set[str]:
|
||||
"""Which stage produces the final artifact that must survive cleanup."""
|
||||
if config.transcribe.enabled:
|
||||
return {"transcribe"}
|
||||
if config.separate.enabled:
|
||||
return {"separate"}
|
||||
if config.convert.enabled:
|
||||
return {"convert"}
|
||||
return set()
|
||||
|
||||
|
||||
_ARTIFACT_NAMES = {
|
||||
"convert": "convert.wav",
|
||||
"separate": "vocals.wav",
|
||||
"transcribe": "text.txt",
|
||||
}
|
||||
|
||||
|
||||
def run_pipeline(
|
||||
config: PipelineConfig,
|
||||
log: LogFn = print,
|
||||
on_stage: Optional[StageFn] = None,
|
||||
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)
|
||||
total = len(files)
|
||||
results: list[FileResult] = []
|
||||
|
||||
if total == 0:
|
||||
log("Аудиофайлы не найдены.")
|
||||
return results
|
||||
|
||||
log(f"Найдено треков: {total}")
|
||||
|
||||
for idx, input_path in enumerate(files, start=1):
|
||||
_check_cancel(stop_event)
|
||||
on_progress and on_progress(idx - 1, total)
|
||||
log(f"\n▶ [{idx}/{total}] {input_path.name}")
|
||||
try:
|
||||
results.append(_process_one(input_path, config, log, on_stage, stop_event))
|
||||
except Cancelled:
|
||||
log("⏹ Отменено пользователем")
|
||||
break
|
||||
except Exception as exc: # noqa: BLE001 — we log per-file and continue
|
||||
log(f" ✗ Ошибка: {exc}")
|
||||
results.append(FileResult(input_path, None, None, str(exc)))
|
||||
finally:
|
||||
on_progress and on_progress(idx, total)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _process_one(
|
||||
input_path: Path,
|
||||
config: PipelineConfig,
|
||||
log: LogFn,
|
||||
on_stage: Optional[StageFn],
|
||||
stop_event: Optional[threading.Event],
|
||||
) -> FileResult:
|
||||
track_stem = input_path.stem
|
||||
track_out = Path(config.output_dir) / track_stem
|
||||
track_out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
tmp_work = Path(tempfile.mkdtemp(prefix=f"at_{track_stem}_"))
|
||||
finals = _final_stages(config)
|
||||
|
||||
try:
|
||||
# ---------- 1. convert ----------
|
||||
source_wav: Path = input_path
|
||||
if config.convert.enabled:
|
||||
_check_cancel(stop_event)
|
||||
on_stage and on_stage("convert", input_path.name)
|
||||
log(" [1] Конвертация в WAV...")
|
||||
target = (
|
||||
track_out / "convert.wav"
|
||||
if ("convert" in finals or config.keep_intermediate)
|
||||
else tmp_work / "convert.wav"
|
||||
)
|
||||
if config.overwrite or not target.exists():
|
||||
convert_to_wav(
|
||||
input_path, target,
|
||||
sample_rate=config.convert.sample_rate,
|
||||
channels=config.convert.channels,
|
||||
ffmpeg_path=config.convert.ffmpeg_path,
|
||||
)
|
||||
source_wav = target
|
||||
log(f" ✓ {target.name}")
|
||||
|
||||
# ---------- 2. separate ----------
|
||||
vocal_path: Path = source_wav
|
||||
if config.separate.enabled:
|
||||
_check_cancel(stop_event)
|
||||
on_stage and on_stage("separate", input_path.name)
|
||||
log(" [2] Отделение вокала...")
|
||||
raw_vocal = separate_vocals(
|
||||
source_wav, tmp_work,
|
||||
model=config.separate.model,
|
||||
isolate_stem=config.separate.isolate_stem,
|
||||
unblend_path=config.separate.unblend_path,
|
||||
custom_kernels=config.separate.custom_kernels,
|
||||
)
|
||||
if "separate" in finals or config.keep_intermediate:
|
||||
vocal_path = track_out / "vocals.wav"
|
||||
shutil.move(str(raw_vocal), str(vocal_path))
|
||||
else:
|
||||
vocal_path = raw_vocal # stays in tmp_work
|
||||
log(f" ✓ {vocal_path.name}")
|
||||
|
||||
# ---------- 3. transcribe ----------
|
||||
text: Optional[str] = None
|
||||
if config.transcribe.enabled:
|
||||
_check_cancel(stop_event)
|
||||
on_stage and on_stage("transcribe", input_path.name)
|
||||
log(" [3] Транскрибация...")
|
||||
text = transcribe(
|
||||
vocal_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,
|
||||
)
|
||||
(track_out / "text.txt").write_text(text, encoding="utf-8")
|
||||
log(f" ✓ text.txt ({len(text)} символов)")
|
||||
|
||||
if not config.keep_intermediate:
|
||||
_cleanup_keeping(track_out, finals)
|
||||
|
||||
return FileResult(input_path, text, track_out)
|
||||
|
||||
finally:
|
||||
shutil.rmtree(tmp_work, ignore_errors=True)
|
||||
|
||||
|
||||
def _cleanup_keeping(track_out: Path, finals: set[str]) -> None:
|
||||
if not track_out.exists():
|
||||
return
|
||||
keep = {_ARTIFACT_NAMES[s] for s in finals if s in _ARTIFACT_NAMES}
|
||||
for p in track_out.iterdir():
|
||||
if p.name in keep:
|
||||
continue
|
||||
try:
|
||||
if p.is_file():
|
||||
p.unlink()
|
||||
elif p.is_dir():
|
||||
shutil.rmtree(p, ignore_errors=True)
|
||||
except OSError:
|
||||
pass
|
||||
Reference in New Issue
Block a user