import argparse
import torch
import importlib.util
import os
import shutil
from whisper import load_model

# 动态加载用户上传的转换器
# 注意：前提是你必须将 'models/convert-whisper-to-coreml.py' 上传到 Space 的 models 文件夹下
spec = importlib.util.spec_from_file_location('whisper_to_coreml', 'models/convert-whisper-to-coreml.py')
whisper_to_coreml = importlib.util.module_from_spec(spec)
spec.loader.exec_module(whisper_to_coreml)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, required=True, help="Path to the intermediate .pt file")
    parser.add_argument("--model-name", type=str, default="model")
    parser.add_argument("--output-dir", type=str, required=True)
    parser.add_argument("--quantize", action="store_true")
    args = parser.parse_args()

    print(f"Loading model from {args.model_path}")
    whisper = load_model(args.model_path).cpu()
    hparams = whisper.dims
    print(hparams)

    encoder = whisper.encoder
    decoder = whisper.decoder

    print("Converting Encoder...")
    encoder = whisper_to_coreml.convert_encoder(hparams, encoder, quantize=args.quantize)
    encoder.save(os.path.join(args.output_dir, f"coreml-encoder-{args.model_name}.mlpackage"))

    print("Converting Decoder...")
    decoder = whisper_to_coreml.convert_decoder(hparams, decoder, quantize=args.quantize)
    decoder.save(os.path.join(args.output_dir, f"coreml-decoder-{args.model_name}.mlpackage"))
    
    print("CoreML conversion done.")