142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
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
|
|
return 2
|
|
|
|
if args.kind == "music":
|
|
if synthesize_audiocraft("music", prompt, output, args.duration, args.model):
|
|
return 0
|
|
return 2
|
|
|
|
if synthesize_audiocraft("sfx", prompt, output, args.duration, args.model):
|
|
return 0
|
|
return 2
|
|
|
|
|
|
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
|
|
|
|
return False
|
|
|
|
|
|
def synthesize_audiocraft(kind: str, prompt: str, output: Path, duration: float, model_name: str) -> bool:
|
|
model_path = Path(model_name).expanduser() if model_name else None
|
|
if torch is None or model_path is None or not model_path.exists():
|
|
return False
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
try:
|
|
if kind == "music":
|
|
from audiocraft.models import MusicGen # type: ignore
|
|
|
|
model = MusicGen.get_pretrained(str(model_path.resolve()))
|
|
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(str(model_path.resolve()))
|
|
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 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_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())
|