"""Orchestrate identify -> convert -> separate -> transcribe -> lyrics-match.""" from __future__ import annotations import difflib import json import shutil import sys import tempfile import threading from concurrent.futures import ThreadPoolExecutor, as_completed import re from dataclasses import dataclass, field from pathlib import Path from typing import Callable, Optional from .config import PipelineConfig from .convert import convert_to_wav from .lyrics import LyricsMatch, search_by_identity, search_by_text from .metadata import TrackMetadata, read_metadata from .recognize import ShazamResult, recognize from .separate import separate_vocals, separate_vocals_batch from .transcribe import ( CudaOOMError, preload_transcription_model, transcribe, ) from .lrc import is_lrc, read_embedded_lyrics AUDIO_EXTENSIONS = { ".mp3", ".wav", ".flac", ".ogg", ".m4a", ".aac", ".opus", ".wma", } LogFn = Callable[[str], None] ProgressFn = Callable[[str, int, int], None] # stage, current, total_for_stage OverallFn = Callable[[int, int], None] # completed_files, total_files @dataclass class FileResult: input_path: Path text: Optional[str] output_dir: Optional[Path] error: Optional[str] = None source: Optional[str] = None metadata: Optional[dict] = field(default=None) class Cancelled(Exception): """Raised when the stop event is set.""" # ---------------------------------------------------------------- helpers def _free_memory() -> None: """Free unreferenced RAM and VRAM. Keeps model caches alive. Never import torch here: if the current run does not use ASR, importing it would initialize a CUDA context and steal VRAM from separation. If torch is already loaded by the ASR stage, we can safely empty its cache. """ import gc gc.collect() torch = sys.modules.get("torch") if torch is None: return try: if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception: # noqa: BLE001 — never fail cleanup on a broken CUDA context pass 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 _common_parent(paths: list[Path]) -> Path: if not paths: raise ValueError("empty paths") parts = [p.resolve().parts for p in paths] common = [] for chunk in zip(*parts): if all(c == chunk[0] for c in chunk): common.append(chunk[0]) else: break if not common: return paths[0].parent return Path(*common) def _check_cancel(stop_event: Optional[threading.Event]) -> None: if stop_event is not None and stop_event.is_set(): raise Cancelled() def _text_head(text: str, n_words: int = 8) -> str: return " ".join(text.split()[:n_words]) def _build_metadata_dict( metadata: TrackMetadata, shazam: Optional[ShazamResult], source: Optional[str], input_path: Optional[Path] = None, ) -> Optional[dict]: d: dict = {} if input_path is not None: try: d["path"] = str(input_path.resolve()) except (OSError, RuntimeError): d["path"] = str(input_path) if metadata.has_identity: d["tags"] = { "title": metadata.title, "artist": metadata.artist, "album": metadata.album, "duration": metadata.duration, } if shazam: d["shazam"] = { "title": shazam.title, "artist": shazam.artist, "album": shazam.album, "year": shazam.year, } if source: d["source"] = source return d if d else None def _artifact_belongs_to(track_out: Path, input_path: Path) -> bool: """Check that an output folder was created for this exact input file. ``metadata.json`` now stores the source path. If an output folder with the same stem belongs to a different file (e.g. two albums both contain ``01 - Intro.mp3``), its artifacts must not be treated as "ready" for this track. Legacy folders without a ``path`` key are still accepted. """ meta_file = track_out / "metadata.json" if not meta_file.is_file(): return True try: data = json.loads(meta_file.read_text(encoding="utf-8")) except (OSError, ValueError): return True stored = data.get("path") if isinstance(data, dict) else None if not stored: return True try: return Path(str(stored)).resolve() == input_path.resolve() except (OSError, RuntimeError): return Path(str(stored)) == input_path def _write_metadata_json(path: Path, data: Optional[dict]) -> None: if not data: return path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8") def _merge_metadata_file(existing: Path, data: Optional[dict]) -> Optional[dict]: """Keep tags/shazam from an earlier run when the new data has less info. Reprocessing a track for one stage only (e.g. separation) must not replace a rich metadata.json with a path-only dict. """ if not data or not existing.is_file(): return data try: old = json.loads(existing.read_text(encoding="utf-8")) except (OSError, ValueError): return data if not isinstance(old, dict): return data merged = dict(old) merged.update(data) return merged def _output_dir_for(f: Path, input_root: Path, output_root: Path) -> Path: """Path-based track folder: output//.""" try: rel = f.relative_to(input_root) except ValueError: rel = Path(f.name) return output_root / rel.with_suffix("") def _existing_artifacts( input_root: Path, output_root: Path, input_path: Path, ) -> tuple[bool, bool, bool]: """Return (has_text, has_vocals, has_convert) for the matching output folder. Artifacts are ignored when ``metadata.json`` says the folder belongs to a different input file with the same stem. """ track_out = _output_dir_for(input_path, input_root, output_root) if not _artifact_belongs_to(track_out, input_path): return False, False, False return ( (track_out / "text.txt").is_file(), (track_out / "vocals.wav").is_file(), (track_out / "convert.wav").is_file(), ) def _final_artifact_present(track_out: Path, config: PipelineConfig) -> bool: """True if the artifact of the last enabled stage already exists. ``skip_existing`` must not treat an old ``text.txt`` as "done" when the current run only needs separation: the final stage here is ``vocals.wav``, and without it the track still has work to do. """ if config.transcribe.enabled: return (track_out / "text.txt").is_file() if config.separate.enabled: return (track_out / "vocals.wav").is_file() if config.convert.enabled: return (track_out / "convert.wav").is_file() return False def _is_track_finished( track_out: Path, config: PipelineConfig, belongs: bool = True, ) -> bool: """True when ``skip_existing`` may skip this output folder entirely.""" if not belongs: return False if _final_artifact_present(track_out, config): return True # Identification-only runs still finish with text.txt (no other stage # produces an artifact of its own). if (not config.convert.enabled and not config.separate.enabled and not config.transcribe.enabled and (config.shazam.enabled or config.lyrics.enabled)): return (track_out / "text.txt").is_file() return False def _needs_work( config: PipelineConfig, input_root: Path, output_root: Path, f: Path, ) -> bool: """True if at least one enabled stage still has to do something for ``f``. This mirrors the resume rules used by the staged/sequential pipelines: * final artifact present + skip_existing -> resolved, nothing to do * vocals.wav + resume_partial -> convert/separate are skipped * convert.wav + resume_partial -> convert is skipped Files with no pending work are filtered out before batching, so batches stay full instead of being filled by already-finished tracks. """ track_out = _output_dir_for(f, input_root, output_root) belongs = _artifact_belongs_to(track_out, f) has_text, has_vocals, has_convert = _existing_artifacts( input_root, output_root, f, ) if config.skip_existing and _is_track_finished(track_out, config, belongs): return False vocals_ready = config.resume_partial and has_vocals convert_ready = config.resume_partial and (has_vocals or has_convert) if config.shazam.enabled or config.lyrics.enabled: return True # identification still has to run if config.convert.enabled and not convert_ready: return True if config.separate.enabled and not vocals_ready: return True if config.transcribe.enabled: return True # text is missing or skip_existing is disabled return False def _save_unrecognized( config: PipelineConfig, f: Path, text: Optional[str], source: Optional[str], metadata: Optional[TrackMetadata] = None, lrc_path: Optional[Path] = None, ) -> None: """Append a record to the JSONL file for texts NOT confirmed by LRCLIB.""" if not config.lyrics.save_unrecognized or not text: return if source and "+lrclib" in source: return if config.lyrics.unrecognized_path: out_path = Path(config.lyrics.unrecognized_path) else: out_path = Path(config.output_dir) / "unrecognized_lyrics.jsonl" out_path.parent.mkdir(parents=True, exist_ok=True) record = { "file": str(f), "trackName": (metadata.title if metadata else None) or f.stem, "artistName": metadata.artist if metadata else None, "albumName": metadata.album if metadata else None, "duration": metadata.duration if metadata else None, "plainLyrics": text if not is_lrc(text) else None, "syncedLyrics": text if is_lrc(text) else None, "source": source, } try: with out_path.open("a", encoding="utf-8") as fh: fh.write(json.dumps(record, ensure_ascii=False) + "\n") except OSError: pass @dataclass class IdentifyResult: metadata: TrackMetadata shazam: Optional[ShazamResult] = None text: Optional[str] = None source: Optional[str] = None needs_asr: bool = True lrclib_attempted: bool = False # author/title LRCLIB search was attempted def _identify_file( input_path: Path, work_dir: Path, config: PipelineConfig, log: LogFn, ) -> IdentifyResult: """Run cheap identification steps. May populate .text and .source, and set .needs_asr=False if the track is fully resolved. Order: 0. embedded lyrics in file tags (USLT/LYRICS/©lyr) — cheapest, no network 1. Shazam fingerprint 2. LRCLIB identity search — exactly once: Shazam result if available, otherwise artist/title from file tags """ result = IdentifyResult(metadata=read_metadata(input_path)) # --- 0) embedded lyrics в самом файле --- if config.lyrics.enabled and config.lyrics.use_embedded_lyrics: embedded = read_embedded_lyrics(input_path) if embedded and embedded.strip(): if is_lrc(embedded): result.text = embedded result.source = "embedded+lrc" log(" ✓ встроенный LRC найден") else: result.text = embedded result.source = "embedded" log(" ✓ встроенный текст найден") result.needs_asr = False return result # --- 1) Shazam --- if config.shazam.enabled: log(" Shazam...") shazam_wav = work_dir / "shazam_16k.wav" try: work_dir.mkdir(parents=True, exist_ok=True) convert_to_wav( input_path, shazam_wav, sample_rate=16000, channels=1, ffmpeg_path=config.convert.ffmpeg_path, ) result.shazam = recognize(shazam_wav) except Exception as e: # noqa: BLE001 log(f" ⚠ Shazam error: {e}") result.shazam = None finally: # 16 kHz WAV для Shazam больше не нужен — удаляем сразу, # чтобы не копить десятки гигабайт в /tmp на длинных батчах. try: shazam_wav.unlink(missing_ok=True) except OSError: pass if result.shazam: log(f" ✓ {result.shazam.artist} — {result.shazam.title}") else: log(" ⚠ не распознан") # --- 2) LRCLIB: ровно один поиск по автору/названию --- # Приоритет — результат Shazam. Если Shazam ничего не нашёл (или его # результат не используется), берём artist/title из тегов файла. if config.lyrics.enabled: if result.shazam and config.lyrics.use_shazam_result: log(" LRCLIB по Shazam...") result.lrclib_attempted = True match = search_by_identity( title=result.shazam.title, artist=result.shazam.artist, album=result.shazam.album, duration=result.metadata.duration, user_agent=config.lyrics.user_agent, timeout=config.lyrics.timeout, ) if match: result.text = match.lyrics result.source = "shazam+lrclib" result.needs_asr = False log(f" ✓ {match.artist} — {match.title}") return result log(" ⚠ LRCLIB: не найдено") elif config.lyrics.use_tags_first and result.metadata.has_identity: log(" LRCLIB по тегам...") result.lrclib_attempted = True match = search_by_identity( title=result.metadata.title, artist=result.metadata.artist, album=result.metadata.album, duration=result.metadata.duration, user_agent=config.lyrics.user_agent, timeout=config.lyrics.timeout, ) if match: result.text = match.lyrics result.source = "tag+lrclib" result.needs_asr = False log(f" ✓ {match.artist} — {match.title}") return result log(" ⚠ LRCLIB: не найдено") # --- 3) Identified, but no lyrics --- if result.shazam: result.source = "shazam" if not config.lyrics.transcribe_if_identified_no_lyrics: result.needs_asr = False if not result.lrclib_attempted: log(" ⚠ текста нет") return result def _try_post_asr_match(text: str, config: PipelineConfig, log: LogFn) -> Optional[str]: if not text: return None ignore_phrases = ["субтитры", "редактор", "перевод", "сделал", "озвучка"] lines = [] for line in text.splitlines(): line = line.strip() if not line or any(p in line.lower() for p in ignore_phrases): continue lines.append(line) if not lines: return None query = "" for line in lines: words = line.split() if len(words) >= 3: clean_words = [re.sub(r'[^\w\s]', '', w) for w in words] query = " ".join(clean_words[:6]) break if not query: return None log(f" LRCLIB по тексту: {query!r}...") try: candidates = search_by_text( query, limit=5, user_agent=config.lyrics.user_agent, timeout=config.lyrics.timeout, ) except Exception as e: # noqa: BLE001 log(f" ⚠ LRCLIB error: {e}") return None if not candidates: log(" ⚠ нет кандидатов") return None best: Optional[LyricsMatch] = None best_ratio = 0.0 for c in candidates: ratio = difflib.SequenceMatcher(None, text.lower(), c.lyrics.lower()).ratio() if ratio > best_ratio: best_ratio = ratio best = c if best is not None and best_ratio >= config.lyrics.fuzzy_threshold: log(f" ✓ совпадение {best_ratio:.2f}: {best.artist} — {best.title}") return best.lyrics log(f" ⚠ лучшее совпадение {best_ratio:.2f} < {config.lyrics.fuzzy_threshold}") return None def _persist_artifacts( work_dir: Path, persist_dir: Path, config: PipelineConfig, has_text: bool, has_metadata: bool, log: LogFn, ) -> None: if work_dir.resolve() == persist_dir.resolve(): return # nothing to move — same directory persist_dir.mkdir(parents=True, exist_ok=True) names: list[str] = [] if has_metadata and (work_dir / "metadata.json").exists(): names.append("metadata.json") if has_text and (work_dir / "text.txt").exists(): names.append("text.txt") if has_text and (work_dir / "text.lrc").exists(): names.append("text.lrc") if config.keep_intermediate: for n in ("convert.wav", "vocals.wav"): if (work_dir / n).exists(): names.append(n) elif not has_text: if config.separate.enabled and (work_dir / "vocals.wav").exists(): names.append("vocals.wav") elif config.convert.enabled and (work_dir / "convert.wav").exists(): names.append("convert.wav") if config.separate.save_vocals_copy and (work_dir / "vocals.wav").exists(): if "vocals.wav" not in names: names.append("vocals.wav") for name in names: src = work_dir / name dst = persist_dir / name if dst.exists(): dst.unlink() shutil.move(str(src), str(dst)) log(f" [persist] {name}") # ---------------------------------------------------------------- entry def _run_identification_parallel( config: PipelineConfig, files: list[Path], log: LogFn, on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event], work_root: Path, ) -> dict[Path, Optional[IdentifyResult]]: if not (config.shazam.enabled or config.lyrics.enabled): return {} concurrency = max(1, int(config.shazam.identify_concurrency)) total = len(files) if concurrency == 1 or total == 1: results: dict[Path, Optional[IdentifyResult]] = {} for i, f in enumerate(files, 1): if stop_event is not None and stop_event.is_set(): break if on_progress: on_progress("identify", i, total) try: results[f] = _identify_file(f, work_root / f.stem, config, log) except Exception as e: # noqa: BLE001 log(f" ✗ {f.name} (identify): {e}") results[f] = None return results log(f" Идентификация: {total} файлов, конкурентность {concurrency}") def worker(f: Path) -> tuple[Path, Optional[IdentifyResult]]: if stop_event is not None and stop_event.is_set(): return f, None def prefixed(msg: str) -> None: log(f" [{f.name}] {msg.strip()}") try: return f, _identify_file(f, work_root / f.stem, config, prefixed) except Exception as e: # noqa: BLE001 log(f" ✗ {f.name} (identify): {e}") return f, None results: dict[Path, Optional[IdentifyResult]] = {} with ThreadPoolExecutor(max_workers=concurrency) as pool: futures = {pool.submit(worker, f): f for f in files} done = 0 for fut in as_completed(futures): f = futures[fut] try: _, ident = fut.result() results[f] = ident except Exception as e: # noqa: BLE001 log(f" ✗ {f.name} (identify): {e}") results[f] = None done += 1 if on_progress: on_progress("identify", done, total) return results def _run_convert_parallel( config: PipelineConfig, files: list[Path], state: dict, log: LogFn, on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event], temp_registry: Optional[list[Path]] = None, ) -> None: concurrency = max(1, int(config.convert.convert_concurrency)) total = len(files) def worker(f: Path) -> tuple[Path, Optional[Path], Optional[str]]: if stop_event is not None and stop_event.is_set(): return f, None, None try: temp_dir = Path(tempfile.mkdtemp(prefix=f"at_{f.stem}_")) if temp_registry is not None: temp_registry.append(temp_dir) target = temp_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, ) return f, target, None except Exception as e: # noqa: BLE001 return f, None, f"{type(e).__name__}: {e}" if concurrency == 1 or total == 1: for i, f in enumerate(files, 1): if stop_event is not None and stop_event.is_set(): break if on_progress: on_progress("convert", i, total) f2, target, err = worker(f) if err: state[f2]["error"] = err log(f" ✗ {f2.name}: {err}") else: state[f2]["audio_path"] = target log(f" ✓ {f2.name}") return log(f" Конвертация: {total} файлов, конкурентность {concurrency}") with ThreadPoolExecutor(max_workers=concurrency) as pool: futures = {pool.submit(worker, f): f for f in files} done = 0 for fut in as_completed(futures): f2, target, err = fut.result() if err: state[f2]["error"] = err log(f" ✗ {f2.name}: {err}") else: state[f2]["audio_path"] = target log(f" ✓ {f2.name}") done += 1 if on_progress: on_progress("convert", done, total) def run_pipeline( config: PipelineConfig, log: LogFn = print, on_progress: Optional[ProgressFn] = None, stop_event: Optional[threading.Event] = None, on_overall_progress: Optional[OverallFn] = 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() total_all = len(files) log(f"Найдено треков: {total_all} | режим: {mode} | " f"keep_intermediate={config.keep_intermediate}") if config.separate.save_vocals_copy: log(" ✓ Вокал будет сохраняться всегда (save_vocals_copy=True)") if config.shazam.enabled: log(" Shazam: включён") if config.lyrics.enabled: log(f" LRCLIB: включён (порог fuzzy={config.lyrics.fuzzy_threshold:.2f})") if config.skip_existing: log(" Resume: готовые файлы (результат последнего этапа) будут пропущены") if config.resume_partial: log(" Resume: существующие vocals.wav / convert.wav будут переиспользованы") # Pre-scan before chunking: remove files that are already complete for the # current set of enabled stages. Without this, skipped tracks consumed # batch slots and a staged run produced many half-empty batches. pending = [f for f in files if _needs_work(config, input_dir, output_dir, f)] already_done = total_all - len(pending) if already_done: log(f" ✓ Уже готово: {already_done} — в батчи не попадут") log(f" → К обработке: {len(pending)}") files = pending if not files: log(" Все файлы уже готовы — работа не требуется.") if on_overall_progress: on_overall_progress(total_all, total_all) return [] if on_overall_progress: on_overall_progress(already_done, total_all) if mode == "staged": return _run_staged( config, files, log, on_progress, stop_event, on_overall_progress, overall_base=already_done, overall_total=total_all, input_root=input_dir, ) return _run_sequential( config, files, log, on_progress, stop_event, on_overall_progress, overall_base=already_done, overall_total=total_all, input_root=input_dir, ) def _run_sequential( config: PipelineConfig, files: list[Path], log: LogFn, on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event], on_overall_progress: Optional[OverallFn] = None, overall_base: int = 0, overall_total: Optional[int] = None, input_root: Optional[Path] = None, ) -> list[FileResult]: total = len(files) if overall_total is None: overall_total = overall_base + total if input_root is None: input_root = Path(config.input_dir) results: list[FileResult] = [] output_root = Path(config.output_dir) log("Раскладка: output///") ident_tmp = Path(tempfile.mkdtemp(prefix="at_ident_")) ident_files = files if config.skip_existing: ident_files = [ f for f in files if not (_output_dir_for(f, input_root, output_root) / "text.txt").is_file() ] try: ident_map = _run_identification_parallel( config, ident_files, log, on_progress, stop_event, ident_tmp, ) finally: shutil.rmtree(ident_tmp, ignore_errors=True) 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("file", idx - 1, total) log(f"\n▶ [{idx}/{total}] {f.name}") track_out = _output_dir_for(f, input_root, output_root) track_out.mkdir(parents=True, exist_ok=True) belongs = _artifact_belongs_to(track_out, f) if not belongs and any( (track_out / name).exists() for name in ("text.txt", "vocals.wav", "convert.wav") ): log(f" ⚠ коллизия раскладки: {track_out} занят другим треком — обрабатываю заново") if config.skip_existing and _is_track_finished(track_out, config, belongs): existing_text = track_out / "text.txt" existing: Optional[str] = None if existing_text.is_file(): try: existing = existing_text.read_text(encoding="utf-8") except OSError: existing = None log(" ⊘ результат последнего этапа уже есть — пропуск") if on_progress: on_progress("file", idx, total) if on_overall_progress: on_overall_progress( min(overall_base + idx, overall_total), overall_total, ) results.append(FileResult(f, existing, track_out, None, "existing")) continue tmp_root = Path(tempfile.mkdtemp(prefix=f"at_{f.stem}_")) text: Optional[str] = None source: Optional[str] = None error: Optional[str] = None metadata_dict: Optional[dict] = None appended = False was_cancelled = False try: ident = ident_map.get(f) or IdentifyResult(metadata=TrackMetadata()) if not ident.needs_asr: text = ident.text source = ident.source metadata_dict = _build_metadata_dict( ident.metadata, ident.shazam, source, input_path=f, ) if text: (tmp_root / "text.txt").write_text(text, encoding="utf-8") if metadata_dict: _write_metadata_json(tmp_root / "metadata.json", metadata_dict) _persist_artifacts( tmp_root, track_out, config, has_text=text is not None, has_metadata=bool(metadata_dict), log=log, ) _save_unrecognized( config, f, text, source, metadata=ident.metadata, ) log(f" ✓ {f.name} [{source}]") results.append(FileResult(f, text, track_out, None, source, metadata_dict)) appended = True continue audio_path: Path = f skip_convert = False skip_separate = False if config.resume_partial and belongs: vocals = track_out / "vocals.wav" convert = track_out / "convert.wav" if vocals.is_file(): audio_path = vocals skip_convert = True skip_separate = True log(" ↻ используем готовый vocals.wav") elif convert.is_file(): audio_path = convert skip_convert = True log(" ↻ используем готовый convert.wav") if config.convert.enabled and not skip_convert: _check_cancel(stop_event) log(" [1] Конвертация в WAV...") target = tmp_root / "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 and not skip_separate: _check_cancel(stop_event) log(" [2] Отделение вокала...") target = tmp_root / "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, ) text = (text or "").strip() source = config.transcribe.engine if not text: text = None source = "empty_result" error = "EmptyResult: ASR returned no text" log(f" ⚠ {f.name}: текст не извлечён (пустой результат ASR)") else: log(f" ✓ распознано {len(text)} символов") if (config.lyrics.enabled and config.lyrics.try_after_asr): matched = _try_post_asr_match(text, config, log) if matched: text = matched source = f"{source}+lrclib" (tmp_root / "text.txt").write_text(text, encoding="utf-8") metadata_dict = _build_metadata_dict( ident.metadata, ident.shazam, source, input_path=f, ) if belongs: metadata_dict = _merge_metadata_file( track_out / "metadata.json", metadata_dict, ) if metadata_dict: _write_metadata_json(tmp_root / "metadata.json", metadata_dict) _persist_artifacts( tmp_root, track_out, config, has_text=text is not None, has_metadata=bool(metadata_dict), log=log, ) _save_unrecognized( config, f, text, source, metadata=ident.metadata, ) except Cancelled: log("⏹ Отменено") was_cancelled = True break except CudaOOMError as e: error = f"CudaOOMError: {e}" source = "oom_skipped" log(f" ⊘ {f.name}: OOM после ретраев — пропущен, идём дальше") except Exception as e: # noqa: BLE001 error = f"{type(e).__name__}: {e}" log(f" ✗ Ошибка: {error}") finally: # Even on error — preserve vocals if the user asked for it. if error is not None and config.separate.save_vocals_copy: vocal_src = tmp_root / "vocals.wav" if vocal_src.exists(): try: dst = track_out / "vocals.wav" if dst.exists(): dst.unlink() shutil.move(str(vocal_src), str(dst)) log(f" ✓ {f.name} (vocals сохранён, несмотря на ошибку)") except OSError: pass shutil.rmtree(tmp_root, ignore_errors=True) if was_cancelled: # The file was not processed to completion — do not count it # in the overall bar and do not report a fake "success" row. pass else: if not appended: results.append(FileResult( f, text, track_out if not error else None, error, source, metadata_dict, )) if on_progress: on_progress("file", idx, total) if on_overall_progress: on_overall_progress( min(overall_base + idx, overall_total), overall_total, ) _log_summary(results, log) return results def _run_staged( config: PipelineConfig, files: list[Path], log: LogFn, on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event], on_overall_progress: Optional[OverallFn] = None, overall_base: int = 0, overall_total: Optional[int] = None, input_root: Optional[Path] = None, ) -> list[FileResult]: """Chunked staged mode: process `chunk_size` files, free memory, repeat.""" output_root = Path(config.output_dir) if input_root is None: input_root = Path(config.input_dir) total = len(files) if overall_total is None: overall_total = overall_base + total chunk_size = max(1, int(config.chunk_size or total)) log("Раскладка: output///") if chunk_size >= total: log(f"Один батч: {total} файлов (chunk_size={chunk_size})") results = _run_staged_batch( config, files, output_root, log, on_progress, stop_event, overall_offset=overall_base, overall_total=overall_total, on_overall_progress=on_overall_progress, overall_counter=[overall_base], input_root=input_root, ) _free_memory() _log_summary(results, log) return results n_batches = (total + chunk_size - 1) // chunk_size log(f"Батчей: {n_batches} × {chunk_size} файлов (всего {total})") all_results: list[FileResult] = [] overall_counter = [overall_base] for b in range(n_batches): if stop_event is not None and stop_event.is_set(): log("⏹ Отменено") break start = b * chunk_size end = min(start + chunk_size, total) chunk = files[start:end] log(f"\n╔═══════════════════════════════════════════════════════╗") log(f"║ Батч {b+1}/{n_batches}: файлы {start+1}–{end} из {total}") log(f"╚═══════════════════════════════════════════════════════╝") try: batch_results = _run_staged_batch( config, chunk, output_root, log, on_progress, stop_event, overall_offset=overall_base + start, overall_total=overall_total, on_overall_progress=on_overall_progress, overall_counter=overall_counter, input_root=input_root, ) all_results.extend(batch_results) except Exception as e: # noqa: BLE001 log(f" ✗ Батч {b+1} упал: {type(e).__name__}: {e}") _free_memory() log(f" [batch] память очищена") _log_summary(all_results, log) return all_results def _run_staged_batch( config: PipelineConfig, batch_files: list[Path], output_root: Path, log: LogFn, on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event], overall_offset: int = 0, overall_total: Optional[int] = None, on_overall_progress: Optional[OverallFn] = None, overall_counter: Optional[list[int]] = None, input_root: Optional[Path] = None, ) -> list[FileResult]: """One chunk of the staged pipeline. All temps are tracked and removed in the `finally` block at the end of this call.""" if input_root is None: input_root = Path(config.input_dir) if overall_total is None: overall_total = len(batch_files) if overall_counter is None: overall_counter = [overall_offset] def mark_overall(complete: bool) -> None: if not complete: return overall_counter[0] += 1 if on_overall_progress: on_overall_progress( min(overall_counter[0], overall_total), overall_total, ) def is_complete(s: dict) -> bool: return bool( s["resolved"] or s["error"] is not None or s["text"] is not None or not config.transcribe.enabled ) def progress(stage: str, done: int, total: int) -> None: """Report progress for the current stage only. The GUI bar is reset for every stage, so the counters never have to be mixed across stages (identify sees all files, transcribe only the ones that actually need ASR, etc.). """ if on_progress: on_progress(stage, done, total) temp_registry: list[Path] = [] def _mkdtemp(prefix: str) -> Path: p = Path(tempfile.mkdtemp(prefix=prefix)) temp_registry.append(p) return p def active() -> list[Path]: return [f for f in batch_files if state[f]["error"] is None and not state[f]["resolved"]] def cancelled() -> bool: return stop_event is not None and stop_event.is_set() # Seed the state from artifacts already present in output_root/. # This is what makes "Resume" work: a file that already has text.txt is # resolved, a file with vocals.wav skips convert+separate, and a file with # only convert.wav skips just the conversion step. state: dict[Path, dict] = {} skipped_existing = 0 resumed_partial = 0 mismatched_folders = 0 for f in batch_files: track_out = _output_dir_for(f, input_root, output_root) audio_path: Path = f skip_convert = False skip_separate = False text: Optional[str] = None source: Optional[str] = None resolved = False belongs = _artifact_belongs_to(track_out, f) if not belongs: has_artifacts = any( (track_out / name).exists() for name in ("text.txt", "vocals.wav", "convert.wav") ) if has_artifacts: mismatched_folders += 1 log(f" ⚠ коллизия раскладки: {track_out} занят другим треком — обрабатываю заново") text_file = track_out / "text.txt" if config.skip_existing and _is_track_finished(track_out, config, belongs): resolved = True skipped_existing += 1 source = "existing" if text_file.is_file(): try: text = text_file.read_text(encoding="utf-8") except OSError: text = None elif config.resume_partial and belongs: vocals = track_out / "vocals.wav" convert = track_out / "convert.wav" if vocals.is_file(): audio_path = vocals skip_convert = True skip_separate = True resumed_partial += 1 elif convert.is_file(): audio_path = convert skip_convert = True resumed_partial += 1 state[f] = { "audio_path": audio_path, "ident": None, "text": text, "source": source, "metadata_dict": None, "error": None, "resolved": resolved, "skip_convert": skip_convert, "skip_separate": skip_separate, } if skipped_existing: log(f" ⊘ Пропущено готовых (результат последнего этапа): {skipped_existing}") if resumed_partial: log(f" ↻ Подхвачено промежуточных (vocals/convert): {resumed_partial}") if mismatched_folders: log(f" ⚠ Папок с чужим path: {mismatched_folders} (пересчитаны заново)") if on_overall_progress: on_overall_progress(min(overall_counter[0], overall_total), overall_total) try: # --- Stage 0: identification --- if config.shazam.enabled or config.lyrics.enabled: to_identify = [f for f in batch_files if not state[f]["resolved"]] if to_identify: log(f"\n══ Этап 0 · Идентификация ({len(to_identify)} файлов) ══") ident_tmp = _mkdtemp("at_ident_") try: ident_map = _run_identification_parallel( config, to_identify, log, progress, stop_event, ident_tmp, ) finally: shutil.rmtree(ident_tmp, ignore_errors=True) try: temp_registry.remove(ident_tmp) except ValueError: pass for f in to_identify: ident = ident_map.get(f) if ident is None: state[f]["ident"] = IdentifyResult(metadata=TrackMetadata()) log(f" → {f.name} (identification failed)") continue state[f]["ident"] = ident if not ident.needs_asr: state[f]["resolved"] = True state[f]["text"] = ident.text state[f]["source"] = ident.source state[f]["metadata_dict"] = _build_metadata_dict( ident.metadata, ident.shazam, ident.source, input_path=f, ) track_out = _output_dir_for(f, input_root, output_root) track_out.mkdir(parents=True, exist_ok=True) if ident.text: (track_out / "text.txt").write_text(ident.text, encoding="utf-8") if state[f]["metadata_dict"]: _write_metadata_json(track_out / "metadata.json", state[f]["metadata_dict"]) _save_unrecognized( config, f, ident.text, ident.source, metadata=ident.metadata, ) log(f" ✓ {f.name} [{ident.source}]") else: log(f" → {f.name} (нужен полный пайплайн)") # --- Stage 1: convert --- if config.convert.enabled and not cancelled(): todo = [f for f in active() if not state[f]["skip_convert"]] log(f"\n══ Этап 1 · Конвертация ({len(todo)} файлов) ══") progress("convert", 0, len(todo)) _run_convert_parallel( config, todo, state, log, progress, stop_event, temp_registry, ) # --- Stage 2: separate --- if config.separate.enabled and not cancelled(): todo = [f for f in active() if not state[f]["skip_separate"]] log(f"\n══ Этап 2 · Отделение вокала ({len(todo)} файлов) ══") if todo: def separate_single(f: Path) -> None: try: src = state[f]["audio_path"] # Never write next to the original input: with # convert.enabled=False src == f, so every track would # overwrite input_dir/vocals.wav. Only reuse the track # output folder when the source already lives there. track_dir = _output_dir_for(f, input_root, output_root) if src.parent == track_dir: # Already in the persistent track folder (resume). target = src.parent / "vocals.wav" elif src.parent == f.parent: # src is the original input file (convert disabled). # Never write next to it: all tracks share input_dir # and would overwrite input_dir/vocals.wav. target = _mkdtemp("at_sep_") / "vocals.wav" else: # Freshly converted file in its own unique temp dir. target = src.parent / "vocals.wav" with tempfile.TemporaryDirectory() as tmp: raw = separate_vocals( src, 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}") # Resumed convert.wav files already live in output_root/, # while freshly converted files live in unique temp dirs. # Original inputs (convert disabled) must also go through the # one-by-one path: batching them shares input_dir and can # collide on vocals.wav. # unblend's `{parent}` batch layout would mix those roots, so # separate resumed/original files one by one and batch the rest. resumed = [ f for f in todo if state[f]["audio_path"].parent == _output_dir_for(f, input_root, output_root) or state[f]["audio_path"].parent == f.parent ] resumed_set = set(resumed) batchable = [f for f in todo if f not in resumed_set] done = 0 progress("separate", 0, len(todo)) if len(batchable) == 1: separate_single(batchable[0]) done = 1 progress("separate", done, len(todo)) elif batchable: input_wavs = [state[f]["audio_path"] for f in batchable] common_root = _common_parent(input_wavs) log(f" Батч: {len(batchable)} файлов, один вызов unblend") try: mapping = separate_vocals_batch( input_wavs, common_root, model=config.separate.model, isolate_stem=config.separate.isolate_stem, unblend_path=config.separate.unblend_path, custom_kernels=config.separate.custom_kernels, progress_callback=lambda p: progress( "separate", p, len(todo), ), ) for f in batchable: src = state[f]["audio_path"] vocal = mapping.get(src) if vocal is None: state[f]["error"] = "batch produced no output" log(f" ✗ {f.name}: нет результата в батче") continue state[f]["audio_path"] = vocal log(f" ✓ {f.name}") done = len(batchable) progress("separate", done, len(todo)) except Exception as e: # noqa: BLE001 log(f" ✗ Батч упал: {e}") log(" ↻ Fallback: по одному...") for i, f in enumerate(batchable, 1): if cancelled(): break progress("separate", i - 1, len(todo)) separate_single(f) done = i progress("separate", i, len(todo)) for f in resumed: if cancelled(): break done += 1 progress("separate", done, len(todo)) separate_single(f) else: progress("separate", 0, 0) # --- Stage 3: transcribe --- if config.transcribe.enabled and not cancelled(): todo = active() log(f"\n══ Этап 3 · Транскрибация ({len(todo)} файлов) ══") if todo: # Load once, after separation has finished. Every later file # and every later batch reuses this instance. try: preload_transcription_model( engine=config.transcribe.engine, model=config.transcribe.model, device=config.transcribe.device, compute_type=config.transcribe.compute_type, gigaam_provider=config.transcribe.gigaam_provider, local_files_only=config.transcribe.local_files_only, ) except Exception as e: # noqa: BLE001 log(f" ⚠ Не удалось загрузить модель заранее: {e}") progress("transcribe", 0, len(todo)) for i, f in enumerate(todo, 1): if cancelled(): log("⏹ Отменено") break 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, ) text = (text or "").strip() if not text: state[f]["text"] = None state[f]["source"] = "empty_result" state[f]["error"] = "EmptyResult: ASR returned no text" log(f" ⚠ {f.name}: текст не извлечён (пустой результат ASR)") else: state[f]["text"] = text state[f]["source"] = config.transcribe.engine log(f" ✓ {f.name} ({len(text)} символов)") except CudaOOMError as e: state[f]["error"] = f"CudaOOMError: {e}" state[f]["source"] = "oom_skipped" log(f" ⊘ {f.name}: OOM — пропущен") except Exception as e: # noqa: BLE001 state[f]["error"] = f"{type(e).__name__}: {e}" log(f" ✗ {f.name}: {e}") finally: progress("transcribe", i, len(todo)) # --- Stage 4: finalize --- # NOTE: runs even if cancelled — we don't want to lose texts that # were already transcribed before the user hit "Stop". Only the # network call to LRCLIB (fuzzy match) is skipped when cancelled. log("\n══ Финализация ══") for f in batch_files: s = state[f] # Errored file: no text, but if separate succeeded we may still # have a vocals.wav worth keeping (user asked for save_vocals_copy). if s["error"] is not None: if config.separate.save_vocals_copy: src_audio = s["audio_path"] if isinstance(src_audio, Path): vocal = src_audio.parent / "vocals.wav" if vocal.exists(): track_out = _output_dir_for(f, input_root, output_root) track_out.mkdir(parents=True, exist_ok=True) dst = track_out / "vocals.wav" if dst.exists(): dst.unlink() try: shutil.move(str(vocal), str(dst)) log(f" ✓ {f.name} (vocals сохранён, несмотря на ошибку)") except OSError: pass mark_overall(True) continue if s["resolved"]: mark_overall(True) continue text = s["text"] source = s["source"] ident = s["ident"] # LRCLIB fuzzy-match is the only slow/network step here — # skip it on cancel, but still save the text we already have. # This is a *text* search, not a duplicate author/title search. if (text and config.lyrics.enabled and config.lyrics.try_after_asr and not cancelled()): matched = _try_post_asr_match(text, config, log) if matched: text = matched source = f"{source}+lrclib" s["text"] = text s["source"] = source metadata_dict = _build_metadata_dict( ident.metadata if ident else TrackMetadata(), ident.shazam if ident else None, source, input_path=f, ) s["metadata_dict"] = metadata_dict track_out = _output_dir_for(f, input_root, output_root) track_out.mkdir(parents=True, exist_ok=True) if _artifact_belongs_to(track_out, f): metadata_dict = _merge_metadata_file( track_out / "metadata.json", metadata_dict, ) s["metadata_dict"] = metadata_dict audio_work = s["audio_path"] if isinstance(audio_work, Path) and audio_work != f: work_dir = audio_work.parent if text: (work_dir / "text.txt").write_text(text, encoding="utf-8") if metadata_dict: _write_metadata_json(work_dir / "metadata.json", metadata_dict) _persist_artifacts( work_dir, track_out, config, has_text=text is not None, has_metadata=bool(metadata_dict), log=log, ) else: if text: (track_out / "text.txt").write_text(text, encoding="utf-8") if metadata_dict: _write_metadata_json(track_out / "metadata.json", metadata_dict) _save_unrecognized( config, f, text, source, metadata=(ident.metadata if ident else None), ) if text is not None or s["resolved"]: log(f" ✓ {f.name}") else: log(f" ⊘ {f.name}: пропущен (нет текста)") mark_overall(is_complete(s)) finally: # Clean every temp dir created in this batch, regardless of success/error for p in temp_registry: try: if p.exists(): shutil.rmtree(p, ignore_errors=True) except OSError: pass log(f" [batch] удалено временных папок: {len(temp_registry)}") return [ FileResult( f, state[f]["text"], _output_dir_for(f, input_root, output_root) if state[f]["error"] is None else None, state[f]["error"], state[f]["source"], state[f]["metadata_dict"], ) for f in batch_files ] def _log_summary(results: list[FileResult], log: LogFn) -> None: from collections import Counter counts = Counter(r.source or ("error" if r.error else "no_result") for r in results) log("\n══ Итог ══") for src, n in sorted(counts.items(), key=lambda kv: -kv[1]): log(f" {src}: {n}")