238 lines
8.5 KiB
Python
238 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Local cinematic asset synthesizer.
|
|
|
|
Generates a single audio asset (voiceover | music | sfx) from a text prompt using only
|
|
locally-resident models. No network is used at inference time.
|
|
|
|
voiceover -> Piper (ONNX voice model passed via --model, run through --piper-binary)
|
|
music -> MusicGen (transformers), model directory derived from --model
|
|
sfx -> AudioLDM2 (diffusers), model directory derived from --model
|
|
|
|
Contract (unchanged): exit 0 on success, exit 2 on failure; writes a WAV to --output.
|
|
The --model argument is a licensed, regular file inside the model directory (the Java caller
|
|
enforces the license sidecar); for music/sfx the model directory is that file's parent.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import wave
|
|
from array import array
|
|
from pathlib import Path
|
|
|
|
# Force offline model loading: weights are pre-provisioned local directories.
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
|
|
|
OUTPUT_SAMPLE_RATE = 48000 # pipeline mastering target
|
|
|
|
try:
|
|
import soundfile as sf # type: ignore
|
|
except Exception: # pragma: no cover - optional runtime dependency
|
|
sf = None
|
|
|
|
try:
|
|
import numpy as np # type: ignore
|
|
except Exception: # pragma: no cover - optional runtime dependency
|
|
np = None
|
|
|
|
try:
|
|
import torch # type: ignore
|
|
except Exception: # pragma: no cover - optional runtime dependency
|
|
torch = 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)
|
|
|
|
try:
|
|
if args.kind == "voiceover":
|
|
ok = synthesize_voiceover(prompt, output, args.piper_binary, args.model)
|
|
elif args.kind == "music":
|
|
ok = synthesize_musicgen(prompt, output, args.duration, args.model)
|
|
else:
|
|
ok = synthesize_audioldm2(prompt, output, args.duration, args.model)
|
|
except Exception as exc: # fail closed, never emit a partial asset
|
|
_cleanup(output)
|
|
print(f"error: {args.kind} synthesis failed: {exc}", flush=True)
|
|
return 2
|
|
|
|
if ok and output.is_file() and output.stat().st_size > 0:
|
|
return 0
|
|
_cleanup(output)
|
|
return 2
|
|
|
|
|
|
def _cleanup(output: Path) -> None:
|
|
try:
|
|
if output.exists():
|
|
output.unlink()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _model_dir(model_arg: str) -> Path | None:
|
|
"""The Java caller passes a licensed regular file inside the model directory."""
|
|
if not model_arg:
|
|
return None
|
|
marker = Path(model_arg).expanduser()
|
|
directory = marker.parent if marker.is_file() else marker
|
|
return directory if directory.is_dir() else None
|
|
|
|
|
|
# --------------------------------------------------------------------------- voiceover
|
|
|
|
def synthesize_voiceover(prompt: str, output: Path, piper_binary: str, model: str) -> bool:
|
|
model_path = model.strip()
|
|
if not model_path or not Path(model_path).is_file() or not shutil.which(piper_binary):
|
|
return False
|
|
command = [piper_binary, "--model", model_path, "--output_file", str(output)]
|
|
process = subprocess.run(
|
|
command,
|
|
input=prompt.encode("utf-8"),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
check=False,
|
|
)
|
|
if process.returncode != 0 or not output.is_file():
|
|
return False
|
|
# Normalise to the pipeline sample rate / mono.
|
|
_resample_in_place(output)
|
|
return output.is_file()
|
|
|
|
|
|
# --------------------------------------------------------------------------- music
|
|
|
|
def synthesize_musicgen(prompt: str, output: Path, duration: float, model: str) -> bool:
|
|
directory = _model_dir(model)
|
|
if torch is None or directory is None:
|
|
return False
|
|
from transformers import MusicgenForConditionalGeneration, AutoProcessor # lazy
|
|
|
|
processor = AutoProcessor.from_pretrained(directory)
|
|
net = MusicgenForConditionalGeneration.from_pretrained(directory, torch_dtype=torch.float32)
|
|
net.to("cpu")
|
|
frame_rate = int(getattr(net.config.audio_encoder, "frame_rate", 50) or 50)
|
|
src_sr = int(net.config.audio_encoder.sampling_rate)
|
|
max_new_tokens = max(frame_rate, int(round(max(1.0, float(duration)) * frame_rate)))
|
|
inputs = processor(text=[prompt], padding=True, return_tensors="pt")
|
|
with torch.no_grad():
|
|
audio = net.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, guidance_scale=3.0)
|
|
wav = audio[0, 0].detach().cpu().numpy()
|
|
_write_audio(wav, src_sr, output)
|
|
return output.is_file()
|
|
|
|
|
|
# --------------------------------------------------------------------------- sfx
|
|
|
|
def synthesize_audioldm2(prompt: str, output: Path, duration: float, model: str) -> bool:
|
|
directory = _model_dir(model)
|
|
if torch is None or directory is None:
|
|
return False
|
|
from diffusers import AudioLDM2Pipeline # lazy
|
|
|
|
steps = int(os.getenv("LOCAL_ASSET_SFX_STEPS", "40"))
|
|
pipe = AudioLDM2Pipeline.from_pretrained(directory, torch_dtype=torch.float32)
|
|
pipe = pipe.to("cpu")
|
|
generator = torch.Generator("cpu").manual_seed(0)
|
|
length = max(1.0, float(duration))
|
|
result = pipe(
|
|
prompt,
|
|
negative_prompt="low quality, average quality, noise",
|
|
num_inference_steps=steps,
|
|
audio_length_in_s=length,
|
|
generator=generator,
|
|
)
|
|
wav = result.audios[0]
|
|
_write_audio(wav, 16000, output)
|
|
return output.is_file()
|
|
|
|
|
|
# --------------------------------------------------------------------------- io helpers
|
|
|
|
def _to_mono_f32(audio, src_sr: int):
|
|
audio = np.asarray(audio, dtype="float32")
|
|
if audio.ndim == 3:
|
|
audio = audio[0]
|
|
if audio.ndim == 2:
|
|
# (channels, samples) -> mono
|
|
if audio.shape[0] < audio.shape[1]:
|
|
audio = audio.mean(axis=0)
|
|
else:
|
|
audio = audio.mean(axis=1)
|
|
return audio.reshape(-1)
|
|
|
|
|
|
def _resample(audio, src_sr: int, dst_sr: int):
|
|
if src_sr == dst_sr:
|
|
return audio
|
|
if torch is not None:
|
|
import torchaudio # lazy
|
|
tensor = torch.from_numpy(np.asarray(audio, dtype="float32")).unsqueeze(0)
|
|
out = torchaudio.functional.resample(tensor, src_sr, dst_sr)
|
|
return out.squeeze(0).numpy()
|
|
# Fallback: linear interpolation
|
|
ratio = dst_sr / float(src_sr)
|
|
idx = np.arange(int(len(audio) * ratio)) / ratio
|
|
return np.interp(idx, np.arange(len(audio)), audio).astype("float32")
|
|
|
|
|
|
def _peak_normalize(audio, target_peak: float = 0.89):
|
|
peak = float(np.max(np.abs(audio))) if audio.size else 0.0
|
|
if peak > 1e-6:
|
|
audio = audio * (target_peak / peak)
|
|
return np.clip(audio, -1.0, 1.0)
|
|
|
|
|
|
def _write_audio(audio, src_sr: int, output: Path) -> None:
|
|
if np is None:
|
|
raise RuntimeError("numpy is required to write model-generated audio")
|
|
mono = _to_mono_f32(audio, src_sr)
|
|
mono = _resample(mono, src_sr, OUTPUT_SAMPLE_RATE)
|
|
mono = _peak_normalize(mono)
|
|
if sf is not None:
|
|
sf.write(str(output), mono, OUTPUT_SAMPLE_RATE, subtype="PCM_16")
|
|
return
|
|
_write_wave_file(output, mono, OUTPUT_SAMPLE_RATE)
|
|
|
|
|
|
def _resample_in_place(output: Path) -> None:
|
|
"""Resample an existing WAV (e.g. Piper 22.05 kHz) to the pipeline rate."""
|
|
if sf is None or np is None:
|
|
return
|
|
data, sr = sf.read(str(output), dtype="float32", always_2d=False)
|
|
if sr == OUTPUT_SAMPLE_RATE and data.ndim == 1:
|
|
return
|
|
mono = _to_mono_f32(data if data.ndim > 1 else data.reshape(1, -1), sr)
|
|
mono = _resample(mono, sr, OUTPUT_SAMPLE_RATE)
|
|
sf.write(str(output), mono, OUTPUT_SAMPLE_RATE, subtype="PCM_16")
|
|
|
|
|
|
def _write_wave_file(output: Path, waveform, sample_rate: int) -> None:
|
|
waveform = np.clip(np.asarray(waveform, dtype="float32"), -1.0, 1.0)
|
|
pcm = (waveform * 32767.0).astype("int16")
|
|
frames = array("h", pcm.reshape(-1).tolist())
|
|
with wave.open(str(output), "wb") as wav:
|
|
wav.setnchannels(1)
|
|
wav.setsampwidth(2)
|
|
wav.setframerate(sample_rate)
|
|
wav.writeframes(frames.tobytes())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|