feat: add remote audiobook generation
This commit is contained in:
Курнат Андрей
@@ -0,0 +1,688 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import hashlib
|
||||
import hmac
|
||||
import html
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import shutil
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import zipfile
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Iterator
|
||||
from xml.etree import ElementTree
|
||||
|
||||
import imageio_ffmpeg
|
||||
import soundfile as sf
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
from num2words import num2words
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
|
||||
ROOT = Path(os.environ.get("AUDIOBOOK_SERVICE_ROOT", "/home/sevenhill/apps/aletheia-audiobook"))
|
||||
SERVICE_ROOT = ROOT / "service"
|
||||
DATA_ROOT = ROOT / "audiobook-jobs"
|
||||
TOKEN_FILE = SERVICE_ROOT / "service-token.txt"
|
||||
DATABASE = DATA_ROOT / "jobs.db3"
|
||||
PERFORMANCE_LOG = DATA_ROOT / "performance.jsonl"
|
||||
MODEL_ID = os.environ.get("QWEN_TTS_MODEL_ID", "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
|
||||
MAX_UPLOAD_BYTES = 200 * 1024 * 1024
|
||||
GPU_IDS = tuple(
|
||||
int(value.strip())
|
||||
for value in os.environ.get("QWEN_TTS_GPU_IDS", "0,1,2").split(",")
|
||||
if value.strip()
|
||||
)
|
||||
SPEAKER = os.environ.get("QWEN_TTS_SPEAKER", "Ryan")
|
||||
VOICE_INSTRUCTION = (
|
||||
"Read in a calm, clear, natural audiobook style with steady pacing and distinct diction."
|
||||
)
|
||||
|
||||
SERVICE_ROOT.mkdir(parents=True, exist_ok=True)
|
||||
DATA_ROOT.mkdir(parents=True, exist_ok=True)
|
||||
OLLAMA_PAUSE_MARKER = SERVICE_ROOT / "ollama-paused-for-audiobook"
|
||||
if not TOKEN_FILE.exists():
|
||||
TOKEN_FILE.write_text(secrets.token_urlsafe(48), encoding="ascii")
|
||||
API_TOKEN = TOKEN_FILE.read_text(encoding="ascii").strip()
|
||||
|
||||
|
||||
def pause_ollama_for_audiobook() -> None:
|
||||
"""Release every CMP GPU before loading the three Qwen workers."""
|
||||
subprocess.run(["systemctl", "--user", "stop", "ollama.service"], check=True)
|
||||
subprocess.run(["systemctl", "--user", "mask", "--runtime", "ollama.service"], check=True)
|
||||
OLLAMA_PAUSE_MARKER.write_text("1", encoding="ascii")
|
||||
|
||||
|
||||
def restore_ollama_after_audiobook() -> None:
|
||||
if not OLLAMA_PAUSE_MARKER.exists():
|
||||
return
|
||||
subprocess.run(["systemctl", "--user", "unmask", "--runtime", "ollama.service"], check=True)
|
||||
subprocess.run(["systemctl", "--user", "start", "ollama.service"], check=True)
|
||||
OLLAMA_PAUSE_MARKER.unlink(missing_ok=True)
|
||||
|
||||
|
||||
# The worker terminates after every job to release CUDA/Triton contexts. The
|
||||
# next lightweight receiver process restores Ollama before accepting new work.
|
||||
restore_ollama_after_audiobook()
|
||||
|
||||
app = FastAPI(title="Aletheia CMP Audiobook Service", version="1.1.0")
|
||||
queue_condition = threading.Condition()
|
||||
queued_jobs: list[str] = []
|
||||
worker_state = {"activeJobId": None, "modelLoaded": False}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def database() -> Iterator[sqlite3.Connection]:
|
||||
connection = sqlite3.connect(DATABASE, timeout=30)
|
||||
connection.row_factory = sqlite3.Row
|
||||
try:
|
||||
yield connection
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def initialize_database() -> None:
|
||||
with database() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL,
|
||||
author TEXT NOT NULL,
|
||||
source_name TEXT NOT NULL,
|
||||
input_path TEXT NOT NULL,
|
||||
output_path TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
stage TEXT NOT NULL,
|
||||
processed_characters INTEGER NOT NULL DEFAULT 0,
|
||||
total_characters INTEGER NOT NULL DEFAULT 0,
|
||||
current_chapter TEXT,
|
||||
duration_ms INTEGER NOT NULL DEFAULT 0,
|
||||
chapters_json TEXT NOT NULL DEFAULT '[]',
|
||||
error_message TEXT,
|
||||
cancel_requested INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.execute(
|
||||
"UPDATE jobs SET status='queued', stage='queued' "
|
||||
"WHERE status IN ('extracting','synthesizing','encoding')"
|
||||
)
|
||||
rows = connection.execute(
|
||||
"SELECT id FROM jobs WHERE status='queued' ORDER BY created_at"
|
||||
).fetchall()
|
||||
queued_jobs.extend(row["id"] for row in rows)
|
||||
|
||||
|
||||
def require_token(authorization: Annotated[str | None, Header()] = None) -> None:
|
||||
expected = f"Bearer {API_TOKEN}"
|
||||
if not authorization or not hmac.compare_digest(authorization, expected):
|
||||
raise HTTPException(status_code=401, detail="Invalid API token")
|
||||
|
||||
|
||||
def now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def get_job(job_id: str) -> sqlite3.Row:
|
||||
with database() as connection:
|
||||
row = connection.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail="Audiobook job not found")
|
||||
return row
|
||||
|
||||
|
||||
def job_payload(row: sqlite3.Row) -> dict:
|
||||
total = row["total_characters"]
|
||||
processed = row["processed_characters"]
|
||||
progress = min(100, int(processed * 100 / total)) if total else 0
|
||||
if row["status"] == "ready":
|
||||
progress = 100
|
||||
return {
|
||||
"id": row["id"],
|
||||
"status": row["status"],
|
||||
"stage": row["stage"],
|
||||
"progress": progress,
|
||||
"processedCharacters": processed,
|
||||
"totalCharacters": total,
|
||||
"currentChapter": row["current_chapter"],
|
||||
"durationMs": row["duration_ms"],
|
||||
"chapters": json.loads(row["chapters_json"] or "[]"),
|
||||
"error": row["error_message"],
|
||||
"createdAt": row["created_at"],
|
||||
"updatedAt": row["updated_at"],
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> dict:
|
||||
return {
|
||||
"status": "ok",
|
||||
"workerMode": "lazy-model",
|
||||
"activeJobId": worker_state["activeJobId"],
|
||||
"modelLoaded": worker_state["modelLoaded"],
|
||||
}
|
||||
|
||||
|
||||
@app.post("/v1/audiobooks", dependencies=[Depends(require_token)])
|
||||
async def create_audiobook(
|
||||
book: Annotated[UploadFile, File()],
|
||||
title: Annotated[str, Form()],
|
||||
author: Annotated[str, Form()] = "",
|
||||
) -> dict:
|
||||
extension = Path(book.filename or "book.epub").suffix.lower()
|
||||
if extension not in {".epub", ".fb2"}:
|
||||
raise HTTPException(status_code=400, detail="Only EPUB and FB2 are supported")
|
||||
job_id = uuid.uuid4().hex
|
||||
job_root = DATA_ROOT / job_id
|
||||
job_root.mkdir(parents=True)
|
||||
input_path = job_root / f"source{extension}"
|
||||
received = 0
|
||||
try:
|
||||
with input_path.open("wb") as output:
|
||||
while chunk := await book.read(1024 * 1024):
|
||||
received += len(chunk)
|
||||
if received > MAX_UPLOAD_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Book file is too large")
|
||||
output.write(chunk)
|
||||
except Exception:
|
||||
shutil.rmtree(job_root, ignore_errors=True)
|
||||
raise
|
||||
timestamp = now_ms()
|
||||
with database() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs (
|
||||
id,title,author,source_name,input_path,output_path,status,stage,created_at,updated_at
|
||||
) VALUES (?,?,?,?,?,?,?,?,?,?)
|
||||
""",
|
||||
(
|
||||
job_id,
|
||||
title.strip() or Path(book.filename or "Книга").stem,
|
||||
author.strip(),
|
||||
book.filename or input_path.name,
|
||||
str(input_path),
|
||||
str(job_root / "audiobook.m4a"),
|
||||
"queued",
|
||||
"queued",
|
||||
timestamp,
|
||||
timestamp,
|
||||
),
|
||||
)
|
||||
with queue_condition:
|
||||
queued_jobs.append(job_id)
|
||||
queue_condition.notify()
|
||||
return job_payload(get_job(job_id))
|
||||
|
||||
|
||||
@app.get("/v1/audiobooks/{job_id}", dependencies=[Depends(require_token)])
|
||||
def audiobook_status(job_id: str) -> dict:
|
||||
return job_payload(get_job(job_id))
|
||||
|
||||
|
||||
@app.get("/v1/audiobooks/{job_id}/file", dependencies=[Depends(require_token)])
|
||||
def audiobook_file(job_id: str) -> FileResponse:
|
||||
row = get_job(job_id)
|
||||
output = Path(row["output_path"])
|
||||
if row["status"] != "ready" or not output.is_file():
|
||||
raise HTTPException(status_code=409, detail="Audiobook is not ready")
|
||||
safe_title = re.sub(r"[^0-9A-Za-zА-Яа-яЁё._ -]+", "_", row["title"]).strip() or "audiobook"
|
||||
return FileResponse(output, media_type="audio/mp4", filename=f"{safe_title}.m4a")
|
||||
|
||||
|
||||
@app.delete("/v1/audiobooks/{job_id}", dependencies=[Depends(require_token)])
|
||||
def cancel_or_delete(job_id: str) -> dict:
|
||||
row = get_job(job_id)
|
||||
if row["status"] == "queued":
|
||||
with queue_condition:
|
||||
queued_jobs[:] = [queued_id for queued_id in queued_jobs if queued_id != job_id]
|
||||
delete_job(job_id, Path(row["input_path"]).parent)
|
||||
return {"status": "deleted"}
|
||||
if row["status"] in {"extracting", "synthesizing", "encoding"}:
|
||||
with database() as connection:
|
||||
connection.execute(
|
||||
"UPDATE jobs SET cancel_requested=1, stage='cancelling', updated_at=? WHERE id=?",
|
||||
(now_ms(), job_id),
|
||||
)
|
||||
return {"status": "cancelling"}
|
||||
delete_job(job_id, Path(row["input_path"]).parent)
|
||||
return {"status": "deleted"}
|
||||
|
||||
|
||||
def delete_job(job_id: str, job_root: Path) -> None:
|
||||
shutil.rmtree(job_root, ignore_errors=True)
|
||||
with database() as connection:
|
||||
connection.execute("DELETE FROM jobs WHERE id=?", (job_id,))
|
||||
|
||||
|
||||
def finish_cancellation(job_id: str, job_root: Path) -> None:
|
||||
update_job(job_id, status="cancelled", stage="cancelled")
|
||||
delete_job(job_id, job_root)
|
||||
|
||||
|
||||
def local_name(tag: str) -> str:
|
||||
return tag.rsplit("}", 1)[-1].lower()
|
||||
|
||||
|
||||
def clean_text(value: str) -> str:
|
||||
value = html.unescape(re.sub(r"<[^>]+>", " ", value))
|
||||
value = value.replace("\u00a0", " ")
|
||||
return re.sub(r"\s+", " ", value).strip()
|
||||
|
||||
|
||||
def extract_epub(path: Path) -> list[tuple[str, str]]:
|
||||
sections: list[tuple[str, str]] = []
|
||||
with zipfile.ZipFile(path) as archive:
|
||||
container = ElementTree.fromstring(archive.read("META-INF/container.xml"))
|
||||
rootfile = next(node for node in container.iter() if local_name(node.tag) == "rootfile")
|
||||
opf_path = rootfile.attrib["full-path"]
|
||||
opf = ElementTree.fromstring(archive.read(opf_path))
|
||||
manifest = {
|
||||
node.attrib["id"]: node.attrib["href"]
|
||||
for node in opf.iter()
|
||||
if local_name(node.tag) == "item" and "id" in node.attrib and "href" in node.attrib
|
||||
}
|
||||
spine = [
|
||||
node.attrib["idref"]
|
||||
for node in opf.iter()
|
||||
if local_name(node.tag) == "itemref" and "idref" in node.attrib
|
||||
]
|
||||
base = PurePosixPath(opf_path).parent
|
||||
for index, item_id in enumerate(spine, start=1):
|
||||
href = manifest.get(item_id)
|
||||
if not href:
|
||||
continue
|
||||
entry = str(base / href.split("#", 1)[0])
|
||||
raw = archive.read(entry).decode("utf-8", errors="replace")
|
||||
heading_match = re.search(r"<h[1-3]\b[^>]*>(.*?)</h[1-3]>", raw, re.I | re.S)
|
||||
title = clean_text(heading_match.group(1)) if heading_match else f"Раздел {index}"
|
||||
paragraphs = [clean_text(match) for match in re.findall(r"<p\b[^>]*>(.*?)</p>", raw, re.I | re.S)]
|
||||
text = "\n".join(paragraph for paragraph in paragraphs if paragraph)
|
||||
if text:
|
||||
sections.append((title, text))
|
||||
return sections
|
||||
|
||||
|
||||
def element_text_without_nested_sections(element: ElementTree.Element) -> str:
|
||||
parts: list[str] = []
|
||||
if element.text:
|
||||
parts.append(element.text)
|
||||
for child in element:
|
||||
if local_name(child.tag) != "section":
|
||||
parts.append(" ".join(child.itertext()))
|
||||
if child.tail:
|
||||
parts.append(child.tail)
|
||||
return clean_text(" ".join(parts))
|
||||
|
||||
|
||||
def extract_fb2(path: Path) -> list[tuple[str, str]]:
|
||||
root = ElementTree.parse(path).getroot()
|
||||
sections: list[tuple[str, str]] = []
|
||||
for section in root.iter():
|
||||
if local_name(section.tag) != "section":
|
||||
continue
|
||||
title_element = next((child for child in section if local_name(child.tag) == "title"), None)
|
||||
title = clean_text(" ".join(title_element.itertext())) if title_element is not None else ""
|
||||
text = element_text_without_nested_sections(section)
|
||||
if text:
|
||||
sections.append((title or f"Раздел {len(sections) + 1}", text))
|
||||
return sections
|
||||
|
||||
|
||||
def normalize_numbers(text: str) -> str:
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
value = match.group(0)
|
||||
try:
|
||||
return num2words(int(value), lang="ru")
|
||||
except (ValueError, OverflowError):
|
||||
return value
|
||||
|
||||
return re.sub(r"(?<![\w+])\d{1,12}(?!\w)", replace, text)
|
||||
|
||||
|
||||
def split_chunks(sections: list[tuple[str, str]], limit: int = 520) -> list[dict]:
|
||||
chunks: list[dict] = []
|
||||
for chapter, text in sections:
|
||||
sentences = re.split(r"(?<=[.!?…])\s+|\n+", text)
|
||||
current = ""
|
||||
for sentence in sentences:
|
||||
sentence = normalize_numbers(sentence.strip())
|
||||
if not sentence or not re.search(r"[А-Яа-яЁё]", sentence):
|
||||
continue
|
||||
if len(sentence) > limit:
|
||||
pieces = re.split(r"(?<=[,;:])\s+", sentence)
|
||||
else:
|
||||
pieces = [sentence]
|
||||
for piece in pieces:
|
||||
if current and len(current) + 1 + len(piece) > limit:
|
||||
chunks.append({"chapter": chapter, "text": current})
|
||||
current = piece
|
||||
else:
|
||||
current = f"{current} {piece}".strip()
|
||||
if current:
|
||||
chunks.append({"chapter": chapter, "text": current})
|
||||
return chunks
|
||||
|
||||
|
||||
def update_job(job_id: str, **values: object) -> None:
|
||||
values["updated_at"] = now_ms()
|
||||
assignments = ",".join(f"{key}=?" for key in values)
|
||||
with database() as connection:
|
||||
connection.execute(
|
||||
f"UPDATE jobs SET {assignments} WHERE id=?",
|
||||
(*values.values(), job_id),
|
||||
)
|
||||
|
||||
|
||||
def cancellation_requested(job_id: str) -> bool:
|
||||
with database() as connection:
|
||||
row = connection.execute("SELECT cancel_requested FROM jobs WHERE id=?", (job_id,)).fetchone()
|
||||
return row is None or bool(row["cancel_requested"])
|
||||
|
||||
|
||||
def load_model(device_id: int) -> "Qwen3TTSModel":
|
||||
import torch
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA is unavailable")
|
||||
torch.set_float32_matmul_precision("high")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
MODEL_ID,
|
||||
device_map=f"cuda:{device_id}",
|
||||
dtype=torch.bfloat16,
|
||||
attn_implementation="sdpa",
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def load_models() -> list["Qwen3TTSModel"]:
|
||||
if not GPU_IDS:
|
||||
raise RuntimeError("No Qwen GPU IDs are configured")
|
||||
models = [load_model(device_id) for device_id in GPU_IDS]
|
||||
worker_state["modelLoaded"] = True
|
||||
return models
|
||||
|
||||
|
||||
def generate_custom_voice_with_metrics(
|
||||
model: Any,
|
||||
batch_chunks: list[dict],
|
||||
device_id: int,
|
||||
) -> tuple[list[Any], int]:
|
||||
"""Measure Qwen talker and codec stages without changing generated audio."""
|
||||
import torch
|
||||
|
||||
timings: dict[str, float] = {}
|
||||
original_generate = model.model.generate
|
||||
speech_tokenizer = model.model.speech_tokenizer
|
||||
original_decode = speech_tokenizer.decode
|
||||
|
||||
def synchronize() -> None:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize(device_id)
|
||||
|
||||
def timed_generate(*args: Any, **kwargs: Any) -> Any:
|
||||
synchronize()
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
return original_generate(*args, **kwargs)
|
||||
finally:
|
||||
synchronize()
|
||||
timings["talkerGenerateSeconds"] = time.perf_counter() - started
|
||||
|
||||
def timed_decode(*args: Any, **kwargs: Any) -> Any:
|
||||
synchronize()
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
return original_decode(*args, **kwargs)
|
||||
finally:
|
||||
synchronize()
|
||||
timings["codecDecodeSeconds"] = time.perf_counter() - started
|
||||
|
||||
model.model.generate = timed_generate
|
||||
speech_tokenizer.decode = timed_decode
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
wavs, sample_rate = model.generate_custom_voice(
|
||||
text=[chunk["text"] for chunk in batch_chunks],
|
||||
language=["Russian"] * len(batch_chunks),
|
||||
speaker=[SPEAKER] * len(batch_chunks),
|
||||
instruct=[VOICE_INSTRUCTION] * len(batch_chunks),
|
||||
)
|
||||
synchronize()
|
||||
finally:
|
||||
model.model.generate = original_generate
|
||||
speech_tokenizer.decode = original_decode
|
||||
|
||||
total_seconds = time.perf_counter() - started
|
||||
talker_seconds = timings.get("talkerGenerateSeconds", 0.0)
|
||||
codec_seconds = timings.get("codecDecodeSeconds", 0.0)
|
||||
record = {
|
||||
"timestampMs": now_ms(),
|
||||
"batchSize": len(batch_chunks),
|
||||
"characters": sum(len(chunk["text"]) for chunk in batch_chunks),
|
||||
"totalSeconds": round(total_seconds, 4),
|
||||
"talkerGenerateSeconds": round(talker_seconds, 4),
|
||||
"codecDecodeSeconds": round(codec_seconds, 4),
|
||||
"wrapperSeconds": round(max(0.0, total_seconds - talker_seconds - codec_seconds), 4),
|
||||
"cudaPeakAllocatedMiB": round(
|
||||
torch.cuda.max_memory_allocated(device_id) / (1024 * 1024), 1
|
||||
),
|
||||
}
|
||||
with PERFORMANCE_LOG.open("a", encoding="utf-8") as output:
|
||||
output.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
return wavs, sample_rate
|
||||
|
||||
|
||||
def unload_models(models: list[Any]) -> None:
|
||||
if models:
|
||||
models.clear()
|
||||
import torch
|
||||
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
for device_id in GPU_IDS:
|
||||
torch.cuda.empty_cache()
|
||||
else:
|
||||
gc.collect()
|
||||
worker_state["modelLoaded"] = False
|
||||
|
||||
|
||||
def encode_m4a(chunk_files: list[Path], output: Path, job_root: Path) -> None:
|
||||
concat = job_root / "concat.txt"
|
||||
concat.write_text(
|
||||
"\n".join(f"file '{path.as_posix()}'" for path in chunk_files),
|
||||
encoding="utf-8",
|
||||
)
|
||||
command = [
|
||||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||||
"-y",
|
||||
"-f",
|
||||
"concat",
|
||||
"-safe",
|
||||
"0",
|
||||
"-i",
|
||||
str(concat),
|
||||
"-vn",
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-b:a",
|
||||
"96k",
|
||||
"-movflags",
|
||||
"+faststart",
|
||||
str(output),
|
||||
]
|
||||
completed = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", errors="replace")
|
||||
if completed.returncode != 0:
|
||||
raise RuntimeError(completed.stderr[-2000:] or "FFmpeg failed")
|
||||
|
||||
|
||||
def run_job(job_id: str) -> None:
|
||||
row = get_job(job_id)
|
||||
job_root = Path(row["input_path"]).parent
|
||||
chunks_root = job_root / "chunks"
|
||||
chunks_root.mkdir(exist_ok=True)
|
||||
models: list[Any] = []
|
||||
try:
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
update_job(job_id, status="extracting", stage="extracting", error_message=None)
|
||||
input_path = Path(row["input_path"])
|
||||
sections = extract_epub(input_path) if input_path.suffix.lower() == ".epub" else extract_fb2(input_path)
|
||||
chunks = split_chunks(sections)
|
||||
if not chunks:
|
||||
raise RuntimeError("В книге не найден русский текст для озвучивания")
|
||||
total = sum(len(chunk["text"]) for chunk in chunks)
|
||||
manifest_path = job_root / "chunks.json"
|
||||
manifest_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
update_job(job_id, status="synthesizing", stage="loading_model", total_characters=total)
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
pause_ollama_for_audiobook()
|
||||
models = load_models()
|
||||
import torch
|
||||
|
||||
processed = 0
|
||||
duration_ms = 0
|
||||
chapter_markers: list[dict] = []
|
||||
last_chapter: str | None = None
|
||||
chunk_files: list[Path] = []
|
||||
index = 0
|
||||
while index < len(chunks):
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
output = chunks_root / f"{index:06d}.wav"
|
||||
completed_indices: list[int]
|
||||
if output.is_file():
|
||||
completed_indices = [index]
|
||||
else:
|
||||
batch_indices: list[int] = []
|
||||
cursor = index
|
||||
while cursor < len(chunks) and len(batch_indices) < len(models):
|
||||
candidate = chunks_root / f"{cursor:06d}.wav"
|
||||
if candidate.is_file():
|
||||
break
|
||||
batch_indices.append(cursor)
|
||||
cursor += 1
|
||||
with ThreadPoolExecutor(max_workers=len(batch_indices)) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
generate_custom_voice_with_metrics,
|
||||
model,
|
||||
[chunks[batch_index]],
|
||||
device_id,
|
||||
)
|
||||
for batch_index, model, device_id in zip(
|
||||
batch_indices,
|
||||
models[: len(batch_indices)],
|
||||
GPU_IDS[: len(batch_indices)],
|
||||
strict=True,
|
||||
)
|
||||
]
|
||||
rendered = [future.result() for future in futures]
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
for batch_index, (wavs, sample_rate) in zip(batch_indices, rendered, strict=True):
|
||||
sf.write(
|
||||
chunks_root / f"{batch_index:06d}.wav",
|
||||
wavs[0],
|
||||
sample_rate,
|
||||
subtype="PCM_16",
|
||||
)
|
||||
completed_indices = batch_indices
|
||||
|
||||
for completed_index in completed_indices:
|
||||
completed_chunk = chunks[completed_index]
|
||||
completed_output = chunks_root / f"{completed_index:06d}.wav"
|
||||
if completed_chunk["chapter"] != last_chapter:
|
||||
chapter_markers.append({"title": completed_chunk["chapter"], "startMs": duration_ms})
|
||||
last_chapter = completed_chunk["chapter"]
|
||||
info = sf.info(completed_output)
|
||||
duration_ms += round(info.frames * 1000 / info.samplerate)
|
||||
chunk_files.append(completed_output)
|
||||
processed += len(completed_chunk["text"])
|
||||
update_job(
|
||||
job_id,
|
||||
status="synthesizing",
|
||||
stage="synthesizing",
|
||||
processed_characters=processed,
|
||||
current_chapter=completed_chunk["chapter"],
|
||||
duration_ms=duration_ms,
|
||||
chapters_json=json.dumps(chapter_markers, ensure_ascii=False),
|
||||
)
|
||||
index = completed_indices[-1] + 1
|
||||
|
||||
unload_models(models)
|
||||
models = []
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
update_job(job_id, status="encoding", stage="encoding", current_chapter=None)
|
||||
output = Path(row["output_path"])
|
||||
encode_m4a(chunk_files, output, job_root)
|
||||
if cancellation_requested(job_id):
|
||||
finish_cancellation(job_id, job_root)
|
||||
return
|
||||
if not output.is_file() or output.stat().st_size == 0:
|
||||
raise RuntimeError("Итоговый M4A-файл не создан")
|
||||
update_job(
|
||||
job_id,
|
||||
status="ready",
|
||||
stage="ready",
|
||||
processed_characters=total,
|
||||
current_chapter=None,
|
||||
chapters_json=json.dumps(chapter_markers, ensure_ascii=False),
|
||||
)
|
||||
except Exception as error:
|
||||
update_job(
|
||||
job_id,
|
||||
status="failed",
|
||||
stage="failed",
|
||||
error_message=f"{type(error).__name__}: {error}",
|
||||
)
|
||||
finally:
|
||||
unload_models(models)
|
||||
|
||||
|
||||
def worker_loop() -> None:
|
||||
while True:
|
||||
with queue_condition:
|
||||
while not queued_jobs:
|
||||
queue_condition.wait()
|
||||
job_id = queued_jobs.pop(0)
|
||||
worker_state["activeJobId"] = job_id
|
||||
try:
|
||||
try:
|
||||
run_job(job_id)
|
||||
except HTTPException as error:
|
||||
if error.status_code != 404:
|
||||
raise
|
||||
finally:
|
||||
worker_state["activeJobId"] = None
|
||||
# Qwen/Triton can retain a CUDA context after Python references are
|
||||
# released. Let systemd start a clean idle receiver so CMP VRAM is
|
||||
# fully available between audiobook requests. Unfinished jobs are
|
||||
# restored from SQLite by initialize_database() after restart.
|
||||
os._exit(0)
|
||||
|
||||
|
||||
initialize_database()
|
||||
threading.Thread(target=worker_loop, name="audiobook-worker", daemon=True).start()
|
||||
Reference in New Issue
Block a user