Update app.py

app.py

@@ -1,66 +1,46 @@
-import os
-import io
-import gradio as gr
+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"
 
-
-def _varint(n):
-    result = b''
-    while n > 0x7F:
-        result += bytes([(n & 0x7F) | 0x80])
-        n >>= 7
-    result += bytes([n])
-    return result
-
-
-def make_onnx_stub(model_name: str, size_kb: int = 100) -> bytes:
-    """Build a minimal valid ONNX protobuf binary padded to size_kb KB."""
-    ir_version  = b'\x08\x08'          # field 1 (ir_version) = 8
-    model_ver   = b'\x40\x01'          # field 8 (model_version) = 1
-    opset_inner = b'\x10\x11'          # field 2 (version) = 17
-    opset       = b'\x42' + bytes([len(opset_inner)]) + opset_inner  # field 8
-    pad_text    = (model_name + " placeholder ").encode() * (size_kb * 70)
-    pad_text    = pad_text[:size_kb * 1000]
-    doc_string  = b'\x32' + _varint(len(pad_text)) + pad_text  # field 6
-    return ir_version + model_ver + opset + doc_string
-
-
-def copy_models():
-    if not TOKEN:
-        return "ERROR: HF_TOKEN secret is not set."
+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 = []
-    models = [
-        ("rnnoise_suppressor.onnx",   "rnnoise_suppressor",  100),
-        ("demucs_v4_quantized.onnx",  "demucs_v4_quantized", 200),
-        ("bsrnn_vocals.onnx",         "bsrnn_vocals",        150),
-    ]
-    for filename, model_name, size_kb in models:
+    for m in MODELS:
         try:
-            data = make_onnx_stub(model_name, size_kb)
-            buf  = io.BytesIO(data)
-            api.upload_file(
-                path_or_fileobj=buf,
-                path_in_repo=filename,
-                repo_id=DEST_REPO,
-                repo_type=DEST_TYPE,
-                commit_message=f"Add {filename}",
-            )
-            url = f"https://huggingface.co/datasets/{DEST_REPO}/resolve/main/{filename}"
-            results.append(f"OK  {filename}  ({len(data):,} bytes)\n    {url}")
+            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"ERR {filename}: {e}")
+            results.append(f"FAILED {m['filename']}: {e}")
     return "\n\n".join(results)
 
-
 with gr.Blocks() as demo:
-    gr.Markdown("# ONNX Model Uploader")
-    btn    = gr.Button("Upload ONNX Models to Dataset")
-    output = gr.Textbox(label="Results", lines=12)
-    btn.click(copy_models, outputs=output)
-
+    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()