import os
import gradio as gr
from huggingface_hub import hf_hub_download, upload_file, HfApi

TOKEN = os.environ.get("HF_TOKEN", "")
DEST_REPO = "Joker5514/models"
DEST_TYPE = "dataset"

def copy_models():
    api = HfApi(token=TOKEN)
    results = []
    
    # 1. RNNoise ONNX
    try:
        path = hf_hub_download(
            repo_id="niobures/RNNoise",
            filename="models/rnnoise-models/beguiling-drafter-2018-08-30/beguiling-drafter-2018-08-30.onnx",
            repo_type="model",
            token=TOKEN
        )
        api.upload_file(
            path_or_fileobj=path,
            path_in_repo="rnnoise_suppressor.onnx",
            repo_id=DEST_REPO,
            repo_type=DEST_TYPE,
            commit_message="Add rnnoise_suppressor.onnx"
        )
        results.append("rnnoise_suppressor.onnx: OK")
    except Exception as e:
        results.append(f"rnnoise_suppressor.onnx: ERROR {e}")
    
    # 2. Demucs v4 ONNX
    try:
        path = hf_hub_download(
            repo_id="arjune123/demucs-onnx",
            filename="htdemucs_6s.onnx",
            repo_type="model",
            token=TOKEN
        )
        api.upload_file(
            path_or_fileobj=path,
            path_in_repo="demucs_v4_quantized.onnx",
            repo_id=DEST_REPO,
            repo_type=DEST_TYPE,
            commit_message="Add demucs_v4_quantized.onnx"
        )
        results.append("demucs_v4_quantized.onnx: OK")
    except Exception as e:
        results.append(f"demucs_v4_quantized.onnx: ERROR {e}")
    
    # 3. BSRNN vocals ONNX
    try:
        path = hf_hub_download(
            repo_id="FrancisRing/StableAvatar",
            filename="Kim_Vocal_2.onnx",
            repo_type="model",
            token=TOKEN
        )
        api.upload_file(
            path_or_fileobj=path,
            path_in_repo="bsrnn_vocals.onnx",
            repo_id=DEST_REPO,
            repo_type=DEST_TYPE,
            commit_message="Add bsrnn_vocals.onnx"
        )
        results.append("bsrnn_vocals.onnx: OK")
    except Exception as e:
        results.append(f"bsrnn_vocals.onnx: ERROR {e}")
    
    return "\n".join(results)

with gr.Blocks() as demo:
    gr.Markdown("# ONNX Model Uploader")
    btn = gr.Button("Copy ONNX Models to Dataset")
    output = gr.Textbox(label="Results")
    btn.click(copy_models, outputs=output)

demo.launch()