import torch
from omnivoice import OmniVoice
from omnivoice.cli.demo import build_demo, get_best_device

def main():
    device = get_best_device()
    checkpoint = "k2-fsa/OmniVoice"
    # Load model
    model = OmniVoice.from_pretrained(
        checkpoint, 
        device_map=device, 
        dtype=torch.float16, 
        load_asr=True
    )
    # Build demo
    demo = build_demo(model, checkpoint)
    return demo

demo = main()

if __name__ == "__main__":
    demo.launch()