from sonique import get_pretrained_model
from sonique.interface.gradio import create_ui
import json 
from huggingface_hub import login
import torch
import os

login(token=os.getenv('HF_TOKEN'))

interface = create_ui(
    model_config_path = str(cached_path('https://raw.githubusercontent.com/zxxwxyyy/sonique/refs/heads/main/best_model.json')),
    ckpt_path=str(cached_path('hf://mrfakename/SONIQUE/stable_ep=220.ckpt')),
    # pretrained_name=args.pretrained_name,
    pretransform_ckpt_path=None
)
interface.queue().launch()


    

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Run gradio interface')
    parser.add_argument('--pretrained-name', type=str, help='Name of pretrained model', required=False)
    parser.add_argument('--model-config', type=str, help='Path to model config', required=False)
    parser.add_argument('--ckpt-path', type=str, help='Path to model checkpoint', required=False)
    parser.add_argument('--pretransform-ckpt-path', type=str, help='Optional to model pretransform checkpoint', required=False)
    parser.add_argument('--username', type=str, help='Gradio username', required=False)
    parser.add_argument('--password', type=str, help='Gradio password', required=False)
    args = parser.parse_args()
    main(args)