import os, io, requests, gradio as gr
from huggingface_hub import HfApi

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

MODELS = [
    {"filename": "rnnoise_suppressor.onnx",
     "url": "https://huggingface.co/niobures/RNNoise/resolve/main/models/rnnoise.onnx",
     "min_bytes": 100_000},
    {"filename": "demucs_v4_quantized.onnx",
     "url": "https://huggingface.co/MrCitron/demucs-v4-onnx/resolve/main/htdemucs.onnx",
     "fallback": "https://huggingface.co/timcsy/demucs-web-onnx/resolve/main/demucs.onnx",
     "min_bytes": 10_000_000},
    {"filename": "bsrnn_vocals.onnx",
     "url": "https://huggingface.co/facebook/bsrnn/resolve/main/bsrnn_vocals.onnx",
     "fallback": "https://huggingface.co/Randell/bsrnn-onnx/resolve/main/bsrnn_vocals.onnx",
     "min_bytes": 5_000_000},
]

def upload_models():
    if not TOKEN: return "ERROR: HF_TOKEN not set"
    api = HfApi(token=TOKEN)
    results = []
    for m in MODELS:
        try:
            try:
                r = requests.get(m["url"], timeout=300); r.raise_for_status(); data = r.content
                if len(data) < m["min_bytes"]: raise ValueError("too small")
                src = m["url"]
            except:
                r = requests.get(m.get("fallback", m["url"]), timeout=300); r.raise_for_status(); data = r.content
                src = m.get("fallback", m["url"])
            api.upload_file(path_or_fileobj=io.BytesIO(data), path_in_repo=m["filename"],
                repo_id=DEST_REPO, repo_type=DEST_TYPE,
                commit_message=f"Upload real ONNX: {m['filename']} ({len(data):,} bytes)")
            results.append(f"SUCCESS {m['filename']} ({len(data)/1e6:.1f} MB) from {src}")
        except Exception as e:
            results.append(f"FAILED {m['filename']}: {e}")
    return "\n\n".join(results)

with gr.Blocks() as demo:
    gr.Markdown("# ONNX Model Uploader\nDownloads real models and uploads to Joker5514/models")
    gr.Button("Upload Real ONNX Models", variant="primary").click(upload_models, outputs=gr.Textbox(lines=15))
demo.launch()