64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
#!/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())
|