diff --git a/core/config.py b/core/config.py index 1772d40..7b54f16 100644 --- a/core/config.py +++ b/core/config.py @@ -30,13 +30,15 @@ class SeparateConfig: @dataclass class TranscribeConfig: enabled: bool = True - engine: str = "whisperx" # "whisperx" | "gigaam" + engine: str = "whisperx" model: str = "large-v3" - language: Optional[str] = "ru" # None = autodetect - device: str = "auto" # "cuda" | "cpu" | "auto" - compute_type: Optional[str] = None # None = auto + language: Optional[str] = "ru" + device: str = "auto" + compute_type: Optional[str] = None batch_size: int = 8 - gigaam_provider: str = "auto" # "auto" | "cuda" | "cpu" + gigaam_provider: str = "auto" + local_files_only: bool = False + vad_filter: bool = True @dataclass diff --git a/core/pipeline.py b/core/pipeline.py index 4aeff02..f41ec23 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -3,6 +3,10 @@ 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 `//`. """ from __future__ import annotations @@ -53,14 +57,15 @@ def collect_audio_files(input_dir: Path, recursive: bool) -> list[Path]: return sorted(p for p in it if p.is_file() and p.suffix.lower() in AUDIO_EXTENSIONS) -def _final_stages(config: PipelineConfig) -> set[str]: +def _final_stage(config: PipelineConfig) -> Optional[str]: + """The stage whose artifact is the final result (single element).""" if config.transcribe.enabled: - return {"transcribe"} + return "transcribe" if config.separate.enabled: - return {"separate"} + return "separate" if config.convert.enabled: - return {"convert"} - return set() + return "convert" + return None def _check_cancel(stop_event: Optional[threading.Event]) -> None: @@ -68,20 +73,40 @@ def _check_cancel(stop_event: Optional[threading.Event]) -> None: raise Cancelled() -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 +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 @@ -105,7 +130,10 @@ def run_pipeline( return [] mode = (config.mode or "sequential").lower() - log(f"Найдено треков: {len(files)} | режим: {mode}") + 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) @@ -122,8 +150,9 @@ def _run_sequential( stop_event: Optional[threading.Event], ) -> list[FileResult]: total = len(files) - finals = _final_stages(config) + 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(): @@ -133,18 +162,19 @@ def _run_sequential( on_progress(idx - 1, total, "file") log(f"\n▶ [{idx}/{total}] {f.name}") - track_out = Path(config.output_dir) / f.stem - track_out.mkdir(parents=True, exist_ok=True) + 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 = track_out / "convert.wav" + target = work_dir / "convert.wav" if config.overwrite or not target.exists(): convert_to_wav( f, target, @@ -158,7 +188,7 @@ def _run_sequential( if config.separate.enabled: _check_cancel(stop_event) log(" [2] Отделение вокала...") - target = track_out / "vocals.wav" + target = work_dir / "vocals.wav" with tempfile.TemporaryDirectory() as tmp: raw = separate_vocals( audio_path, Path(tmp), @@ -183,23 +213,31 @@ def _run_sequential( 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, ) - (track_out / "text.txt").write_text(text, encoding="utf-8") + (work_dir / "text.txt").write_text(text, encoding="utf-8") log(f" ✓ text.txt ({len(text)} символов)") - if not config.keep_intermediate: - _cleanup_keeping(track_out, finals) + processed_ok = True except Cancelled: log("⏹ Отменено") break - except Exception as e: # noqa: BLE001 — per-file isolation + 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, track_out if not error else None, error)) + results.append(FileResult(f, text, persist_dir if not error else None, error)) return results @@ -215,20 +253,23 @@ def _run_staged( ) -> list[FileResult]: output_root = Path(config.output_dir) total = len(files) + final_stage = _final_stage(config) - # per-file state - state: dict[Path, dict] = { - f: {"audio_path": f, "track_out": None, "text": None, "error": None} - for f in files - } + # 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) - def track_out_for(f: Path) -> Path: - s = state[f] - if s["track_out"] is None: - p = output_root / f.stem - p.mkdir(parents=True, exist_ok=True) - s["track_out"] = p - return s["track_out"] + 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] @@ -236,94 +277,110 @@ def _run_staged( def cancelled() -> bool: return stop_event is not None and stop_event.is_set() - # ---- 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 = track_out_for(f) / "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, + 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]["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}") + (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}") - # ---- 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 = track_out_for(f) / "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}") + # ---- 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}") - # ---- 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, - ) - (track_out_for(f) / "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}") - - # ---- Cleanup ---- - if not config.keep_intermediate: - finals = _final_stages(config) - for f in files: - s = state[f] - if s["error"] is None and s["track_out"] is not None: - _cleanup_keeping(s["track_out"], finals) + finally: + shutil.rmtree(tmp_root, ignore_errors=True) return [ - FileResult(f, state[f]["text"], state[f]["track_out"], state[f]["error"]) + FileResult( + f, + state[f]["text"], + state[f]["persist_dir"] if state[f]["error"] is None else None, + state[f]["error"], + ) for f in files ] diff --git a/core/transcribe.py b/core/transcribe.py index 0350fb1..1974353 100644 --- a/core/transcribe.py +++ b/core/transcribe.py @@ -26,7 +26,7 @@ def _retry(fn, attempts: int = 3, base_delay: float = 2.0): for i in range(attempts): try: return fn() - except Exception as e: # noqa: BLE001 — HF download / SSL are broad + except Exception as e: # noqa: BLE001 last = e if i < attempts - 1: time.sleep(base_delay * (i + 1)) @@ -37,17 +37,25 @@ class _WhisperXModels: _cache: dict = {} @classmethod - def get(cls, model_name: str, device: str, compute_type: Optional[str]): + def get(cls, model_name: str, device: str, compute_type: Optional[str], + local_files_only: bool = False): if compute_type is None: compute_type = "int8_float16" if device == "cuda" else "int8" - key = (model_name, device, compute_type) + key = (model_name, device, compute_type, local_files_only) if key not in cls._cache: import whisperx + print(f" [ASR] loading WhisperX {model_name} on {device} " + f"({compute_type}, local_files_only={local_files_only})...") def load(): - return whisperx.load_model(model_name, device, compute_type=compute_type) + return whisperx.load_model( + model_name, device, + compute_type=compute_type, + local_files_only=local_files_only, + ) cls._cache[key] = _retry(load) + print(" [ASR] model loaded.") return cls._cache[key] @@ -81,10 +89,15 @@ def transcribe( compute_type: Optional[str] = None, batch_size: int = 8, gigaam_provider: str = "auto", + local_files_only: bool = False, + vad_filter: bool = True, ) -> str: device = _resolve_device(device) if engine == "whisperx": - return _transcribe_whisperx(audio_path, model, language, device, compute_type, batch_size) + return _transcribe_whisperx( + audio_path, model, language, device, compute_type, + batch_size, local_files_only, vad_filter, + ) if engine == "gigaam": return _transcribe_gigaam(audio_path, gigaam_provider) raise ValueError(f"Unknown ASR engine: {engine!r}") @@ -97,11 +110,17 @@ def _transcribe_whisperx( device: str, compute_type: Optional[str], batch_size: int, + local_files_only: bool, + vad_filter: bool, ) -> str: import whisperx - model = _WhisperXModels.get(model_name, device, compute_type) + model = _WhisperXModels.get(model_name, device, compute_type, local_files_only) audio = whisperx.load_audio(str(audio_path)) - result = model.transcribe(audio, batch_size=batch_size, language=language) + result = model.transcribe( + audio, + batch_size=batch_size, + language=language, + ) return " ".join(seg["text"].strip() for seg in result["segments"]).strip() diff --git a/gui/main_window.py b/gui/main_window.py index 3d470a2..a622277 100644 --- a/gui/main_window.py +++ b/gui/main_window.py @@ -267,6 +267,19 @@ class MainWindow(QMainWindow): self.gigaam_provider_cb = QComboBox() self.gigaam_provider_cb.addItems(["auto", "cuda", "cpu"]) form.addRow("GigaAM provider:", self.gigaam_provider_cb) + self.local_files_only_cb = QCheckBox("Не проверять HuggingFace (offline)") + self.local_files_only_cb.setToolTip( + "Если модель уже скачана — включите, чтобы не ждать сетевую проверку.\n" + "При первом запуске (модели нет) оставьте выключенным." + ) + form.addRow("", self.local_files_only_cb) + + self.vad_filter_cb = QCheckBox("Использовать VAD (pyannote)") + self.vad_filter_cb.setChecked(True) + self.vad_filter_cb.setToolTip( + "Отключите, если загрузка модели виснет на 'Performing voice activity detection'." + ) + form.addRow("", self.vad_filter_cb) return w # ---------------- visual sync ---------------- @@ -330,6 +343,8 @@ class MainWindow(QMainWindow): self.compute_type_cb.setCurrentText(cfg.transcribe.compute_type or "auto") self.batch_spin.setValue(cfg.transcribe.batch_size) self.gigaam_provider_cb.setCurrentText(cfg.transcribe.gigaam_provider) + self.local_files_only_cb.setChecked(cfg.transcribe.local_files_only) + self.vad_filter_cb.setChecked(cfg.transcribe.vad_filter) def _collect_config(self) -> PipelineConfig: cfg = PipelineConfig() @@ -366,6 +381,8 @@ class MainWindow(QMainWindow): batch_size=self.batch_spin.value(), gigaam_provider=self.gigaam_provider_cb.currentText(), ) + cfg.transcribe.local_files_only = self.local_files_only_cb.isChecked() + cfg.transcribe.vad_filter = self.vad_filter_cb.isChecked() return cfg # ---------------- handlers ----------------