video_editing_poc/tools/local_asset_worker.py

181 lines
6.2 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import argparse
import math
import os
import shutil
import subprocess
import wave
from array import array
from pathlib import Path
try:
import soundfile as sf # type: ignore
except Exception: # pragma: no cover - optional runtime dependency
sf = None
try:
import torch # type: ignore
except Exception: # pragma: no cover - optional runtime dependency
torch = None
try:
import numpy as np # type: ignore
except Exception: # pragma: no cover - optional runtime dependency
np = None
def main() -> int:
parser = argparse.ArgumentParser(description="Local cinematic asset synthesizer")
parser.add_argument("kind", choices=["voiceover", "music", "sfx"])
parser.add_argument("--prompt-file", required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--duration", type=float, default=8.0)
parser.add_argument("--model", default="")
parser.add_argument("--piper-binary", default=os.getenv("LOCAL_ASSET_PIPER_BINARY", "piper"))
args = parser.parse_args()
prompt = Path(args.prompt_file).read_text(encoding="utf-8").strip()
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
if args.kind == "voiceover":
if synthesize_voiceover(prompt, output, args.piper_binary, args.model):
return 0
write_silence(output, max(1.0, args.duration))
return 0
if args.kind == "music":
if synthesize_audiocraft("music", prompt, output, args.duration, args.model):
return 0
write_fallback_tone(output, max(4.0, args.duration), 110.0, 0.02)
return 0
if synthesize_audiocraft("sfx", prompt, output, args.duration, args.model):
return 0
write_fallback_tone(output, max(0.5, min(1.0, args.duration)), 880.0, 0.12)
return 0
def synthesize_voiceover(prompt: str, output: Path, piper_binary: str, model: str) -> bool:
model_path = model.strip()
if model_path and shutil.which(piper_binary):
command = [piper_binary, "--model", model_path, "--output_file", str(output)]
try:
process = subprocess.run(
command,
input=prompt.encode("utf-8"),
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
check=False,
)
return process.returncode == 0 and output.is_file()
except Exception:
return False
say = shutil.which("say")
if say:
process = subprocess.run([say, "-o", str(output), prompt], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
return process.returncode == 0 and output.is_file()
espeak = shutil.which("espeak")
if espeak:
process = subprocess.run([espeak, "-w", str(output), prompt], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
return process.returncode == 0 and output.is_file()
return False
def synthesize_audiocraft(kind: str, prompt: str, output: Path, duration: float, model_name: str) -> bool:
if torch is None:
return False
try:
if kind == "music":
from audiocraft.models import MusicGen # type: ignore
model = MusicGen.get_pretrained(normalize_model_name(model_name, "musicgen-small"))
model.set_generation_params(duration=max(1.0, float(duration)))
audio = model.generate([prompt])[0]
write_tensor_audio(audio, output, getattr(model, "sample_rate", 32000))
return output.is_file()
from audiocraft.models import AudioGen # type: ignore
model = AudioGen.get_pretrained(normalize_model_name(model_name, "audiogen-medium"))
model.set_generation_params(duration=max(0.5, float(duration)))
audio = model.generate([prompt])[0]
write_tensor_audio(audio, output, getattr(model, "sample_rate", 32000))
return output.is_file()
except Exception:
return False
def normalize_model_name(model_name: str, default_name: str) -> str:
value = model_name.strip() if model_name else default_name
if "/" in value:
return value
return f"facebook/{value}"
def write_tensor_audio(audio, output: Path, sample_rate: int) -> None:
if np is None:
raise RuntimeError("numpy is required to write model-generated audio")
if hasattr(audio, "detach"):
audio = audio.detach().cpu().numpy()
elif torch is not None and isinstance(audio, torch.Tensor):
audio = audio.cpu().numpy()
audio = np.asarray(audio)
if audio.ndim == 3:
audio = audio[0]
if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
audio = audio.transpose(1, 0)
if sf is not None:
sf.write(str(output), audio, sample_rate)
return
write_wave_file(output, audio, sample_rate)
def write_fallback_tone(output: Path, duration: float, frequency: float, volume: float) -> None:
sample_rate = 48000
total_samples = int(sample_rate * duration)
frames = array(
"h",
[
int(max(-1.0, min(1.0, math.sin(2.0 * math.pi * frequency * (i / sample_rate))) * volume) * 32767.0)
for i in range(total_samples)
],
)
write_wave_file(output, frames, sample_rate, channels=1)
def write_silence(output: Path, duration: float) -> None:
sample_rate = 24000
total_samples = int(sample_rate * duration)
frames = array("h", [0] * total_samples)
write_wave_file(output, frames, sample_rate, channels=1)
def write_wave_file(output: Path, waveform, sample_rate: int, channels: int | None = None) -> None:
if hasattr(waveform, "ndim") and np is not None:
waveform = np.asarray(waveform)
if waveform.ndim == 1:
waveform = waveform[:, None]
waveform = np.clip(waveform, -1.0, 1.0)
pcm = (waveform * 32767.0).astype(np.int16)
channels = pcm.shape[1]
frames = array("h", pcm.reshape(-1).tolist())
else:
frames = waveform
if channels is None:
channels = 1
with wave.open(str(output), "wb") as wav:
wav.setnchannels(channels)
wav.setsampwidth(2)
wav.setframerate(sample_rate)
wav.writeframes(frames.tobytes())
if __name__ == "__main__":
raise SystemExit(main())