claudeandClaude Opus 5
3dde6e58e4
Poll interval wait_for_completion() polled every 3 s, which quantised every clip to a 3 s boundary — across 2075 existing clips, 91% of the gaps landed exactly on a multiple of 3 s. It now starts at 0.25 s and backs off to 2 s after 10 s, so short clips are picked up immediately without hammering the server on long ones. Measured over 34 clips: 4.75 s actual vs 6.35 s under the old poll, a 25% cut in TTS wait (~18 min on a 693-clip episode). Runs now print real generation time, e.g. "done (4.2s)". A WebSocket would remove polling entirely but needs a third-party client, which would cost this script its dependency-free property. settings.ini voice, speed, temperature, max_tokens, strip_leading_ms, trailing_fade_ms, orig_volume, tts_volume, output_dir and comfy_url moved to an INI file next to the script. Precedence is CLI flag > settings.ini > built-in default, wired by feeding loaded values in as argparse defaults. --settings FILE selects an alternative profile. INI rather than TOML (tomllib is 3.11+, this runs on system Python 3.10) or JSON (no comments). A missing, partial, malformed or corrupt file degrades to built-in defaults with a warning instead of killing a long run. Run-specific flags are deliberately excluded: --start, --end, --seed, --regen, --rebuild, --resume-from, --embedded-subs, --skip-tts. Per-job folders All state for one video now lives in <output_dir>/clips/<video name>/ — clips, tts_input.srt, tts_assembled.flac, embedded_sub.srt and video_fragment.mkv. Only <stem>_lektor.mkv stays at the top level. tts_input.srt moved with the clips because they are a matched pair: block N in the SRT is clip_NNNN.flac. Previously both sat at the top level and each new job overwrote the last one's, so rebuilding an older film silently paired its clips with the newest film's text. Any past job can now be rebuilt. A run against the old flat layout gets an explanatory message. gudowski/ -> voices/ VOICE_PROMPTS_DIR is now resolved relative to the script rather than hardcoded, so moving the project no longer breaks it. Voice filenames are unchanged, so --voice gudowski_short still works. Docs updated throughout for the new layout, settings file and poll behaviour. Known issue, not addressed here: with a non-zero --start the voiceover plays early. extract_video_fragment() stream-copies, so it can only cut on a keyframe and seeks backwards to it, while TTS timestamps are relative to the requested start. On S01_E08 at --start 07:40 the nearest earlier keyframe is 458.375 s, putting the voiceover 1.625 s ahead of the picture. Runs with --start 00:00 are unaffected, since 0 is always a keyframe. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
885 lines
35 KiB
Python
885 lines
35 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
lektor_ai_comfy_v2 — SRT-to-voiceover via ComfyUI + Higgs Audio V3.
|
||
Sends subtitles in batches (default: 5) so the TTS model sees multi-sentence
|
||
context, which can improve prosody and naturalness compared to per-clip generation.
|
||
|
||
Usage:
|
||
python3 lektor_ai_comfy_v2.py input/film.srt input/film.mkv --start 00:00 --end 01:40
|
||
|
||
Requires ComfyUI with TTS-Audio-Suite running at http://127.0.0.1:8189
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
from urllib.request import urlopen, Request
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Config
|
||
# ---------------------------------------------------------------------------
|
||
|
||
COMFY_URL = "http://127.0.0.1:8189"
|
||
COMFY_OUTPUT_DIR = Path("/home/lukas/comfy-tts-app/ComfyUI/output")
|
||
COMFY_INPUT_DIR = Path("/home/lukas/comfy-tts-app/ComfyUI/input")
|
||
VOICE_PROMPTS_DIR = Path("/home/lukas/Programy/lektor_ai/voices")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SRT helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_SRT_TIME_RE = re.compile(r'(\d+):(\d+):(\d+)[,.](\d+)')
|
||
|
||
|
||
def _parse_srt_time(s: str) -> float:
|
||
m = _SRT_TIME_RE.match(s.strip())
|
||
if not m:
|
||
raise ValueError(f"Bad SRT timestamp: {s!r}")
|
||
h, mn, sec, ms = int(m[1]), int(m[2]), int(m[3]), int(m[4])
|
||
return h * 3600 + mn * 60 + sec + ms / 1000
|
||
|
||
|
||
def _fmt_srt_time(seconds: float) -> str:
|
||
h = int(seconds // 3600)
|
||
m = int((seconds % 3600) // 60)
|
||
s = int(seconds % 60)
|
||
ms = int(round((seconds % 1) * 1000))
|
||
if ms == 1000:
|
||
ms = 999
|
||
return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
|
||
|
||
|
||
def _parse_mmss(value: str) -> float:
|
||
parts = value.strip().split(":")
|
||
if len(parts) == 2:
|
||
return int(parts[0]) * 60 + float(parts[1])
|
||
if len(parts) == 3:
|
||
return int(parts[0]) * 3600 + int(parts[1]) * 60 + float(parts[2])
|
||
raise ValueError(f"Cannot parse time: {value!r}")
|
||
|
||
|
||
def _ensure_terminal_punct(text: str, single_dot: bool = False) -> str:
|
||
if not text or text[-1] in ('.', ',', '!', '?', ';', ':'):
|
||
return text
|
||
return text + ('.' if single_dot else '...')
|
||
|
||
|
||
def _clean_text(lines: list) -> str:
|
||
has_dialogue = any(re.match(r"^-\s*\S", l.strip()) for l in lines if l.strip())
|
||
out = []
|
||
for line in lines:
|
||
line = line.strip()
|
||
line = re.sub(r"^-\s*", "", line)
|
||
if line:
|
||
out.append(line)
|
||
result = "\n".join(out) if has_dialogue else " ".join(out)
|
||
letters = [c for c in result if c.isalpha()]
|
||
was_allcaps = bool(letters) and all(c == c.upper() for c in letters)
|
||
if was_allcaps:
|
||
result = result.lower()
|
||
if result and result[0].islower():
|
||
result = result[0].upper() + result[1:]
|
||
if has_dialogue:
|
||
result = "\n".join(_ensure_terminal_punct(p, was_allcaps) for p in result.split("\n"))
|
||
else:
|
||
result = _ensure_terminal_punct(result, was_allcaps)
|
||
return result
|
||
|
||
|
||
def load_tts_input_srt(path: Path) -> list:
|
||
"""Read a pre-processed tts_input.srt into (t_start, t_end, text) tuples."""
|
||
text = path.read_text(encoding="utf-8")
|
||
blocks = re.split(r"\n{2,}", text.strip())
|
||
result = []
|
||
for block in blocks:
|
||
lines = block.strip().splitlines()
|
||
ti = next((i for i, l in enumerate(lines) if "-->" in l), None)
|
||
if ti is None:
|
||
continue
|
||
parts = lines[ti].split("-->")
|
||
if len(parts) != 2:
|
||
continue
|
||
t_start = _parse_srt_time(parts[0])
|
||
t_end = _parse_srt_time(parts[1])
|
||
payload = " ".join(l.strip() for l in lines[ti + 1:] if l.strip())
|
||
if payload:
|
||
result.append((t_start, t_end, payload))
|
||
return result
|
||
|
||
|
||
def parse_subtitle_blocks(srt_path: str, start_sec: float, end_sec: float) -> list:
|
||
"""Return list of (start_relative, end_relative, plain_text) blocks."""
|
||
text = Path(srt_path).read_text(encoding="utf-8-sig")
|
||
blocks = re.split(r"\n{2,}", text.strip())
|
||
result = []
|
||
for block in blocks:
|
||
lines = block.strip().splitlines()
|
||
ti = next((i for i, l in enumerate(lines) if "-->" in l), None)
|
||
if ti is None:
|
||
continue
|
||
parts = lines[ti].split("-->")
|
||
if len(parts) != 2:
|
||
continue
|
||
t_start = _parse_srt_time(parts[0])
|
||
t_end = _parse_srt_time(parts[1])
|
||
if t_start < start_sec or t_start > end_sec:
|
||
continue
|
||
plain = _clean_text(lines[ti + 1:])
|
||
if plain:
|
||
result.append((t_start - start_sec, t_end - start_sec, plain))
|
||
return result
|
||
|
||
|
||
def split_dialogue_blocks(blocks: list) -> list:
|
||
"""Split multi-line blocks into one entry per line with proportional timing."""
|
||
result = []
|
||
for t_start, t_end, text in blocks:
|
||
lines = [l for l in text.split("\n") if l.strip()]
|
||
if len(lines) <= 1:
|
||
result.append((t_start, t_end, text))
|
||
continue
|
||
duration = t_end - t_start
|
||
total_chars = sum(len(l) for l in lines)
|
||
current = t_start
|
||
for i, line in enumerate(lines):
|
||
share = len(line) / total_chars if total_chars else 1 / len(lines)
|
||
end = t_end if i == len(lines) - 1 else current + duration * share
|
||
result.append((current, end, line))
|
||
current = end
|
||
return result
|
||
|
||
|
||
def blocks_to_srt(blocks: list) -> str:
|
||
"""Convert (start, end, text) tuples back to SRT string for inspection."""
|
||
lines = []
|
||
for i, (t_start, t_end, text) in enumerate(blocks, 1):
|
||
lines.append(str(i))
|
||
lines.append(f"{_fmt_srt_time(t_start)} --> {_fmt_srt_time(t_end)}")
|
||
lines.append(text)
|
||
lines.append("")
|
||
return "\n".join(lines)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FFmpeg helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _ffmpeg_ts(seconds: float) -> str:
|
||
h = int(seconds // 3600)
|
||
m = int((seconds % 3600) // 60)
|
||
s = seconds % 60
|
||
return f"{h:02d}:{m:02d}:{s:06.3f}"
|
||
|
||
|
||
def _run(cmd: list, label: str):
|
||
print(f" {label}...")
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
if result.returncode != 0:
|
||
print(f"ffmpeg error ({label}):\n{result.stderr[-2000:]}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
|
||
|
||
def get_audio_duration(path: Path) -> float:
|
||
result = subprocess.run(
|
||
["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", str(path)],
|
||
capture_output=True, text=True,
|
||
)
|
||
info = json.loads(result.stdout)
|
||
for stream in info.get("streams", []):
|
||
if "duration" in stream:
|
||
return float(stream["duration"])
|
||
raise RuntimeError(f"Cannot determine duration of {path}")
|
||
|
||
|
||
def speed_up_audio(input_path: Path, output_path: Path, speed: float,
|
||
fade_out_ms: float = 0):
|
||
"""Apply atempo (and optional fade-out) to a single clip."""
|
||
filters = []
|
||
if speed != 1.0:
|
||
remaining = speed
|
||
while remaining > 2.0:
|
||
filters.append("atempo=2.0")
|
||
remaining /= 2.0
|
||
filters.append(f"atempo={remaining:.4f}")
|
||
if fade_out_ms > 0:
|
||
fade_sec = fade_out_ms / 1000
|
||
raw_dur = get_audio_duration(input_path)
|
||
post_dur = raw_dur / speed if speed != 1.0 else raw_dur
|
||
fade_start = max(0.0, post_dur - fade_sec)
|
||
filters.append(f"afade=t=out:st={fade_start:.3f}:d={fade_sec:.3f}")
|
||
if not filters:
|
||
shutil.copy2(input_path, output_path)
|
||
return
|
||
label = f"Speed {speed:.2f}×"
|
||
if fade_out_ms:
|
||
label += f" + fade-out {fade_out_ms:.0f}ms"
|
||
_run([
|
||
"ffmpeg", "-y",
|
||
"-i", str(input_path),
|
||
"-filter:a", ",".join(filters),
|
||
str(output_path),
|
||
], label)
|
||
|
||
|
||
def strip_leading_audio(input_path: Path, output_path: Path, strip_ms: float):
|
||
"""Remove the first strip_ms milliseconds from a clip."""
|
||
_run([
|
||
"ffmpeg", "-y",
|
||
"-i", str(input_path),
|
||
"-ss", f"{strip_ms / 1000:.3f}",
|
||
"-c", "copy",
|
||
str(output_path),
|
||
], f"Strip {strip_ms:.0f}ms leading audio")
|
||
|
||
|
||
def _assemble_batch(clips: list, output_path: str, label: str):
|
||
"""Mix a batch of (start_sec, path) clips with absolute adelay into output_path."""
|
||
if len(clips) == 1:
|
||
delay_ms = int(clips[0][0] * 1000)
|
||
_run([
|
||
"ffmpeg", "-y",
|
||
"-i", str(clips[0][1]),
|
||
"-af", f"adelay={delay_ms}:all=1",
|
||
output_path,
|
||
], label)
|
||
return
|
||
inputs, filter_parts = [], []
|
||
for i, (start_sec, path) in enumerate(clips):
|
||
inputs += ["-i", str(path)]
|
||
delay_ms = int(start_sec * 1000)
|
||
filter_parts.append(f"[{i}:a]adelay={delay_ms}:all=1[a{i}]")
|
||
mix_inputs = "".join(f"[a{i}]" for i in range(len(clips)))
|
||
filter_parts.append(
|
||
f"{mix_inputs}amix=inputs={len(clips)}:duration=longest:normalize=0[out]"
|
||
)
|
||
_run(
|
||
["ffmpeg", "-y"] + inputs +
|
||
["-filter_complex", ";".join(filter_parts), "-map", "[out]", output_path],
|
||
label,
|
||
)
|
||
|
||
|
||
def assemble_timed_audio(clips: list, output_path: str, batch_size: int = 500):
|
||
"""Place each clip at its start time via adelay then mix.
|
||
clips: list of (start_sec, Path).
|
||
Processes in batches to stay under OS file-descriptor limits.
|
||
"""
|
||
if len(clips) <= batch_size:
|
||
_assemble_batch(clips, output_path, f"Assembling {len(clips)} timed clips")
|
||
return
|
||
|
||
tmp_dir = Path(output_path).parent / "_batch_tmp"
|
||
tmp_dir.mkdir(exist_ok=True)
|
||
batch_files = []
|
||
try:
|
||
for b, start in enumerate(range(0, len(clips), batch_size)):
|
||
batch = clips[start:start + batch_size]
|
||
batch_path = str(tmp_dir / f"batch_{b:04d}.flac")
|
||
_assemble_batch(
|
||
batch, batch_path,
|
||
f"Assembling batch {b + 1}/{(len(clips) - 1) // batch_size + 1} ({len(batch)} clips)",
|
||
)
|
||
batch_files.append(batch_path)
|
||
batch_clips = [(0.0, Path(p)) for p in batch_files]
|
||
_assemble_batch(batch_clips, output_path, f"Merging {len(batch_files)} batches")
|
||
finally:
|
||
for p in batch_files:
|
||
Path(p).unlink(missing_ok=True)
|
||
tmp_dir.rmdir() if tmp_dir.exists() and not any(tmp_dir.iterdir()) else None
|
||
|
||
|
||
def extract_video_fragment(video_path: str, start_sec: float, end_sec: float,
|
||
output_path: str):
|
||
_run([
|
||
"ffmpeg", "-y",
|
||
"-ss", _ffmpeg_ts(start_sec),
|
||
"-i", video_path,
|
||
"-t", str(end_sec - start_sec),
|
||
"-c", "copy",
|
||
output_path,
|
||
], f"Extracting video {_ffmpeg_ts(start_sec)}–{_ffmpeg_ts(end_sec)}")
|
||
|
||
|
||
def merge_voiceover(video_path: str, tts_path: str, output_path: str,
|
||
orig_volume: float = 0.60, tts_volume: float = 1.8):
|
||
"""Produce a video with two audio tracks:
|
||
track 0 — Original (first audio stream from source, re-encoded to AAC)
|
||
track 1 — Polish (lektor mix: TTS at tts_volume + original at orig_volume)
|
||
"""
|
||
probe = subprocess.run(
|
||
["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", video_path],
|
||
capture_output=True, text=True,
|
||
)
|
||
streams = json.loads(probe.stdout).get("streams", [])
|
||
audio_start = next(
|
||
(float(s.get("start_time", 0)) for s in streams if s.get("codec_type") == "audio"),
|
||
0.0,
|
||
)
|
||
pts_reset = ",asetpts=PTS-STARTPTS" if audio_start > 0.01 else ""
|
||
if audio_start > 0.01:
|
||
print(f" Video audio start offset: {audio_start:.3f}s — resetting PTS in lektor mix")
|
||
|
||
_run([
|
||
"ffmpeg", "-y",
|
||
"-i", video_path,
|
||
"-i", tts_path,
|
||
"-filter_complex",
|
||
(
|
||
f"[0:a:0]volume={orig_volume}{pts_reset}[orig_mix];"
|
||
f"[1:a]volume={tts_volume}[tts];"
|
||
"[orig_mix][tts]amix=inputs=2:duration=longest:normalize=0[lektor]"
|
||
),
|
||
"-map", "0:v",
|
||
"-map", "0:a:0",
|
||
"-map", "[lektor]",
|
||
"-c:v", "copy",
|
||
"-c:a:0", "copy",
|
||
"-c:a:1", "aac", "-b:a:1", "192k",
|
||
"-metadata:s:a:0", "title=Original",
|
||
"-metadata:s:a:1", "title=Polish",
|
||
"-metadata:s:a:1", "language=pol",
|
||
output_path,
|
||
], "Merging voiceover with video (2 audio tracks)")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI API helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _comfy_post(path: str, data: dict) -> dict:
|
||
payload = json.dumps(data).encode()
|
||
req = Request(
|
||
f"{COMFY_URL}{path}",
|
||
data=payload,
|
||
headers={"Content-Type": "application/json"},
|
||
)
|
||
with urlopen(req, timeout=30) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def _comfy_get(path: str) -> dict:
|
||
with urlopen(f"{COMFY_URL}{path}", timeout=30) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def upload_voice_wav(wav_path: Path) -> tuple:
|
||
"""Copy voice WAV into ComfyUI's input dir; return (filename, reference_text)."""
|
||
dst = COMFY_INPUT_DIR / wav_path.name
|
||
if not dst.exists() or dst.stat().st_mtime < wav_path.stat().st_mtime:
|
||
shutil.copy2(wav_path, dst)
|
||
print(f" Copied {wav_path.name} → {dst}")
|
||
txt_path = wav_path.with_suffix(".txt")
|
||
ref_text = txt_path.read_text(encoding="utf-8").strip() if txt_path.exists() else ""
|
||
if ref_text:
|
||
print(f" Reference text loaded ({len(ref_text)} chars)")
|
||
else:
|
||
print(" WARNING: no reference .txt found — voice cloning quality will be lower")
|
||
return wav_path.name, ref_text
|
||
|
||
|
||
def build_text_workflow(text: str, voice_filename: str, voice_ref_text: str,
|
||
seed: int, temperature: float = 0.9,
|
||
max_new_tokens: int = 4096) -> dict:
|
||
"""ComfyUI workflow for a single text utterance with proper voice reference."""
|
||
return {
|
||
"1": {
|
||
"class_type": "HiggsAudioV3EngineNode",
|
||
"inputs": {
|
||
"model": "higgs-audio-v3-tts-4b",
|
||
"device": "auto",
|
||
"dtype": "auto",
|
||
"attention": "auto",
|
||
"temperature": temperature,
|
||
"top_p": 0.95,
|
||
"top_k": 50,
|
||
"max_new_tokens": max_new_tokens,
|
||
},
|
||
},
|
||
"2": {
|
||
"class_type": "LoadAudio",
|
||
"inputs": {"audio": voice_filename},
|
||
},
|
||
"3": {
|
||
"class_type": "CharacterVoicesNode",
|
||
"inputs": {
|
||
"voice_name": "none",
|
||
"reference_text": voice_ref_text,
|
||
"opt_audio_input": ["2", 0],
|
||
},
|
||
},
|
||
"4": {
|
||
"class_type": "UnifiedTTSTextNode",
|
||
"inputs": {
|
||
"TTS_engine": ["1", 0],
|
||
"text": text,
|
||
"narrator_voice": "none",
|
||
"seed": seed,
|
||
"opt_narrator": ["3", 0],
|
||
"enable_audio_cache": False,
|
||
},
|
||
},
|
||
"5": {
|
||
"class_type": "SaveAudio",
|
||
"inputs": {
|
||
"audio": ["4", 0],
|
||
"filename_prefix": "audio/lektor_clip",
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
def submit_workflow(workflow: dict) -> str:
|
||
client_id = str(uuid.uuid4())
|
||
resp = _comfy_post("/prompt", {"prompt": workflow, "client_id": client_id})
|
||
return resp["prompt_id"]
|
||
|
||
|
||
def wait_for_completion(prompt_id: str, poll_interval: int = 3) -> dict:
|
||
print(f" Waiting for {prompt_id[:8]}…", end="", flush=True)
|
||
while True:
|
||
try:
|
||
history = _comfy_get(f"/history/{prompt_id}")
|
||
except Exception as e:
|
||
print(f"\n Poll error: {e}", file=sys.stderr)
|
||
time.sleep(poll_interval)
|
||
continue
|
||
|
||
if prompt_id in history:
|
||
entry = history[prompt_id]
|
||
status = entry.get("status", {})
|
||
for msg_type, msg_data in status.get("messages", []):
|
||
if msg_type == "execution_error":
|
||
print(f"\nComfyUI error: {msg_data}", file=sys.stderr)
|
||
sys.exit(1)
|
||
if status.get("completed"):
|
||
print(" done.")
|
||
return entry
|
||
|
||
print(".", end="", flush=True)
|
||
time.sleep(poll_interval)
|
||
|
||
|
||
def extract_output_audio(history_entry: dict) -> Path:
|
||
outputs = history_entry.get("outputs", {})
|
||
for node_output in outputs.values():
|
||
for audio_info in node_output.get("audio", []):
|
||
subfolder = audio_info.get("subfolder", "")
|
||
path = COMFY_OUTPUT_DIR / subfolder / audio_info["filename"]
|
||
if path.exists():
|
||
return path
|
||
raise RuntimeError(f"Output audio not found in history: {outputs}")
|
||
|
||
|
||
def _generate_one(text: str, voice_filename: str, voice_ref_text: str,
|
||
seed: int, speed: float, out_dir: Path, label: str,
|
||
strip_leading_ms: float, temperature: float = 0.9,
|
||
trailing_fade_ms: float = 0, max_new_tokens: int = 4096) -> tuple:
|
||
"""Submit one TTS job and return (path, duration_after_speedup)."""
|
||
workflow = build_text_workflow(text, voice_filename, voice_ref_text, seed,
|
||
temperature, max_new_tokens)
|
||
prompt_id = submit_workflow(workflow)
|
||
history_entry = wait_for_completion(prompt_id)
|
||
raw_clip = extract_output_audio(history_entry)
|
||
|
||
if strip_leading_ms > 0:
|
||
stripped = out_dir / f"{label}_s.flac"
|
||
strip_leading_audio(raw_clip, stripped, strip_leading_ms)
|
||
raw_clip = stripped
|
||
|
||
fast_clip = out_dir / f"{label}.flac"
|
||
speed_up_audio(raw_clip, fast_clip, speed, fade_out_ms=trailing_fade_ms)
|
||
|
||
return fast_clip, get_audio_duration(fast_clip)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Batch generation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _group_into_batches(blocks: list, batch_size: int) -> list:
|
||
"""Group (t_start, t_end, text) blocks into batches of batch_size.
|
||
|
||
Returns list of dicts:
|
||
batch_num — 1-based batch index
|
||
t_start — SRT start of the first block in the batch
|
||
t_end — SRT end of the last block in the batch
|
||
block_indices — 1-based indices of the blocks that belong to this batch
|
||
texts — list of individual block texts (already have terminal punct)
|
||
"""
|
||
batches = []
|
||
for i in range(0, len(blocks), batch_size):
|
||
group = blocks[i:i + batch_size]
|
||
batches.append({
|
||
"batch_num": len(batches) + 1,
|
||
"t_start": group[0][0],
|
||
"t_end": group[-1][1],
|
||
"block_indices": list(range(i + 1, i + len(group) + 1)),
|
||
"texts": [text for _, _, text in group],
|
||
})
|
||
return batches
|
||
|
||
|
||
def _join_batch_texts(texts: list) -> str:
|
||
"""Join multiple subtitle texts into a single TTS prompt.
|
||
|
||
Each text already ends with punctuation from _clean_text. We join with a
|
||
space so the model reads them as one continuous multi-sentence utterance.
|
||
"""
|
||
return " ".join(texts)
|
||
|
||
|
||
def generate_batch_clips(
|
||
blocks: list, batch_size: int,
|
||
voice_filename: str, voice_ref_text: str,
|
||
base_seed: int, speed: float, clips_dir: Path,
|
||
strip_leading_ms: float = 0.0,
|
||
resume_from: int = 0,
|
||
temperature: float = 0.9,
|
||
trailing_fade_ms: float = 0,
|
||
max_new_tokens: int = 4096,
|
||
) -> tuple:
|
||
"""
|
||
Generate one audio clip per batch of `batch_size` consecutive subtitle blocks.
|
||
The texts of all blocks in a batch are joined into one TTS call, giving the
|
||
model sentence-level context to produce more natural prosody.
|
||
|
||
resume_from: skip batches 1..(resume_from-1) by loading existing clip files.
|
||
|
||
Returns (clips, batches) where:
|
||
clips — list of (actual_start_sec, clip_path) for assembly
|
||
batches — list of batch metadata dicts (saved to batches.json)
|
||
"""
|
||
batches = _group_into_batches(blocks, batch_size)
|
||
total = len(batches)
|
||
clips = []
|
||
cursor = 0.0
|
||
|
||
for batch in batches:
|
||
batch_num = batch["batch_num"]
|
||
t_start = batch["t_start"]
|
||
label = f"clip_{batch_num:04d}"
|
||
canonical = clips_dir / f"{label}.flac"
|
||
|
||
if resume_from > 0 and batch_num < resume_from:
|
||
if not canonical.exists():
|
||
sys.exit(
|
||
f"Resume from batch {resume_from}: {canonical} not found — "
|
||
f"lower --resume-from or regenerate from an earlier batch."
|
||
)
|
||
duration = get_audio_duration(canonical)
|
||
blocks_str = f"blocks {batch['block_indices'][0]}–{batch['block_indices'][-1]}"
|
||
print(f" [batch {batch_num}/{total}] skip (resuming) — "
|
||
f"{blocks_str} {canonical.name} {duration:.2f}s")
|
||
actual_start = max(t_start, cursor)
|
||
clips.append((actual_start, canonical))
|
||
cursor = actual_start + duration
|
||
continue
|
||
|
||
blocks_str = (
|
||
f"block {batch['block_indices'][0]}"
|
||
if len(batch["block_indices"]) == 1
|
||
else f"blocks {batch['block_indices'][0]}–{batch['block_indices'][-1]}"
|
||
)
|
||
print(f"\n [batch {batch_num}/{total}] {blocks_str} @ {_fmt_srt_time(t_start)}")
|
||
for i, text in enumerate(batch["texts"]):
|
||
print(f" [{i + 1}] {text[:80]!r}")
|
||
|
||
joined = _join_batch_texts(batch["texts"])
|
||
seed = base_seed + batch_num - 1
|
||
fast_clip, duration = _generate_one(
|
||
joined, voice_filename, voice_ref_text,
|
||
seed, speed, clips_dir, label, strip_leading_ms, temperature,
|
||
trailing_fade_ms=trailing_fade_ms, max_new_tokens=max_new_tokens,
|
||
)
|
||
|
||
actual_start = max(t_start, cursor)
|
||
clips.append((actual_start, fast_clip))
|
||
cursor = actual_start + duration
|
||
|
||
delay = actual_start - t_start
|
||
delay_str = f" [+{delay:.2f}s behind SRT]" if delay > 0.05 else ""
|
||
print(f" → placed at {_fmt_srt_time(actual_start)}, "
|
||
f"duration {duration:.2f}s{delay_str}")
|
||
|
||
return clips, batches
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rebuild / Regen helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _rebuild_video(args, out_dir: Path):
|
||
"""Re-assemble audio from existing batch clips and rebuild the final video."""
|
||
batches_json = out_dir / "batches.json"
|
||
if not batches_json.exists():
|
||
sys.exit(f"No batches.json found in {out_dir} — run a full v2 generation first.")
|
||
|
||
batches = json.loads(batches_json.read_text(encoding="utf-8"))
|
||
total = len(batches)
|
||
clips_dir = out_dir / "clips"
|
||
|
||
missing = [
|
||
clips_dir / f"clip_{b['batch_num']:04d}.flac"
|
||
for b in batches
|
||
if not (clips_dir / f"clip_{b['batch_num']:04d}.flac").exists()
|
||
]
|
||
if missing:
|
||
sys.exit("Missing batch clip files:\n "
|
||
+ "\n ".join(str(p) for p in missing[:10]))
|
||
|
||
clips = []
|
||
cursor = 0.0
|
||
for batch in batches:
|
||
clip_path = clips_dir / f"clip_{batch['batch_num']:04d}.flac"
|
||
duration = get_audio_duration(clip_path)
|
||
t_start = batch["t_start"]
|
||
actual_start = max(t_start, cursor)
|
||
clips.append((actual_start, clip_path))
|
||
cursor = actual_start + duration
|
||
|
||
assembled_tts = out_dir / "tts_assembled.flac"
|
||
print(f"\n Assembling {total} batch clips…")
|
||
assemble_timed_audio(clips, str(assembled_tts))
|
||
print(f" Assembled: {assembled_tts} ({assembled_tts.stat().st_size // 1024} KB)")
|
||
|
||
if args.end:
|
||
start_sec = _parse_mmss(args.start)
|
||
end_sec = _parse_mmss(args.end)
|
||
video_fragment = out_dir / "video_fragment.mkv"
|
||
if not video_fragment.exists():
|
||
extract_video_fragment(args.video, start_sec, end_sec, str(video_fragment))
|
||
video_src = str(video_fragment)
|
||
else:
|
||
video_src = args.video
|
||
|
||
stem = Path(args.video).stem
|
||
final_output = out_dir / f"{stem}_lektor.mp4"
|
||
print(f"\n Building 2-track video "
|
||
f"(TTS {int(args.tts_volume * 100)}% / orig {int(args.orig_volume * 100)}%)…")
|
||
merge_voiceover(video_src, str(assembled_tts), str(final_output),
|
||
orig_volume=args.orig_volume, tts_volume=args.tts_volume)
|
||
|
||
print(f"\nDone! Output: {final_output}")
|
||
|
||
|
||
def _regen_clips(args, out_dir: Path):
|
||
"""Regenerate specific batch clips by batch number.
|
||
Only replaces FLAC files — run --rebuild afterwards to reassemble audio and video.
|
||
"""
|
||
batches_json = out_dir / "batches.json"
|
||
if not batches_json.exists():
|
||
sys.exit(f"No batches.json found in {out_dir} — run a full v2 generation first.")
|
||
|
||
batches = json.loads(batches_json.read_text(encoding="utf-8"))
|
||
total = len(batches)
|
||
regen_set = set(args.regen)
|
||
clips_dir = out_dir / "clips"
|
||
clips_dir.mkdir(exist_ok=True)
|
||
|
||
bad = [n for n in regen_set if n < 1 or n > total]
|
||
if bad:
|
||
sys.exit(f"Batch number(s) out of range (1–{total}): {bad}")
|
||
|
||
wav_path = Path(args.voice)
|
||
if not wav_path.is_file():
|
||
wav_path = VOICE_PROMPTS_DIR / f"{args.voice}.wav"
|
||
if not wav_path.is_file():
|
||
sys.exit(f"Voice WAV not found: {wav_path}")
|
||
voice_filename, voice_ref_text = upload_voice_wav(wav_path)
|
||
|
||
base_seed = args.seed if args.seed != 0 else random.randint(1, 2**31 - 1)
|
||
print(f"\n Regenerating batches {sorted(regen_set)} with seed {base_seed}…")
|
||
|
||
for batch in batches:
|
||
batch_num = batch["batch_num"]
|
||
if batch_num not in regen_set:
|
||
continue
|
||
|
||
joined = _join_batch_texts(batch["texts"])
|
||
label = f"clip_{batch_num:04d}"
|
||
blocks_str = f"blocks {batch['block_indices'][0]}–{batch['block_indices'][-1]}"
|
||
print(f"\n REGEN [batch {batch_num}/{total}] {blocks_str}")
|
||
for i, text in enumerate(batch["texts"]):
|
||
print(f" [{i + 1}] {text[:80]!r}")
|
||
|
||
seed = base_seed + batch_num - 1
|
||
clip_path, duration = _generate_one(
|
||
joined, voice_filename, voice_ref_text,
|
||
seed, args.speed, clips_dir, label,
|
||
args.strip_leading_ms, args.temperature,
|
||
trailing_fade_ms=args.trailing_fade_ms,
|
||
max_new_tokens=args.max_tokens,
|
||
)
|
||
print(f" → saved as {clip_path.name} ({duration:.2f}s)")
|
||
|
||
print(f"\nDone. Run --rebuild to reassemble audio and rebuild the video.")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main():
|
||
global COMFY_URL
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"Generate AI voiceover via ComfyUI + Higgs Audio V3, in batches of "
|
||
"N subtitle blocks per TTS call for contextual multi-sentence generation."
|
||
)
|
||
)
|
||
parser.add_argument("srt", help="Input SRT subtitle file")
|
||
parser.add_argument("video", help="Input video file (MKV, MP4, …)")
|
||
parser.add_argument("--start", default="00:00", metavar="MM:SS",
|
||
help="Fragment start time (default: 00:00)")
|
||
parser.add_argument("--end", default=None, metavar="MM:SS",
|
||
help="Fragment end time")
|
||
parser.add_argument("--voice", default="gudowski_short",
|
||
help="Voice prompt name or WAV path (default: gudowski_short)")
|
||
parser.add_argument("--batch-size", type=int, default=5, metavar="N",
|
||
help="Subtitle blocks per TTS call (default: 5). "
|
||
"Larger batches give the model more context but produce "
|
||
"longer clips that are harder to regen individually.")
|
||
parser.add_argument("--speed", type=float, default=1.2, metavar="FACTOR",
|
||
help="Per-clip playback speed via atempo (default: 1.2)")
|
||
parser.add_argument("--seed", type=int, default=0,
|
||
help="Base generation seed (default: 0 = random each run)")
|
||
parser.add_argument("--orig-volume", type=float, default=0.60, metavar="0-1",
|
||
help="Original audio volume in mix (default: 0.60)")
|
||
parser.add_argument("--tts-volume", type=float, default=1.8, metavar="0-2",
|
||
help="Voiceover volume in mix (default: 1.8)")
|
||
parser.add_argument("--output-dir", default="output",
|
||
help="Output directory (default: output/)")
|
||
parser.add_argument("--comfy-url", default=COMFY_URL,
|
||
help=f"ComfyUI URL (default: {COMFY_URL})")
|
||
parser.add_argument("--skip-tts", metavar="AUDIO_FILE",
|
||
help="Skip TTS and use this pre-assembled audio file")
|
||
parser.add_argument("--temperature", type=float, default=0.9, metavar="T",
|
||
help="Sampling temperature (default: 0.9)")
|
||
parser.add_argument("--strip-leading-ms", type=float, default=0.0, metavar="MS",
|
||
help="Strip N ms from the start of each clip (default: 0)")
|
||
parser.add_argument("--trailing-fade-ms", type=float, default=80.0, metavar="MS",
|
||
help="Fade out the last N ms of each clip (default: 80)")
|
||
parser.add_argument("--max-tokens", type=int, default=4096, metavar="N",
|
||
help="Max audio tokens per TTS call (default: 4096). "
|
||
"Raise if batch clips are cut off mid-sentence.")
|
||
parser.add_argument("--resume-from", type=int, default=0, metavar="N",
|
||
help="Resume a crashed run: skip batches 1..(N-1) by loading "
|
||
"their existing clip files, generate from batch N onwards.")
|
||
parser.add_argument("--rebuild", action="store_true",
|
||
help="Re-assemble audio from existing clips and rebuild the "
|
||
"final video without any TTS generation.")
|
||
parser.add_argument("--regen", nargs="+", type=int, metavar="N",
|
||
help="Regenerate specific batch clip(s) by batch number "
|
||
"(as shown during generation and in batches.json). "
|
||
"Run --rebuild afterwards to reassemble the video.")
|
||
args = parser.parse_args()
|
||
|
||
COMFY_URL = args.comfy_url
|
||
|
||
out_dir = Path(args.output_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
clips_dir = out_dir / "clips"
|
||
clips_dir.mkdir(exist_ok=True)
|
||
|
||
# ── Rebuild mode ──────────────────────────────────────────────────────────
|
||
if args.rebuild:
|
||
_rebuild_video(args, out_dir)
|
||
return
|
||
|
||
# ── Regen mode ────────────────────────────────────────────────────────────
|
||
if args.regen:
|
||
_regen_clips(args, out_dir)
|
||
return
|
||
|
||
start_sec = _parse_mmss(args.start)
|
||
end_sec = _parse_mmss(args.end) if args.end else None
|
||
|
||
# ── Step 1: Parse subtitles ───────────────────────────────────────────────
|
||
print(f"\n[1/4] Parsing subtitles {args.start}–{args.end or 'end'}…")
|
||
if end_sec is None:
|
||
end_sec = 99 * 3600
|
||
blocks = parse_subtitle_blocks(args.srt, start_sec, end_sec)
|
||
if not blocks:
|
||
sys.exit("No subtitle blocks found in the specified range.")
|
||
blocks = split_dialogue_blocks(blocks)
|
||
n_batches = (len(blocks) + args.batch_size - 1) // args.batch_size
|
||
print(f" {len(blocks)} subtitle blocks → {n_batches} batches "
|
||
f"of up to {args.batch_size}.")
|
||
|
||
srt_preview = out_dir / "tts_input.srt"
|
||
srt_preview.write_text(blocks_to_srt(blocks), encoding="utf-8")
|
||
print(f" Saved individual blocks to {srt_preview} for reference.")
|
||
|
||
# ── Step 2: Generate TTS per batch ────────────────────────────────────────
|
||
if args.skip_tts:
|
||
assembled_tts = Path(args.skip_tts)
|
||
print(f"\n[2/4] Using existing TTS audio: {assembled_tts}")
|
||
else:
|
||
print(f"\n[2/4] Generating TTS ({n_batches} batches × up to "
|
||
f"{args.batch_size} blocks, speed {args.speed}×)…")
|
||
|
||
wav_path = Path(args.voice)
|
||
if not wav_path.is_file():
|
||
wav_path = VOICE_PROMPTS_DIR / f"{args.voice}.wav"
|
||
if not wav_path.is_file():
|
||
sys.exit(f"Voice WAV not found: {wav_path}")
|
||
print(f" Voice: {wav_path.name}")
|
||
|
||
base_seed = args.seed if args.seed != 0 else random.randint(1, 2**31 - 1)
|
||
print(f" Base seed: {base_seed}")
|
||
voice_filename, voice_ref_text = upload_voice_wav(wav_path)
|
||
|
||
clips, batches = generate_batch_clips(
|
||
blocks, args.batch_size,
|
||
voice_filename, voice_ref_text,
|
||
base_seed, args.speed, clips_dir,
|
||
strip_leading_ms=args.strip_leading_ms,
|
||
resume_from=args.resume_from,
|
||
temperature=args.temperature,
|
||
trailing_fade_ms=args.trailing_fade_ms,
|
||
max_new_tokens=args.max_tokens,
|
||
)
|
||
|
||
batches_json = out_dir / "batches.json"
|
||
batches_json.write_text(
|
||
json.dumps(batches, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
print(f"\n Batch map saved to {batches_json}")
|
||
|
||
assembled_tts = out_dir / "tts_assembled.flac"
|
||
assemble_timed_audio(clips, str(assembled_tts))
|
||
print(f" Assembled: {assembled_tts} ({assembled_tts.stat().st_size // 1024} KB)")
|
||
print(f" Clips saved to {clips_dir}")
|
||
|
||
# ── Step 3: Extract video fragment ───────────────────────────────────────
|
||
print(f"\n[3/4] Extracting video fragment…")
|
||
real_end = _parse_mmss(args.end) if args.end else None
|
||
if real_end is not None:
|
||
video_fragment = out_dir / "video_fragment.mkv"
|
||
extract_video_fragment(args.video, start_sec, real_end, str(video_fragment))
|
||
else:
|
||
video_fragment = Path(args.video)
|
||
|
||
# ── Step 4: Mix voiceover with video ─────────────────────────────────────
|
||
stem = Path(args.video).stem
|
||
final_output = out_dir / f"{stem}_lektor.mp4"
|
||
print(f"\n[4/4] Building 2-track video (Original + Polish lektor "
|
||
f"TTS {int(args.tts_volume * 100)}% / orig {int(args.orig_volume * 100)}%)…")
|
||
merge_voiceover(str(video_fragment), str(assembled_tts), str(final_output),
|
||
orig_volume=args.orig_volume, tts_volume=args.tts_volume)
|
||
|
||
print(f"\nDone! Output: {final_output}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|