#!/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())