689 lines
25 KiB
Python
689 lines
25 KiB
Python
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()
|