lukas/lektor-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
fc4b761f52724024acfd703b50fa46e7fbcd428d
lektor-ai/lektor_ai_comfy_v2.py
T

885 lines
35 KiB
Python
Raw Normal View History

Initial commit: Polish AI dubbing pipeline (ComfyUI + Higgs Audio V3)
2026-08-09 14:29:53 +02:00
#!/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/gudowski")
# ---------------------------------------------------------------------------
# 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()
Reference in New Issue Copy Permalink