pre-scan finished tracks and load CUDA lazily
- pre-scan all inputs before chunking so ready files stay out of batches\n- start overall progress at the already-finished count\n- count skipped/resumed files correctly in sequential and staged runs\n- keep batches full and log ready/pending counts before starting\n- avoid importing Torch at startup or in cleanup for separation-only runs\n- move TF32 setup into lazy WhisperX model loading\n- detect the GPU through nvidia-smi for the GUI status line
This commit is contained in:
106
core/pipeline.py
106
core/pipeline.py
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|||||||
import difflib
|
import difflib
|
||||||
import json
|
import json
|
||||||
import shutil
|
import shutil
|
||||||
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import threading
|
import threading
|
||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
@@ -48,11 +49,18 @@ class Cancelled(Exception):
|
|||||||
# ---------------------------------------------------------------- helpers
|
# ---------------------------------------------------------------- helpers
|
||||||
|
|
||||||
def _free_memory() -> None:
|
def _free_memory() -> None:
|
||||||
"""Free unreferenced RAM and VRAM. Keeps model caches alive."""
|
"""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
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
torch = sys.modules.get("torch")
|
||||||
|
if torch is None:
|
||||||
|
return
|
||||||
try:
|
try:
|
||||||
import torch
|
|
||||||
if torch.cuda.is_available():
|
if torch.cuda.is_available():
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
except Exception: # noqa: BLE001 — never fail cleanup on a broken CUDA context
|
except Exception: # noqa: BLE001 — never fail cleanup on a broken CUDA context
|
||||||
@@ -150,6 +158,53 @@ def _write_metadata_json(path: Path, data: Optional[dict]) -> None:
|
|||||||
path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
|
path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
def _existing_artifacts(
|
||||||
|
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_root / input_path.stem
|
||||||
|
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 _needs_work(config: PipelineConfig, 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:
|
||||||
|
* text.txt + 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.
|
||||||
|
"""
|
||||||
|
has_text, has_vocals, has_convert = _existing_artifacts(output_root, f)
|
||||||
|
|
||||||
|
if config.skip_existing and has_text:
|
||||||
|
return False # resolved by the same rule the pipelines use
|
||||||
|
|
||||||
|
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 (no resolved text)
|
||||||
|
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(
|
def _save_unrecognized(
|
||||||
config: PipelineConfig,
|
config: PipelineConfig,
|
||||||
f: Path,
|
f: Path,
|
||||||
@@ -532,7 +587,8 @@ def run_pipeline(
|
|||||||
log("Аудиофайлы не найдены.")
|
log("Аудиофайлы не найдены.")
|
||||||
return []
|
return []
|
||||||
mode = (config.mode or "sequential").lower()
|
mode = (config.mode or "sequential").lower()
|
||||||
log(f"Найдено треков: {len(files)} | режим: {mode} | "
|
total_all = len(files)
|
||||||
|
log(f"Найдено треков: {total_all} | режим: {mode} | "
|
||||||
f"keep_intermediate={config.keep_intermediate}")
|
f"keep_intermediate={config.keep_intermediate}")
|
||||||
if config.separate.save_vocals_copy:
|
if config.separate.save_vocals_copy:
|
||||||
log(" ✓ Вокал будет сохраняться всегда (save_vocals_copy=True)")
|
log(" ✓ Вокал будет сохраняться всегда (save_vocals_copy=True)")
|
||||||
@@ -544,14 +600,32 @@ def run_pipeline(
|
|||||||
log(" Resume: готовые файлы (text.txt) будут пропущены")
|
log(" Resume: готовые файлы (text.txt) будут пропущены")
|
||||||
if config.resume_partial:
|
if config.resume_partial:
|
||||||
log(" Resume: существующие vocals.wav / convert.wav будут переиспользованы")
|
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, 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:
|
if on_overall_progress:
|
||||||
on_overall_progress(0, len(files))
|
on_overall_progress(total_all, total_all)
|
||||||
|
return []
|
||||||
|
|
||||||
|
if on_overall_progress:
|
||||||
|
on_overall_progress(already_done, total_all)
|
||||||
if mode == "staged":
|
if mode == "staged":
|
||||||
return _run_staged(
|
return _run_staged(
|
||||||
config, files, log, on_progress, stop_event, on_overall_progress,
|
config, files, log, on_progress, stop_event, on_overall_progress,
|
||||||
|
overall_base=already_done, overall_total=total_all,
|
||||||
)
|
)
|
||||||
return _run_sequential(
|
return _run_sequential(
|
||||||
config, files, log, on_progress, stop_event, on_overall_progress,
|
config, files, log, on_progress, stop_event, on_overall_progress,
|
||||||
|
overall_base=already_done, overall_total=total_all,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -559,8 +633,12 @@ def _run_sequential(
|
|||||||
config: PipelineConfig, files: list[Path], log: LogFn,
|
config: PipelineConfig, files: list[Path], log: LogFn,
|
||||||
on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event],
|
on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event],
|
||||||
on_overall_progress: Optional[OverallFn] = None,
|
on_overall_progress: Optional[OverallFn] = None,
|
||||||
|
overall_base: int = 0,
|
||||||
|
overall_total: Optional[int] = None,
|
||||||
) -> list[FileResult]:
|
) -> list[FileResult]:
|
||||||
total = len(files)
|
total = len(files)
|
||||||
|
if overall_total is None:
|
||||||
|
overall_total = total
|
||||||
results: list[FileResult] = []
|
results: list[FileResult] = []
|
||||||
output_root = Path(config.output_dir)
|
output_root = Path(config.output_dir)
|
||||||
ident_tmp = Path(tempfile.mkdtemp(prefix="at_ident_"))
|
ident_tmp = Path(tempfile.mkdtemp(prefix="at_ident_"))
|
||||||
@@ -600,7 +678,9 @@ def _run_sequential(
|
|||||||
if on_progress:
|
if on_progress:
|
||||||
on_progress("file", idx, total)
|
on_progress("file", idx, total)
|
||||||
if on_overall_progress:
|
if on_overall_progress:
|
||||||
on_overall_progress(idx, total)
|
on_overall_progress(
|
||||||
|
min(overall_base + idx, overall_total), overall_total,
|
||||||
|
)
|
||||||
results.append(FileResult(f, existing, track_out, None, "existing"))
|
results.append(FileResult(f, existing, track_out, None, "existing"))
|
||||||
continue
|
continue
|
||||||
|
|
||||||
@@ -762,7 +842,9 @@ def _run_sequential(
|
|||||||
if on_progress:
|
if on_progress:
|
||||||
on_progress("file", idx, total)
|
on_progress("file", idx, total)
|
||||||
if on_overall_progress:
|
if on_overall_progress:
|
||||||
on_overall_progress(idx, total)
|
on_overall_progress(
|
||||||
|
min(overall_base + idx, overall_total), overall_total,
|
||||||
|
)
|
||||||
|
|
||||||
_log_summary(results, log)
|
_log_summary(results, log)
|
||||||
return results
|
return results
|
||||||
@@ -772,19 +854,23 @@ def _run_staged(
|
|||||||
config: PipelineConfig, files: list[Path], log: LogFn,
|
config: PipelineConfig, files: list[Path], log: LogFn,
|
||||||
on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event],
|
on_progress: Optional[ProgressFn], stop_event: Optional[threading.Event],
|
||||||
on_overall_progress: Optional[OverallFn] = None,
|
on_overall_progress: Optional[OverallFn] = None,
|
||||||
|
overall_base: int = 0,
|
||||||
|
overall_total: Optional[int] = None,
|
||||||
) -> list[FileResult]:
|
) -> list[FileResult]:
|
||||||
"""Chunked staged mode: process `chunk_size` files, free memory, repeat."""
|
"""Chunked staged mode: process `chunk_size` files, free memory, repeat."""
|
||||||
output_root = Path(config.output_dir)
|
output_root = Path(config.output_dir)
|
||||||
total = len(files)
|
total = len(files)
|
||||||
|
if overall_total is None:
|
||||||
|
overall_total = overall_base + total
|
||||||
chunk_size = max(1, int(config.chunk_size or total))
|
chunk_size = max(1, int(config.chunk_size or total))
|
||||||
|
|
||||||
if chunk_size >= total:
|
if chunk_size >= total:
|
||||||
log(f"Один батч: {total} файлов (chunk_size={chunk_size})")
|
log(f"Один батч: {total} файлов (chunk_size={chunk_size})")
|
||||||
results = _run_staged_batch(
|
results = _run_staged_batch(
|
||||||
config, files, output_root, log, on_progress, stop_event,
|
config, files, output_root, log, on_progress, stop_event,
|
||||||
overall_offset=0, overall_total=total,
|
overall_offset=overall_base, overall_total=overall_total,
|
||||||
on_overall_progress=on_overall_progress,
|
on_overall_progress=on_overall_progress,
|
||||||
overall_counter=[0],
|
overall_counter=[overall_base],
|
||||||
)
|
)
|
||||||
_free_memory()
|
_free_memory()
|
||||||
_log_summary(results, log)
|
_log_summary(results, log)
|
||||||
@@ -794,7 +880,7 @@ def _run_staged(
|
|||||||
log(f"Батчей: {n_batches} × {chunk_size} файлов (всего {total})")
|
log(f"Батчей: {n_batches} × {chunk_size} файлов (всего {total})")
|
||||||
|
|
||||||
all_results: list[FileResult] = []
|
all_results: list[FileResult] = []
|
||||||
overall_counter = [0]
|
overall_counter = [overall_base]
|
||||||
for b in range(n_batches):
|
for b in range(n_batches):
|
||||||
if stop_event is not None and stop_event.is_set():
|
if stop_event is not None and stop_event.is_set():
|
||||||
log("⏹ Отменено")
|
log("⏹ Отменено")
|
||||||
@@ -809,7 +895,7 @@ def _run_staged(
|
|||||||
try:
|
try:
|
||||||
batch_results = _run_staged_batch(
|
batch_results = _run_staged_batch(
|
||||||
config, chunk, output_root, log, on_progress, stop_event,
|
config, chunk, output_root, log, on_progress, stop_event,
|
||||||
overall_offset=start, overall_total=total,
|
overall_offset=overall_base + start, overall_total=overall_total,
|
||||||
on_overall_progress=on_overall_progress,
|
on_overall_progress=on_overall_progress,
|
||||||
overall_counter=overall_counter,
|
overall_counter=overall_counter,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -58,6 +58,21 @@ def _free_vram() -> None:
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _enable_tf32() -> None:
|
||||||
|
"""Enable TF32 matmul on Ampere+ GPUs (RTX 30/40).
|
||||||
|
|
||||||
|
Called lazily right before a WhisperX model is loaded, so runs that do not
|
||||||
|
use ASR never import Torch or create a CUDA context at startup.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
import torch
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.backends.cuda.matmul.allow_tf32 = True
|
||||||
|
torch.backends.cudnn.allow_tf32 = True
|
||||||
|
except Exception: # noqa: BLE001 — purely an optimization
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
def _model_transcribe_with_retry(
|
def _model_transcribe_with_retry(
|
||||||
model_obj, audio, language, batch_size, vad_filter, name="<audio>",
|
model_obj, audio, language, batch_size, vad_filter, name="<audio>",
|
||||||
):
|
):
|
||||||
@@ -98,6 +113,8 @@ class _WhisperXModels:
|
|||||||
compute_type = "int8_float16" if device == "cuda" else "int8"
|
compute_type = "int8_float16" if device == "cuda" else "int8"
|
||||||
key = (model_name, device, compute_type, local_files_only)
|
key = (model_name, device, compute_type, local_files_only)
|
||||||
if key not in cls._cache:
|
if key not in cls._cache:
|
||||||
|
if device == "cuda":
|
||||||
|
_enable_tf32()
|
||||||
import whisperx
|
import whisperx
|
||||||
print(f" [ASR] loading WhisperX {model_name} on {device} "
|
print(f" [ASR] loading WhisperX {model_name} on {device} "
|
||||||
f"({compute_type}, local_files_only={local_files_only})...")
|
f"({compute_type}, local_files_only={local_files_only})...")
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import shutil
|
import shutil
|
||||||
|
import subprocess
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
@@ -42,17 +43,30 @@ MODE_LABELS = {
|
|||||||
|
|
||||||
|
|
||||||
def _detect_device_info() -> tuple[str, str]:
|
def _detect_device_info() -> tuple[str, str]:
|
||||||
|
"""Report the GPU via nvidia-smi.
|
||||||
|
|
||||||
|
Torch must not be imported here: importing it initializes a CUDA context
|
||||||
|
and reserves VRAM even for runs that only need separation. The ASR stage
|
||||||
|
imports Torch lazily when it actually loads a model.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
import torch
|
out = subprocess.run(
|
||||||
except ImportError:
|
["nvidia-smi", "--query-gpu=name,memory.total",
|
||||||
return "⚠ PyTorch не установлен", "#f38ba8"
|
"--format=csv,noheader,nounits"],
|
||||||
if torch.cuda.is_available():
|
capture_output=True, text=True, timeout=5, check=True,
|
||||||
|
).stdout.strip()
|
||||||
|
except (OSError, subprocess.SubprocessError):
|
||||||
|
out = ""
|
||||||
|
if out:
|
||||||
|
parts = [p.strip() for p in out.splitlines()[0].split(",")]
|
||||||
|
name = parts[0] or "CUDA GPU"
|
||||||
|
if len(parts) > 1:
|
||||||
try:
|
try:
|
||||||
name = torch.cuda.get_device_name(0)
|
vram = float(parts[1]) / 1024 # MiB -> GiB
|
||||||
vram = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
|
|
||||||
return f"✓ CUDA: {name} · {vram:.1f} GB VRAM", "#a6e3a1"
|
return f"✓ CUDA: {name} · {vram:.1f} GB VRAM", "#a6e3a1"
|
||||||
except Exception: # noqa: BLE001
|
except ValueError:
|
||||||
return "✓ CUDA доступна", "#a6e3a1"
|
pass
|
||||||
|
return f"✓ CUDA: {name}", "#a6e3a1"
|
||||||
return "⚠ CUDA не обнаружена — расчёт пойдёт на CPU (медленно)", "#f9e2af"
|
return "⚠ CUDA не обнаружена — расчёт пойдёт на CPU (медленно)", "#f9e2af"
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
12
main.py
12
main.py
@@ -26,20 +26,8 @@ def preload_cudnn() -> None:
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
def enable_tf32() -> None:
|
|
||||||
"""Enable TF32 for faster matmul on Ampere+ GPUs (RTX 30/40 series)."""
|
|
||||||
try:
|
|
||||||
import torch
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
torch.backends.cuda.matmul.allow_tf32 = True
|
|
||||||
torch.backends.cudnn.allow_tf32 = True
|
|
||||||
except ImportError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
preload_cudnn()
|
preload_cudnn()
|
||||||
enable_tf32()
|
|
||||||
|
|
||||||
from PySide6.QtWidgets import QApplication
|
from PySide6.QtWidgets import QApplication
|
||||||
from gui.theme import DARK_STYLESHEET
|
from gui.theme import DARK_STYLESHEET
|
||||||
|
|||||||
Reference in New Issue
Block a user