forked from jsl/video_editing_poc
Wire local generative audio models into the highlight asset worker
Replace the unusable audiocraft path (requires xformers, which has no Intel-Mac build) with runtimes proven to work offline on this machine: - music: transformers MusicGen (facebook/musicgen-small) - sfx: diffusers AudioLDM2 (cvssp/audioldm2), resampled 16k -> 48k - voiceover: Piper (unchanged), normalized to 48 kHz mono The worker CLI contract and exit codes are preserved, so the Java LocalAssetSynthesizer license gate and fail-closed behavior are unchanged. Add tools/provision_local_models.py to materialize models into models/ from the local HF cache with no network. Models and their license sidecars live under the git-ignored models/ dir; both audio models are non-commercial (CC-BY-NC-4.0 / CC-BY-NC-SA-4.0), recorded for later production review. Add docs/cinematic-highlight-poc-plan.md tracking the PoC plan and milestones. mvn -o verify: 245 tests, 0 failures/errors/skips (unchanged). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bg76sLc43Wc3j5ZcLkboYR
This commit is contained in:
parent
f6dc6b8a83
commit
a95d1fa0ac
|
|
@ -42,3 +42,7 @@ build/
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.pyc
|
*.pyc
|
||||||
yolov*.pt
|
yolov*.pt
|
||||||
|
|
||||||
|
# Local model bundle for the cinematic highlight PoC (large, provisioned out-of-band)
|
||||||
|
/models/
|
||||||
|
.claude/settings.local.json
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,68 @@
|
||||||
|
# Local Cinematic Highlight PoC — Plan & Milestones
|
||||||
|
|
||||||
|
Owner: (Jaden) · Started: 2026-07-21 · Machine: Intel Mac (x86_64, 16 CPU, 32 GB, no GPU)
|
||||||
|
|
||||||
|
## Objective
|
||||||
|
|
||||||
|
Produce ONE genuinely cinematic highlight from a local source video, end-to-end through the existing
|
||||||
|
pipeline, using only local models resident in the service runtime:
|
||||||
|
Piper (voiceover) + MusicGen (music) + AudioLDM2 (SFX), with deliberate visual treatment and pacing,
|
||||||
|
and explicit human approval before any render.
|
||||||
|
|
||||||
|
A valid MP4 or a passing test is NOT success. Success = representative output passes measured media QA
|
||||||
|
and a human creative review.
|
||||||
|
|
||||||
|
## Non-negotiables (still in force during the PoC)
|
||||||
|
|
||||||
|
No external AI services in the media path · no placeholder silence/tones/OS `say` as a finished asset ·
|
||||||
|
no unlicensed assets · no rendering without explicit approval for the specific project · no inference-time
|
||||||
|
network (models load from local dirs, `HF_HUB_OFFLINE=1`). NOTE: the user authorized model DOWNLOADS on
|
||||||
|
2026-07-21 (one-time provisioning), reversing the earlier no-download stance; inference stays offline.
|
||||||
|
|
||||||
|
## Proven capability baseline (2026-07-21)
|
||||||
|
|
||||||
|
| Model | Runtime | Status | Evidence |
|
||||||
|
|---|---|---|---|
|
||||||
|
| Piper `en_US-lessac-medium` | piper-tts 1.5.0 | ✅ works | real speech, 22.05 kHz |
|
||||||
|
| MusicGen `musicgen-small` | transformers 4.44.2 + torch 2.2.2 | ✅ works | 5 s / 44.7 s CPU, mean −18 dB, 32 kHz |
|
||||||
|
| AudioLDM2 `cvssp/audioldm2` | diffusers 0.30.3 | ✅ works | 3 s / 22.3 s CPU, mean −21.8 dB, 16 kHz |
|
||||||
|
|
||||||
|
Environment traps (see memory `local-model-runtime-intel-mac`): torch capped at 2.2.2 (Intel-Mac),
|
||||||
|
numpy<2, transformers must be 4.x, audiocraft/xformers unusable here, `HF_HUB_DISABLE_XET=1` for downloads.
|
||||||
|
|
||||||
|
## Plan
|
||||||
|
|
||||||
|
### Phase 1 — Generate all three assets THROUGH the pipeline
|
||||||
|
- [x] 1.1 Rewrite `tools/local_asset_worker.py`: music→transformers MusicGen, sfx→diffusers AudioLDM2
|
||||||
|
(resample to 48 kHz mono), voiceover→Piper (unchanged). audiocraft path removed. CLI + exit codes preserved.
|
||||||
|
- [x] 1.2 Materialize models into stable `models/` dirs via `save_pretrained` (`tools/provision_local_models.py`):
|
||||||
|
`models/musicgen-small` (2.2 G), `models/audioldm2` (4.2 G); `models/piper` (60 M) already present.
|
||||||
|
- [x] 1.3 License/provenance `.license.txt` sidecars for each model marker file (config.json / model_index.json /
|
||||||
|
voice .onnx). BOTH audio models are NON-COMMERCIAL (MusicGen CC-BY-NC-4.0, AudioLDM2 CC-BY-NC-SA-4.0) —
|
||||||
|
flagged for production review. Honors the `AssetLicensePolicy` regular-file gate; no Java change.
|
||||||
|
- [x] 1.4 Smoke test: voiceover 3.67 s/−15.8 dB, music 4.94 s/−12.0 dB, sfx 3.00 s/−20.6 dB — all 48 kHz mono,
|
||||||
|
real signal, exit 0.
|
||||||
|
- [x] 1.5 `mvn -o verify` → 245 tests / 62 classes / 0 failures / 0 errors / 0 skips (unchanged from baseline).
|
||||||
|
|
||||||
|
**Phase 1 COMPLETE (2026-07-21).**
|
||||||
|
|
||||||
|
### Phase 2 — One approved end-to-end highlight (STOP before render for explicit approval)
|
||||||
|
- [ ] 2.1 Opt-in `localpoc` Spring profile (base/production defaults untouched): venv python, model paths,
|
||||||
|
offline flags, no bootstrap auto-start, heuristic fallback off, isolated PoC I/O dirs, render disabled.
|
||||||
|
- [ ] 2.2 Reprocess the DJI source so `category.json` + `highlight-candidates.json` are written.
|
||||||
|
- [ ] 2.3 Hand-author a director `edit-plan.json` grounded in persisted candidates (manual, local-only).
|
||||||
|
- [ ] 2.4 Explicit approval for the specific project → render → final + manifest + QA report.
|
||||||
|
- [ ] 2.5 Measure technical QA + structured human creative review.
|
||||||
|
|
||||||
|
### Phase 3 — Iterate to quality
|
||||||
|
- [ ] Improve selection, pacing, visual treatment, voice, music/SFX fit; keep experiments reproducible.
|
||||||
|
|
||||||
|
## Milestone log
|
||||||
|
- 2026-07-21: Models provisioned & individually proven (Piper, MusicGen, AudioLDM2). Plan approved. Phase 1 started.
|
||||||
|
- 2026-07-21: **Phase 1 complete.** Worker rewritten (audiocraft→transformers MusicGen + diffusers AudioLDM2),
|
||||||
|
models materialized to `models/` with license sidecars, all 3 asset kinds generate 48 kHz mono real audio
|
||||||
|
through the worker, `mvn -o verify` green (245/0/0/0). Paused for review before Phase 2.
|
||||||
|
|
||||||
|
## Deferred (until output-quality gate passes)
|
||||||
|
Production hardening: Spring Security/OIDC, PostgreSQL/Testcontainers, containers/K8s, CI/CD, distributed
|
||||||
|
ops, digest-bound authenticated approval. Recorded, not deleted.
|
||||||
|
|
@ -1,4 +1,17 @@
|
||||||
#!/usr/bin/env python3
|
#!/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
|
from __future__ import annotations
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
|
|
@ -9,21 +22,28 @@ import wave
|
||||||
from array import array
|
from array import array
|
||||||
from pathlib import Path
|
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:
|
try:
|
||||||
import soundfile as sf # type: ignore
|
import soundfile as sf # type: ignore
|
||||||
except Exception: # pragma: no cover - optional runtime dependency
|
except Exception: # pragma: no cover - optional runtime dependency
|
||||||
sf = None
|
sf = None
|
||||||
|
|
||||||
try:
|
|
||||||
import torch # type: ignore
|
|
||||||
except Exception: # pragma: no cover - optional runtime dependency
|
|
||||||
torch = None
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import numpy as np # type: ignore
|
import numpy as np # type: ignore
|
||||||
except Exception: # pragma: no cover - optional runtime dependency
|
except Exception: # pragma: no cover - optional runtime dependency
|
||||||
np = None
|
np = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
import torch # type: ignore
|
||||||
|
except Exception: # pragma: no cover - optional runtime dependency
|
||||||
|
torch = None
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
parser = argparse.ArgumentParser(description="Local cinematic asset synthesizer")
|
parser = argparse.ArgumentParser(description="Local cinematic asset synthesizer")
|
||||||
|
|
@ -39,99 +59,175 @@ def main() -> int:
|
||||||
output = Path(args.output)
|
output = Path(args.output)
|
||||||
output.parent.mkdir(parents=True, exist_ok=True)
|
output.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
if args.kind == "voiceover":
|
try:
|
||||||
if synthesize_voiceover(prompt, output, args.piper_binary, args.model):
|
if args.kind == "voiceover":
|
||||||
return 0
|
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
|
return 2
|
||||||
|
|
||||||
if args.kind == "music":
|
if ok and output.is_file() and output.stat().st_size > 0:
|
||||||
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 0
|
||||||
|
_cleanup(output)
|
||||||
return 2
|
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:
|
def synthesize_voiceover(prompt: str, output: Path, piper_binary: str, model: str) -> bool:
|
||||||
model_path = model.strip()
|
model_path = model.strip()
|
||||||
if model_path and shutil.which(piper_binary):
|
if not model_path or not Path(model_path).is_file() or not 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
|
return False
|
||||||
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
command = [piper_binary, "--model", model_path, "--output_file", str(output)]
|
||||||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
process = subprocess.run(
|
||||||
try:
|
command,
|
||||||
if kind == "music":
|
input=prompt.encode("utf-8"),
|
||||||
from audiocraft.models import MusicGen # type: ignore
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
model = MusicGen.get_pretrained(str(model_path.resolve()))
|
check=False,
|
||||||
model.set_generation_params(duration=max(1.0, float(duration)))
|
)
|
||||||
audio = model.generate([prompt])[0]
|
if process.returncode != 0 or not output.is_file():
|
||||||
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
|
return False
|
||||||
|
# Normalise to the pipeline sample rate / mono.
|
||||||
|
_resample_in_place(output)
|
||||||
|
return output.is_file()
|
||||||
|
|
||||||
def write_tensor_audio(audio, output: Path, sample_rate: int) -> None:
|
|
||||||
if np is None:
|
# --------------------------------------------------------------------------- music
|
||||||
raise RuntimeError("numpy is required to write model-generated audio")
|
|
||||||
if hasattr(audio, "detach"):
|
def synthesize_musicgen(prompt: str, output: Path, duration: float, model: str) -> bool:
|
||||||
audio = audio.detach().cpu().numpy()
|
directory = _model_dir(model)
|
||||||
elif torch is not None and isinstance(audio, torch.Tensor):
|
if torch is None or directory is None:
|
||||||
audio = audio.cpu().numpy()
|
return False
|
||||||
audio = np.asarray(audio)
|
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:
|
if audio.ndim == 3:
|
||||||
audio = audio[0]
|
audio = audio[0]
|
||||||
if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
|
if audio.ndim == 2:
|
||||||
audio = audio.transpose(1, 0)
|
# (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:
|
if sf is not None:
|
||||||
sf.write(str(output), audio, sample_rate)
|
sf.write(str(output), mono, OUTPUT_SAMPLE_RATE, subtype="PCM_16")
|
||||||
return
|
return
|
||||||
write_wave_file(output, audio, sample_rate)
|
_write_wave_file(output, mono, OUTPUT_SAMPLE_RATE)
|
||||||
|
|
||||||
|
|
||||||
def write_wave_file(output: Path, waveform, sample_rate: int, channels: int | None = None) -> None:
|
def _resample_in_place(output: Path) -> None:
|
||||||
if hasattr(waveform, "ndim") and np is not None:
|
"""Resample an existing WAV (e.g. Piper 22.05 kHz) to the pipeline rate."""
|
||||||
waveform = np.asarray(waveform)
|
if sf is None or np is None:
|
||||||
if waveform.ndim == 1:
|
return
|
||||||
waveform = waveform[:, None]
|
data, sr = sf.read(str(output), dtype="float32", always_2d=False)
|
||||||
waveform = np.clip(waveform, -1.0, 1.0)
|
if sr == OUTPUT_SAMPLE_RATE and data.ndim == 1:
|
||||||
pcm = (waveform * 32767.0).astype(np.int16)
|
return
|
||||||
channels = pcm.shape[1]
|
mono = _to_mono_f32(data if data.ndim > 1 else data.reshape(1, -1), sr)
|
||||||
frames = array("h", pcm.reshape(-1).tolist())
|
mono = _resample(mono, sr, OUTPUT_SAMPLE_RATE)
|
||||||
else:
|
sf.write(str(output), mono, OUTPUT_SAMPLE_RATE, subtype="PCM_16")
|
||||||
frames = waveform
|
|
||||||
if channels is None:
|
|
||||||
channels = 1
|
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:
|
with wave.open(str(output), "wb") as wav:
|
||||||
wav.setnchannels(channels)
|
wav.setnchannels(1)
|
||||||
wav.setsampwidth(2)
|
wav.setsampwidth(2)
|
||||||
wav.setframerate(sample_rate)
|
wav.setframerate(sample_rate)
|
||||||
wav.writeframes(frames.tobytes())
|
wav.writeframes(frames.tobytes())
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,63 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Materialize the local highlight models into stable, offline-loadable directories.
|
||||||
|
|
||||||
|
Loads from the existing Hugging Face cache (no network) and re-saves a minimal,
|
||||||
|
self-contained copy under models/. Run once during provisioning.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||||||
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||||
|
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
||||||
|
|
||||||
|
REPO = Path(__file__).resolve().parent.parent
|
||||||
|
MODELS = REPO / "models"
|
||||||
|
|
||||||
|
|
||||||
|
def provision_musicgen() -> None:
|
||||||
|
import torch
|
||||||
|
from transformers import MusicgenForConditionalGeneration, AutoProcessor
|
||||||
|
|
||||||
|
dest = MODELS / "musicgen-small"
|
||||||
|
if (dest / "config.json").is_file():
|
||||||
|
print(f"musicgen: already materialized at {dest}")
|
||||||
|
return
|
||||||
|
print("musicgen: loading from cache...")
|
||||||
|
processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
|
||||||
|
model = MusicgenForConditionalGeneration.from_pretrained(
|
||||||
|
"facebook/musicgen-small", torch_dtype=torch.float32
|
||||||
|
)
|
||||||
|
dest.mkdir(parents=True, exist_ok=True)
|
||||||
|
model.save_pretrained(dest)
|
||||||
|
processor.save_pretrained(dest)
|
||||||
|
print(f"musicgen: saved -> {dest}")
|
||||||
|
|
||||||
|
|
||||||
|
def provision_audioldm2() -> None:
|
||||||
|
import torch
|
||||||
|
from diffusers import AudioLDM2Pipeline
|
||||||
|
|
||||||
|
dest = MODELS / "audioldm2"
|
||||||
|
if (dest / "model_index.json").is_file():
|
||||||
|
print(f"audioldm2: already materialized at {dest}")
|
||||||
|
return
|
||||||
|
print("audioldm2: loading from cache...")
|
||||||
|
pipe = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2", torch_dtype=torch.float32)
|
||||||
|
dest.mkdir(parents=True, exist_ok=True)
|
||||||
|
pipe.save_pretrained(dest)
|
||||||
|
print(f"audioldm2: saved -> {dest}")
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
provision_musicgen()
|
||||||
|
provision_audioldm2()
|
||||||
|
print("done")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
Loading…
Reference in New Issue