mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 13:40:27 +02:00
refactor(domain): extract audio sink abstraction to domain layer
- New domain/audio_sink.py: AudioSink context manager + open_audio_sink() factory - Supports WAV/FLAC (soundfile) and MP3/Opus/M4B (ffmpeg pipe) - cancel_check, extra_ffmpeg_args, ffmpeg_cmd parameters - 17 tests in test_domain_audio_sink.py - WebUI: replaced local AudioSink + _open_audio_sink with domain module - Removed 35 lines, 3 call sites now use open_audio_sink() - PyQt: replaced 3 inline audio output setups with domain module - New _open_merged_sink() helper encapsulates m4b cover art logic - ExitStack for automatic cleanup in main conversion - Removed ~140 lines of duplicated ffmpeg/soundfile boilerplate
This commit is contained in:
Artem Akymenko
@@ -0,0 +1,131 @@
|
||||
"""Audio sink abstraction for unified audio output.
|
||||
|
||||
Provides a context-manager-based abstraction for writing audio data
|
||||
to various output formats (WAV, FLAC via soundfile; compressed via ffmpeg).
|
||||
|
||||
Usage:
|
||||
with open_audio_sink(path, "wav") as sink:
|
||||
sink.write(audio_data)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from abogen.domain.audio_buffer import SAMPLE_RATE
|
||||
from abogen.domain.audio_helpers import build_ffmpeg_command
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AudioSink:
|
||||
"""Represents an open audio output target."""
|
||||
|
||||
write: Callable[[np.ndarray], None]
|
||||
close: Callable[[], None]
|
||||
|
||||
def __enter__(self) -> AudioSink:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
def _ensure_ffmpeg() -> None:
|
||||
"""Ensure static ffmpeg binaries are on PATH."""
|
||||
import static_ffmpeg # type: ignore
|
||||
|
||||
ffmpeg_cache_root = _get_ffmpeg_cache_root()
|
||||
platform_cache = os.path.join(ffmpeg_cache_root, sys.platform)
|
||||
os.makedirs(platform_cache, exist_ok=True)
|
||||
try:
|
||||
import static_ffmpeg.run as static_ffmpeg_run # type: ignore
|
||||
|
||||
static_ffmpeg_run.LOCK_FILE = os.path.join(ffmpeg_cache_root, "lock.file")
|
||||
except Exception:
|
||||
pass
|
||||
static_ffmpeg.add_paths(weak=True, download_dir=platform_cache)
|
||||
|
||||
|
||||
def _get_ffmpeg_cache_root() -> str:
|
||||
from abogen.infrastructure.cache import get_internal_cache_path
|
||||
|
||||
return get_internal_cache_path("ffmpeg")
|
||||
|
||||
|
||||
def open_audio_sink(
|
||||
path: Path,
|
||||
fmt: str,
|
||||
*,
|
||||
metadata: Optional[dict[str, str]] = None,
|
||||
cancel_check: Optional[Callable[[], bool]] = None,
|
||||
extra_ffmpeg_args: Optional[list[str]] = None,
|
||||
ffmpeg_cmd: Optional[list[str]] = None,
|
||||
) -> AudioSink:
|
||||
"""Open an audio output sink for writing raw float32 PCM samples.
|
||||
|
||||
Args:
|
||||
path: Output file path.
|
||||
fmt: Output format ("wav", "flac", "mp3", "opus", "m4b").
|
||||
metadata: Optional metadata dict (ignored when ffmpeg_cmd is provided).
|
||||
cancel_check: Optional callable; if it returns True, writes are silently skipped.
|
||||
extra_ffmpeg_args: Optional extra args inserted after ffmpeg header (ignored when ffmpeg_cmd is provided).
|
||||
ffmpeg_cmd: Optional pre-built ffmpeg command list (for m4b with cover art etc.).
|
||||
|
||||
Returns:
|
||||
AudioSink with write() and close() methods.
|
||||
"""
|
||||
fmt = fmt.lower()
|
||||
|
||||
if fmt in {"wav", "flac"}:
|
||||
import soundfile as sf
|
||||
|
||||
soundfile_obj = sf.SoundFile(
|
||||
path,
|
||||
mode="w",
|
||||
samplerate=SAMPLE_RATE,
|
||||
channels=1,
|
||||
format=fmt.upper(),
|
||||
)
|
||||
|
||||
def _write_wav(data: np.ndarray) -> None:
|
||||
if cancel_check and cancel_check():
|
||||
return
|
||||
soundfile_obj.write(data)
|
||||
|
||||
def _close_wav() -> None:
|
||||
soundfile_obj.close()
|
||||
|
||||
return AudioSink(write=_write_wav, close=_close_wav)
|
||||
|
||||
# Compressed formats: pipe through ffmpeg
|
||||
_ensure_ffmpeg()
|
||||
|
||||
if ffmpeg_cmd is not None:
|
||||
cmd = list(ffmpeg_cmd)
|
||||
else:
|
||||
cmd = build_ffmpeg_command(path, fmt, metadata=metadata)
|
||||
if extra_ffmpeg_args:
|
||||
cmd[2:2] = extra_ffmpeg_args
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
|
||||
)
|
||||
|
||||
def _write_compressed(data: np.ndarray) -> None:
|
||||
if (cancel_check and cancel_check()) or process.stdin is None or process.stdin.closed:
|
||||
return
|
||||
process.stdin.write(data.tobytes())
|
||||
|
||||
def _close_compressed() -> None:
|
||||
if process.stdin and not process.stdin.closed:
|
||||
process.stdin.close()
|
||||
process.wait()
|
||||
|
||||
return AudioSink(write=_write_compressed, close=_close_compressed)
|
||||
+91
-190
@@ -5,6 +5,7 @@ import hashlib # For generating unique cache filenames
|
||||
from platformdirs import user_desktop_dir
|
||||
from PyQt6.QtCore import QThread, pyqtSignal, Qt, QTimer
|
||||
from PyQt6.QtWidgets import QCheckBox, QVBoxLayout, QDialog, QLabel, QDialogButtonBox
|
||||
from contextlib import ExitStack
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from abogen.utils import (
|
||||
@@ -29,6 +30,7 @@ from abogen.domain.output_paths import (
|
||||
sanitize_output_stem,
|
||||
)
|
||||
from abogen.domain.audio_helpers import build_ffmpeg_command, to_float32
|
||||
from abogen.domain.audio_sink import AudioSink, open_audio_sink
|
||||
from abogen.domain.audio_buffer import (
|
||||
create_silence,
|
||||
mix_audio,
|
||||
@@ -361,6 +363,7 @@ class ConversionThread(QThread):
|
||||
)
|
||||
try:
|
||||
hf_tracker.set_log_callback(lambda msg: self.log_updated.emit(msg))
|
||||
sink_stack = ExitStack()
|
||||
# Show configuration
|
||||
self.log_updated.emit("Configuration:")
|
||||
|
||||
@@ -735,57 +738,12 @@ class ConversionThread(QThread):
|
||||
# SRT numbering fix: use a global counter
|
||||
merged_srt_index = 1 # SRT numbering for merged file
|
||||
# Prepare output file/ffmpeg process for merged output
|
||||
if self.output_format in ["wav", "mp3", "flac"]:
|
||||
merged_out_file = sf.SoundFile(
|
||||
merged_sink = sink_stack.enter_context(
|
||||
self._open_merged_sink(
|
||||
merged_out_path,
|
||||
"w",
|
||||
samplerate=24000,
|
||||
channels=1,
|
||||
format=self.output_format,
|
||||
cancel_check=lambda: self.cancel_requested,
|
||||
)
|
||||
ffmpeg_proc = None
|
||||
elif self.output_format in ("m4b", "opus"):
|
||||
# Real-time generation using FFmpeg pipe
|
||||
static_ffmpeg.add_paths()
|
||||
merged_out_file = None
|
||||
ffmpeg_proc = None
|
||||
metadata_options, cover_path = (
|
||||
self._extract_and_add_metadata_tags_to_ffmpeg_cmd()
|
||||
if self.output_format == "m4b"
|
||||
else ([], None)
|
||||
)
|
||||
cmd = build_ffmpeg_command(
|
||||
Path(merged_out_path),
|
||||
self.output_format,
|
||||
)
|
||||
# Insert thread queue size after ffmpeg header
|
||||
cmd.insert(2, "-thread_queue_size")
|
||||
cmd.insert(3, "32768")
|
||||
if self.output_format == "m4b" and cover_path and os.path.exists(cover_path):
|
||||
# Insert cover image input before the output path
|
||||
output_path = cmd.pop()
|
||||
cmd.extend([
|
||||
"-i", cover_path,
|
||||
"-map", "0:a",
|
||||
"-map", "1",
|
||||
"-c:v", "copy",
|
||||
"-disposition:v", "attached_pic",
|
||||
])
|
||||
cmd.extend(metadata_options)
|
||||
cmd.append(output_path)
|
||||
elif self.output_format == "m4b":
|
||||
output_path = cmd.pop()
|
||||
cmd.extend(metadata_options)
|
||||
cmd.append(output_path)
|
||||
ffmpeg_proc = create_process(cmd, stdin=subprocess.PIPE, text=False)
|
||||
else:
|
||||
self.log_updated.emit(
|
||||
(f"Unsupported output format: {self.output_format}", "red")
|
||||
)
|
||||
self.conversion_finished.emit(
|
||||
("Audio generation failed.", "red"), None
|
||||
)
|
||||
return
|
||||
# Open merged subtitle file for incremental writing if needed
|
||||
merged_subtitle_file = None
|
||||
if self.subtitle_mode != "Disabled":
|
||||
@@ -849,9 +807,8 @@ class ConversionThread(QThread):
|
||||
merged_subtitle_path = None
|
||||
merged_subtitle_file = None
|
||||
else:
|
||||
# If not merging, set merged_out_file and related variables to None
|
||||
merged_out_file = None
|
||||
ffmpeg_proc = None
|
||||
# If not merging, set merged_sink and related variables to None
|
||||
merged_sink = None
|
||||
merged_out_path = None
|
||||
subtitle_entries = []
|
||||
current_time = 0.0
|
||||
@@ -868,8 +825,7 @@ class ConversionThread(QThread):
|
||||
# Instead of processing the whole text, process by chapter
|
||||
for chapter_idx, (chapter_name, voice_segments) in enumerate(chapters, 1):
|
||||
chapter_out_path = None
|
||||
chapter_out_file = None
|
||||
chapter_ffmpeg_proc = None
|
||||
chapter_sink = None
|
||||
chapter_subtitle_file = None
|
||||
chapter_subtitle_path = None
|
||||
if total_chapters > 1:
|
||||
@@ -904,37 +860,12 @@ class ConversionThread(QThread):
|
||||
chapters_out_dir,
|
||||
f"{chapter_filename}.{separate_chapters_format}",
|
||||
)
|
||||
if separate_chapters_format in ["wav", "mp3", "flac"]:
|
||||
chapter_out_file = sf.SoundFile(
|
||||
chapter_out_path,
|
||||
"w",
|
||||
samplerate=24000,
|
||||
channels=1,
|
||||
format=separate_chapters_format,
|
||||
)
|
||||
chapter_ffmpeg_proc = None
|
||||
elif separate_chapters_format == "opus":
|
||||
static_ffmpeg.add_paths()
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-y",
|
||||
"-thread_queue_size",
|
||||
"32768",
|
||||
"-f",
|
||||
"f32le",
|
||||
"-ar",
|
||||
"24000",
|
||||
"-ac",
|
||||
"1",
|
||||
"-i",
|
||||
"pipe:0",
|
||||
]
|
||||
cmd.extend(["-c:a", "libopus", "-b:a", "24000"])
|
||||
cmd.append(chapter_out_path)
|
||||
chapter_ffmpeg_proc = create_process(
|
||||
cmd, stdin=subprocess.PIPE, text=False
|
||||
)
|
||||
chapter_out_file = None
|
||||
if separate_chapters_format in ("wav", "flac", "mp3", "opus"):
|
||||
chapter_sink = sink_stack.enter_context(open_audio_sink(
|
||||
Path(chapter_out_path),
|
||||
separate_chapters_format,
|
||||
cancel_check=lambda: self.cancel_requested,
|
||||
))
|
||||
else:
|
||||
self.log_updated.emit(
|
||||
(
|
||||
@@ -1129,10 +1060,7 @@ class ConversionThread(QThread):
|
||||
# Print the result for debugging
|
||||
# print(f"Result: {result}")
|
||||
if self.cancel_requested:
|
||||
if chapter_out_file:
|
||||
chapter_out_file.close()
|
||||
if merged_out_file:
|
||||
merged_out_file.close()
|
||||
sink_stack.close()
|
||||
self.conversion_finished.emit("Cancelled", None)
|
||||
return
|
||||
current_segment += 1
|
||||
@@ -1147,14 +1075,10 @@ class ConversionThread(QThread):
|
||||
chunk_start = current_time
|
||||
audio_np = to_float32(result.audio)
|
||||
# Write audio directly to merged file ONLY if merging
|
||||
if merge_chapters_at_end and merged_out_file:
|
||||
merged_out_file.write(audio_np)
|
||||
elif merge_chapters_at_end and ffmpeg_proc:
|
||||
ffmpeg_proc.stdin.write(audio_np.tobytes())
|
||||
if chapter_out_file:
|
||||
chapter_out_file.write(audio_np)
|
||||
elif chapter_ffmpeg_proc:
|
||||
chapter_ffmpeg_proc.stdin.write(audio_np.tobytes())
|
||||
if merge_chapters_at_end and merged_sink:
|
||||
merged_sink.write(audio_np)
|
||||
if chapter_sink:
|
||||
chapter_sink.write(audio_np)
|
||||
# Subtitle logic
|
||||
if self.subtitle_mode != "Disabled":
|
||||
tokens_list = getattr(result, "tokens", [])
|
||||
@@ -1187,7 +1111,7 @@ class ConversionThread(QThread):
|
||||
"whitespace": tok.whitespace,
|
||||
}
|
||||
)
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
if chapter_sink:
|
||||
chapter_tokens_with_timestamps.append(
|
||||
{
|
||||
"start": chapter_current_time
|
||||
@@ -1234,8 +1158,8 @@ class ConversionThread(QThread):
|
||||
f"{merged_srt_index}\n{_format_timestamp(start)} --> {_format_timestamp(end)}\n{text}\n\n"
|
||||
)
|
||||
merged_srt_index += 1
|
||||
# Per-chapter subtitle processing for both file and ffmpeg_proc
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
# Per-chapter subtitle processing for both file and sink
|
||||
if chapter_sink:
|
||||
new_chapter_entries = []
|
||||
self._process_subtitle_tokens(
|
||||
chapter_tokens_with_timestamps,
|
||||
@@ -1270,10 +1194,10 @@ class ConversionThread(QThread):
|
||||
chapter_srt_index += 1
|
||||
if merge_chapters_at_end:
|
||||
current_time += chunk_dur
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
if chapter_sink:
|
||||
chapter_current_time += chunk_dur
|
||||
else:
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
if chapter_sink:
|
||||
chapter_current_time += chunk_dur
|
||||
# Calculate percentage based on characters processed
|
||||
percent = min(
|
||||
@@ -1296,29 +1220,22 @@ class ConversionThread(QThread):
|
||||
# Add silence between chapters for merged output (except after the last chapter)
|
||||
if merge_chapters_at_end and chapter_idx < total_chapters:
|
||||
silence_audio = create_silence(self.silence_duration)
|
||||
silence_bytes = silence_audio.tobytes()
|
||||
|
||||
if merged_out_file:
|
||||
merged_out_file.write(silence_audio)
|
||||
elif ffmpeg_proc:
|
||||
ffmpeg_proc.stdin.write(silence_bytes)
|
||||
if merged_sink:
|
||||
merged_sink.write(silence_audio)
|
||||
|
||||
|
||||
# Update timing for the silence
|
||||
current_time += self.silence_duration
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
if chapter_sink:
|
||||
chapter_current_time += self.silence_duration
|
||||
|
||||
# Set chapter end time after processing
|
||||
if merge_chapters_at_end:
|
||||
chapter_time["end"] = current_time
|
||||
# Finalize chapter file for ffmpeg formats
|
||||
if chapter_out_file or chapter_ffmpeg_proc:
|
||||
if chapter_sink:
|
||||
self.log_updated.emit(("\nProcessing chapter audio...", "grey"))
|
||||
if chapter_ffmpeg_proc:
|
||||
chapter_ffmpeg_proc.stdin.close()
|
||||
chapter_ffmpeg_proc.wait()
|
||||
if chapter_out_file:
|
||||
chapter_out_file.close()
|
||||
# Close chapter subtitle file if open
|
||||
if chapter_subtitle_file:
|
||||
chapter_subtitle_file.close()
|
||||
@@ -1344,11 +1261,8 @@ class ConversionThread(QThread):
|
||||
# Finalize merged output file ONLY if merging
|
||||
if merge_chapters_at_end:
|
||||
self.log_updated.emit(("\nFinalizing audio. Please wait...", "grey"))
|
||||
if self.output_format in ["wav", "mp3", "flac"]:
|
||||
merged_out_file.close()
|
||||
elif self.output_format == "m4b":
|
||||
ffmpeg_proc.stdin.close()
|
||||
ffmpeg_proc.wait()
|
||||
merged_sink.close()
|
||||
if self.output_format == "m4b":
|
||||
# Add chapters via fast post-processing
|
||||
if total_chapters > 1:
|
||||
chapters_info_path = f"{base_filepath_no_ext}_chapters.txt"
|
||||
@@ -1412,8 +1326,7 @@ class ConversionThread(QThread):
|
||||
os.replace(tmp_path, orig_path)
|
||||
os.remove(chapters_info_path)
|
||||
elif self.output_format in ["opus"]:
|
||||
ffmpeg_proc.stdin.close()
|
||||
ffmpeg_proc.wait()
|
||||
merged_sink.close()
|
||||
self.progress_updated.emit(100, "00:00:00")
|
||||
# Close merged subtitle file if open
|
||||
if merged_subtitle_file:
|
||||
@@ -1442,23 +1355,15 @@ class ConversionThread(QThread):
|
||||
chapters_dir,
|
||||
)
|
||||
except Exception as e:
|
||||
# Cleanup ffmpeg subprocesses on error
|
||||
# Cleanup all sinks on error
|
||||
try:
|
||||
if "ffmpeg_proc" in locals() and ffmpeg_proc:
|
||||
ffmpeg_proc.stdin.close()
|
||||
ffmpeg_proc.terminate()
|
||||
ffmpeg_proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if "chapter_ffmpeg_proc" in locals() and chapter_ffmpeg_proc:
|
||||
chapter_ffmpeg_proc.stdin.close()
|
||||
chapter_ffmpeg_proc.terminate()
|
||||
chapter_ffmpeg_proc.wait()
|
||||
sink_stack.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.log_updated.emit((f"Error occurred: {str(e)}", "red"))
|
||||
self.conversion_finished.emit(("Audio generation failed.", "red"), None)
|
||||
finally:
|
||||
sink_stack.close()
|
||||
|
||||
def _process_subtitle_file(self, tts, base_path, is_timestamp_text=False):
|
||||
"""Process subtitle files with precise timing and generate output subtitles."""
|
||||
@@ -1527,39 +1432,10 @@ class ConversionThread(QThread):
|
||||
rate = 24000
|
||||
|
||||
# Setup audio output
|
||||
merged_out_file, ffmpeg_proc = None, None
|
||||
if self.output_format in ["wav", "mp3", "flac"]:
|
||||
merged_out_file = sf.SoundFile(
|
||||
merged_sink = self._open_merged_sink(
|
||||
merged_out_path,
|
||||
"w",
|
||||
samplerate=rate,
|
||||
channels=1,
|
||||
format=self.output_format,
|
||||
)
|
||||
else:
|
||||
static_ffmpeg.add_paths()
|
||||
cmd = build_ffmpeg_command(
|
||||
Path(merged_out_path),
|
||||
self.output_format,
|
||||
)
|
||||
cmd.insert(2, "-thread_queue_size")
|
||||
cmd.insert(3, "32768")
|
||||
if self.output_format == "m4b":
|
||||
metadata_options, cover_path = (
|
||||
self._extract_and_add_metadata_tags_to_ffmpeg_cmd()
|
||||
)
|
||||
if cover_path and os.path.exists(cover_path):
|
||||
output_path = cmd.pop()
|
||||
cmd.extend([
|
||||
"-i", cover_path,
|
||||
"-map", "0:a",
|
||||
"-map", "1",
|
||||
"-c:v", "copy",
|
||||
"-disposition:v", "attached_pic",
|
||||
])
|
||||
cmd.append(output_path)
|
||||
cmd.extend(metadata_options)
|
||||
ffmpeg_proc = create_process(cmd, stdin=subprocess.PIPE, text=False)
|
||||
cancel_check=lambda: self.cancel_requested,
|
||||
).__enter__()
|
||||
|
||||
# Always generate subtitles for subtitle input files
|
||||
subtitle_file, subtitle_path = None, None
|
||||
@@ -1851,13 +1727,9 @@ class ConversionThread(QThread):
|
||||
|
||||
# Write the complete audio buffer
|
||||
self.log_updated.emit(("\nFinalizing audio. Please wait...", "grey"))
|
||||
if merged_out_file:
|
||||
merged_out_file.write(audio_buffer)
|
||||
merged_out_file.close()
|
||||
elif ffmpeg_proc:
|
||||
ffmpeg_proc.stdin.write(to_float32(audio_buffer).tobytes())
|
||||
ffmpeg_proc.stdin.close()
|
||||
ffmpeg_proc.wait()
|
||||
if merged_sink:
|
||||
merged_sink.write(to_float32(audio_buffer))
|
||||
merged_sink.close()
|
||||
|
||||
if subtitle_file:
|
||||
subtitle_file.close()
|
||||
@@ -1870,10 +1742,8 @@ class ConversionThread(QThread):
|
||||
|
||||
except Exception as e:
|
||||
try:
|
||||
if "ffmpeg_proc" in locals() and ffmpeg_proc:
|
||||
ffmpeg_proc.stdin.close()
|
||||
ffmpeg_proc.terminate()
|
||||
ffmpeg_proc.wait()
|
||||
if "merged_sink" in locals() and merged_sink:
|
||||
merged_sink.close()
|
||||
if "subtitle_file" in locals() and subtitle_file:
|
||||
subtitle_file.close()
|
||||
except:
|
||||
@@ -1893,6 +1763,52 @@ class ConversionThread(QThread):
|
||||
self._timestamp_response = treat_as_subtitle
|
||||
self._timestamp_response_event.set()
|
||||
|
||||
@contextmanager
|
||||
def _open_merged_sink(self, path, cancel_check=None):
|
||||
"""Open audio sink for the merged output, handling m4b cover art."""
|
||||
if self.output_format in ("wav", "flac", "mp3", "opus"):
|
||||
with open_audio_sink(
|
||||
Path(path),
|
||||
self.output_format,
|
||||
cancel_check=cancel_check,
|
||||
) as sink:
|
||||
yield sink
|
||||
elif self.output_format == "m4b":
|
||||
static_ffmpeg.add_paths()
|
||||
metadata_options, cover_path = (
|
||||
self._extract_and_add_metadata_tags_to_ffmpeg_cmd()
|
||||
)
|
||||
cmd = build_ffmpeg_command(
|
||||
Path(path),
|
||||
self.output_format,
|
||||
)
|
||||
cmd.insert(2, "-thread_queue_size")
|
||||
cmd.insert(3, "32768")
|
||||
if cover_path and os.path.exists(cover_path):
|
||||
output_path = cmd.pop()
|
||||
cmd.extend([
|
||||
"-i", cover_path,
|
||||
"-map", "0:a",
|
||||
"-map", "1",
|
||||
"-c:v", "copy",
|
||||
"-disposition:v", "attached_pic",
|
||||
])
|
||||
cmd.extend(metadata_options)
|
||||
cmd.append(output_path)
|
||||
else:
|
||||
output_path = cmd.pop()
|
||||
cmd.extend(metadata_options)
|
||||
cmd.append(output_path)
|
||||
with open_audio_sink(
|
||||
Path(path),
|
||||
self.output_format,
|
||||
ffmpeg_cmd=cmd,
|
||||
cancel_check=cancel_check,
|
||||
) as sink:
|
||||
yield sink
|
||||
else:
|
||||
raise ValueError(f"Unsupported output format: {self.output_format}")
|
||||
|
||||
def _extract_and_add_metadata_tags_to_ffmpeg_cmd(self):
|
||||
"""Extract metadata tags from text content and add them to ffmpeg command"""
|
||||
# Read text for metadata extraction
|
||||
@@ -1960,21 +1876,6 @@ class ConversionThread(QThread):
|
||||
self.process.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
# Terminate ffmpeg subprocesses if running
|
||||
try:
|
||||
if hasattr(self, "ffmpeg_proc") and self.ffmpeg_proc:
|
||||
self.ffmpeg_proc.stdin.close()
|
||||
self.ffmpeg_proc.terminate()
|
||||
self.ffmpeg_proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if hasattr(self, "chapter_ffmpeg_proc") and self.chapter_ffmpeg_proc:
|
||||
self.chapter_ffmpeg_proc.stdin.close()
|
||||
self.chapter_ffmpeg_proc.terminate()
|
||||
self.chapter_ffmpeg_proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class VoicePreviewThread(QThread):
|
||||
|
||||
@@ -2,8 +2,6 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import gc
|
||||
@@ -14,8 +12,6 @@ from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Mapping, Optional
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import static_ffmpeg
|
||||
|
||||
from abogen.tts_plugin.utils import is_plugin_registered
|
||||
from abogen.infrastructure.exporters import ExportService
|
||||
@@ -124,6 +120,7 @@ from abogen.domain.audio_buffer import (
|
||||
normalize_audio as _normalize_audio,
|
||||
SAMPLE_RATE,
|
||||
)
|
||||
from abogen.domain.audio_sink import AudioSink, open_audio_sink
|
||||
|
||||
|
||||
from .service import Job, JobStatus
|
||||
@@ -149,11 +146,6 @@ class _JobCancelled(Exception):
|
||||
"""Raised internally to abort a conversion when the client cancels."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioSink:
|
||||
write: Callable[[np.ndarray], None]
|
||||
|
||||
|
||||
_APOSTROPHE_CONFIG = ApostropheConfig()
|
||||
|
||||
|
||||
@@ -402,7 +394,14 @@ def run_conversion_job(job: Job) -> None:
|
||||
if merged_required:
|
||||
audio_path = _build_output_path(audio_dir, job.original_filename, job.output_format)
|
||||
meta_for_sink = job.metadata_tags if job.metadata_tags else None
|
||||
audio_sink = _open_audio_sink(audio_path, job, sink_stack, metadata=meta_for_sink)
|
||||
audio_sink = sink_stack.enter_context(
|
||||
open_audio_sink(
|
||||
audio_path,
|
||||
job.output_format,
|
||||
metadata=meta_for_sink,
|
||||
cancel_check=lambda: job.cancel_requested,
|
||||
)
|
||||
)
|
||||
subtitle_writer = _create_subtitle_writer(job, audio_path)
|
||||
job.result.audio_path = audio_path
|
||||
if subtitle_writer:
|
||||
@@ -638,11 +637,12 @@ def run_conversion_job(job: Job) -> None:
|
||||
f"{Path(job.original_filename).stem}_{_slugify(chapter_display_title, idx)}",
|
||||
job.separate_chapters_format,
|
||||
)
|
||||
chapter_sink = _open_audio_sink(
|
||||
chapter_sink = chapter_sink_stack.enter_context(
|
||||
open_audio_sink(
|
||||
chapter_audio_path,
|
||||
job,
|
||||
chapter_sink_stack,
|
||||
fmt=job.separate_chapters_format,
|
||||
job.separate_chapters_format,
|
||||
cancel_check=lambda: job.cancel_requested,
|
||||
)
|
||||
)
|
||||
|
||||
speak_heading = bool(heading_text)
|
||||
@@ -907,11 +907,12 @@ def run_conversion_job(job: Job) -> None:
|
||||
f"{Path(job.original_filename).stem}_outro",
|
||||
job.separate_chapters_format,
|
||||
)
|
||||
chapter_sink = _open_audio_sink(
|
||||
chapter_sink = outro_sink_stack.enter_context(
|
||||
open_audio_sink(
|
||||
outro_audio_path,
|
||||
job,
|
||||
outro_sink_stack,
|
||||
fmt=job.separate_chapters_format,
|
||||
job.separate_chapters_format,
|
||||
cancel_check=lambda: job.cancel_requested,
|
||||
)
|
||||
)
|
||||
|
||||
outro_segments = emit_text(
|
||||
@@ -1148,50 +1149,6 @@ def _prepare_project_layout(job: Job, base_dir: Path) -> tuple[Path, Path, Path,
|
||||
)
|
||||
|
||||
|
||||
def _open_audio_sink(
|
||||
path: Path,
|
||||
job: Job,
|
||||
stack: ExitStack,
|
||||
*,
|
||||
fmt: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
) -> AudioSink:
|
||||
ffmpeg_cache_root = get_internal_cache_path("ffmpeg")
|
||||
platform_cache = os.path.join(ffmpeg_cache_root, sys.platform)
|
||||
os.makedirs(platform_cache, exist_ok=True)
|
||||
try:
|
||||
import static_ffmpeg.run as static_ffmpeg_run # type: ignore
|
||||
|
||||
static_ffmpeg_run.LOCK_FILE = os.path.join(ffmpeg_cache_root, "lock.file")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
static_ffmpeg.add_paths(weak=True, download_dir=platform_cache)
|
||||
fmt_value = (fmt or job.output_format).lower()
|
||||
|
||||
if fmt_value in {"wav", "flac"}:
|
||||
soundfile = stack.enter_context(
|
||||
sf.SoundFile(path, mode="w", samplerate=SAMPLE_RATE, channels=1, format=fmt_value.upper())
|
||||
)
|
||||
return AudioSink(write=lambda data: soundfile.write(data))
|
||||
|
||||
cmd = _build_ffmpeg_command(path, fmt_value, metadata=metadata)
|
||||
process = create_process(cmd, stdin=subprocess.PIPE, text=False)
|
||||
|
||||
def _finalize() -> None:
|
||||
if process.stdin and not process.stdin.closed:
|
||||
process.stdin.close()
|
||||
process.wait()
|
||||
|
||||
stack.callback(_finalize)
|
||||
|
||||
def _write(data: np.ndarray) -> None:
|
||||
if job.cancel_requested or process.stdin is None:
|
||||
return
|
||||
process.stdin.write(data.tobytes()) # type: ignore[arg-type]
|
||||
|
||||
return AudioSink(write=_write)
|
||||
|
||||
|
||||
def _resolve_voice(pipeline, voice_spec: str, use_gpu: bool):
|
||||
if "*" in voice_spec:
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Tests for domain/audio_sink.py"""
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch, call
|
||||
import subprocess
|
||||
|
||||
from abogen.domain.audio_sink import AudioSink, open_audio_sink, _ensure_ffmpeg
|
||||
|
||||
|
||||
class TestAudioSinkDataclass:
|
||||
def test_audio_sink_is_frozen(self):
|
||||
sink = AudioSink(write=lambda d: None, close=lambda: None)
|
||||
try:
|
||||
sink.write = lambda d: None # type: ignore
|
||||
except AttributeError:
|
||||
pass # Expected — frozen=True
|
||||
assert hasattr(sink, "write")
|
||||
assert hasattr(sink, "close")
|
||||
|
||||
def test_audio_sink_stores_callables(self):
|
||||
write_fn = MagicMock()
|
||||
close_fn = MagicMock()
|
||||
sink = AudioSink(write=write_fn, close=close_fn)
|
||||
assert sink.write is write_fn
|
||||
assert sink.close is close_fn
|
||||
|
||||
def test_audio_sink_write_callable(self):
|
||||
calls = []
|
||||
sink = AudioSink(write=lambda d: calls.append(d), close=lambda: None)
|
||||
data = np.zeros(100, dtype="float32")
|
||||
sink.write(data)
|
||||
assert len(calls) == 1
|
||||
np.testing.assert_array_equal(calls[0], data)
|
||||
|
||||
def test_audio_sink_close_callable(self):
|
||||
closed = []
|
||||
sink = AudioSink(write=lambda d: None, close=lambda: closed.append(True))
|
||||
sink.close()
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
class TestOpenAudioSinkWav:
|
||||
def test_wav_creates_file(self, tmp_path: Path):
|
||||
out = tmp_path / "test.wav"
|
||||
with open_audio_sink(out, "wav") as sink:
|
||||
sink.write(np.zeros(100, dtype="float32"))
|
||||
assert out.exists()
|
||||
|
||||
def test_flac_creates_file(self, tmp_path: Path):
|
||||
out = tmp_path / "test.flac"
|
||||
with open_audio_sink(out, "flac") as sink:
|
||||
sink.write(np.zeros(100, dtype="float32"))
|
||||
assert out.exists()
|
||||
|
||||
def test_wav_sink_writes_audio(self, tmp_path: Path):
|
||||
out = tmp_path / "out.wav"
|
||||
audio = np.random.uniform(-0.5, 0.5, 24000).astype("float32") # 1 second
|
||||
with open_audio_sink(out, "wav") as sink:
|
||||
sink.write(audio)
|
||||
data, sr = sf.read(str(out))
|
||||
assert sr == 24000
|
||||
assert len(data) == 24000
|
||||
np.testing.assert_allclose(data, audio, atol=1e-4)
|
||||
|
||||
def test_wav_sink_close_flushes(self, tmp_path: Path):
|
||||
out = tmp_path / "flush.wav"
|
||||
with open_audio_sink(out, "wav") as sink:
|
||||
sink.write(np.ones(1000, dtype="float32"))
|
||||
assert out.exists()
|
||||
info = sf.info(str(out))
|
||||
assert info.samplerate == 24000
|
||||
assert info.channels == 1
|
||||
|
||||
def test_wav_sink_context_manager(self, tmp_path: Path):
|
||||
out = tmp_path / "ctx.wav"
|
||||
with open_audio_sink(out, "wav") as sink:
|
||||
assert isinstance(sink, AudioSink)
|
||||
sink.write(np.zeros(50, dtype="float32"))
|
||||
assert out.exists()
|
||||
|
||||
def test_wav_sink_multiple_writes(self, tmp_path: Path):
|
||||
out = tmp_path / "multi.wav"
|
||||
with open_audio_sink(out, "wav") as sink:
|
||||
sink.write(np.ones(1000, dtype="float32"))
|
||||
sink.write(np.ones(500, dtype="float32"))
|
||||
data, _ = sf.read(str(out))
|
||||
assert len(data) == 1500
|
||||
|
||||
|
||||
class TestCancelCheck:
|
||||
def test_cancel_check_skips_wav_writes(self, tmp_path: Path):
|
||||
out = tmp_path / "cancelled.wav"
|
||||
with open_audio_sink(out, "wav", cancel_check=lambda: True) as sink:
|
||||
sink.write(np.ones(1000, dtype="float32"))
|
||||
sink.write(np.ones(500, dtype="float32"))
|
||||
data, _ = sf.read(str(out))
|
||||
assert len(data) == 0
|
||||
|
||||
def test_cancel_check_none_allows_writes(self, tmp_path: Path):
|
||||
out = tmp_path / "ok.wav"
|
||||
with open_audio_sink(out, "wav", cancel_check=None) as sink:
|
||||
sink.write(np.ones(1000, dtype="float32"))
|
||||
data, _ = sf.read(str(out))
|
||||
assert len(data) == 1000
|
||||
|
||||
|
||||
class TestUnsupportedFormat:
|
||||
def test_unsupported_format_raises(self, tmp_path: Path):
|
||||
out = tmp_path / "bad.xyz"
|
||||
try:
|
||||
with open_audio_sink(out, "xyz") as sink:
|
||||
pass
|
||||
assert False, "Should have raised"
|
||||
except Exception:
|
||||
pass # Expected
|
||||
|
||||
|
||||
class TestOpenAudioSinkCompressed:
|
||||
@patch("abogen.domain.audio_sink._ensure_ffmpeg")
|
||||
@patch("abogen.domain.audio_sink.build_ffmpeg_command")
|
||||
@patch("abogen.domain.audio_sink.subprocess.Popen")
|
||||
def test_mp3_sink_returns_sink(self, mock_popen, mock_build, mock_ensure, tmp_path: Path):
|
||||
mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.mp3"]
|
||||
proc = MagicMock()
|
||||
proc.stdin = MagicMock()
|
||||
proc.stdin.closed = False
|
||||
proc.wait.return_value = 0
|
||||
mock_popen.return_value = proc
|
||||
|
||||
out = tmp_path / "test.mp3"
|
||||
sink = open_audio_sink(out, "mp3")
|
||||
assert isinstance(sink, AudioSink)
|
||||
sink.write(np.zeros(100, dtype="float32"))
|
||||
assert proc.stdin.write.called
|
||||
sink.close()
|
||||
|
||||
@patch("abogen.domain.audio_sink._ensure_ffmpeg")
|
||||
@patch("abogen.domain.audio_sink.build_ffmpeg_command")
|
||||
@patch("abogen.domain.audio_sink.subprocess.Popen")
|
||||
def test_cancel_check_skips_compressed_writes(self, mock_popen, mock_build, mock_ensure, tmp_path: Path):
|
||||
mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.mp3"]
|
||||
proc = MagicMock()
|
||||
proc.stdin = MagicMock()
|
||||
proc.stdin.closed = False
|
||||
proc.wait.return_value = 0
|
||||
mock_popen.return_value = proc
|
||||
|
||||
sink = open_audio_sink(tmp_path / "c.mp3", "mp3", cancel_check=lambda: True)
|
||||
sink.write(np.zeros(100, dtype="float32"))
|
||||
assert not proc.stdin.write.called
|
||||
sink.close()
|
||||
|
||||
@patch("abogen.domain.audio_sink._ensure_ffmpeg")
|
||||
@patch("abogen.domain.audio_sink.build_ffmpeg_command")
|
||||
@patch("abogen.domain.audio_sink.subprocess.Popen")
|
||||
def test_extra_ffmpeg_args_passed(self, mock_popen, mock_build, mock_ensure, tmp_path: Path):
|
||||
mock_build.return_value = ["ffmpeg", "-y", "-header", "-i", "pipe:0", "out.mp3"]
|
||||
proc = MagicMock()
|
||||
proc.stdin = MagicMock()
|
||||
proc.stdin.closed = False
|
||||
proc.wait.return_value = 0
|
||||
mock_popen.return_value = proc
|
||||
|
||||
sink = open_audio_sink(
|
||||
tmp_path / "extra.mp3",
|
||||
"mp3",
|
||||
extra_ffmpeg_args=["-thread_queue_size", "32768"],
|
||||
)
|
||||
args = mock_popen.call_args[0][0]
|
||||
assert "-thread_queue_size" in args
|
||||
assert "32768" in args
|
||||
sink.close()
|
||||
|
||||
@patch("abogen.domain.audio_sink._ensure_ffmpeg")
|
||||
@patch("abogen.domain.audio_sink.build_ffmpeg_command")
|
||||
@patch("abogen.domain.audio_sink.subprocess.Popen")
|
||||
def test_metadata_passed_to_ffmpeg(self, mock_popen, mock_build, mock_ensure, tmp_path: Path):
|
||||
mock_build.return_value = ["ffmpeg", "-y", "-i", "pipe:0", "out.opus"]
|
||||
proc = MagicMock()
|
||||
proc.stdin = MagicMock()
|
||||
proc.stdin.closed = False
|
||||
proc.wait.return_value = 0
|
||||
mock_popen.return_value = proc
|
||||
|
||||
meta = {"title": "Test", "artist": "Author"}
|
||||
open_audio_sink(tmp_path / "meta.opus", "opus", metadata=meta)
|
||||
mock_build.assert_called_once_with(tmp_path / "meta.opus", "opus", metadata=meta)
|
||||
Reference in New Issue
Block a user