#!/usr/bin/env python3
# convert.py : safetensors to GGUF for ACE-Step captioner and transcriber
# Reads from ./<model>/, writes GGUF to ./models/
# Produces two GGUF per model : the text LM (thinker) and the mmproj (vision + audio towers)

import os
import sys
import subprocess

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "models")
CONVERTER  = "/mnt/workspace/llama.cpp/convert_hf_to_gguf.py"

MODELS = [
    "acestep-captioner",
    "acestep-transcriber",
]

def log(tag, msg):
    print("[%s] %s" % (tag, msg), file=sys.stderr, flush=True)

def run(name, src, outfile, mmproj):
    if os.path.exists(outfile):
        log("Skip", os.path.basename(outfile))
        return
    cmd = [
        sys.executable, CONVERTER,
        src,
        "--outtype", "bf16",
        "--outfile", outfile,
    ]
    if mmproj:
        cmd.append("--mmproj")
    log("Convert", os.path.basename(outfile))
    subprocess.run(cmd, check=True, cwd=SCRIPT_DIR)

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for name in MODELS:
        src = os.path.join(SCRIPT_DIR, name)
        if not os.path.isdir(src):
            log("Miss", "source dir not found: " + src)
            continue
        lm_out     = os.path.join(OUTPUT_DIR, "%s-BF16.gguf" % name)
        mmproj_out = os.path.join(OUTPUT_DIR, "mmproj-%s-BF16.gguf" % name)
        run(name, src, lm_out,     mmproj=False)
        run(name, src, mmproj_out, mmproj=True)
    log("Done", "all conversions finished")

if __name__ == "__main__":
    main()