Update utils.py

utils.py

@@ -1,48 +1,102 @@
-# utils.py
+# main.py
 
 import os
-import logging
-import requests
+from pathlib import Path
+from gui import get_gui
+from utils import download_manager
+from mdx_core import MDX
+import json
 
+MDX_DOWNLOAD_LINK = "https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/"
+UVR_MODELS = [
+    "UVR-MDX-NET-Voc_FT.onnx",
+    "UVR_MDXNET_KARA_2.onnx",
+    "Reverb_HQ_By_FoxJoy.onnx",
+    "UVR-MDX-NET-Inst_HQ_4.onnx",
+]
 
-# Basic logger setup
-logger = logging.getLogger("audio_splitter")
-logger.setLevel(logging.INFO)
-ch = logging.StreamHandler()
-ch.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
-logger.addHandler(ch)
+BASE_DIR = Path(__file__).resolve().parent
+mdxnet_models_dir = BASE_DIR / "mdx_models"
 
+# Ensure mdx_models is a directory
+if mdxnet_models_dir.exists() and not mdxnet_models_dir.is_dir():
+    mdxnet_models_dir.unlink()
+mdxnet_models_dir.mkdir(parents=True, exist_ok=True)
 
-def create_directories(path):
-    os.makedirs(path, exist_ok=True)
 
+# 🔧 Auto-generate fallback data.json if needed
+def generate_data_json(mdxnet_models_dir):
+    presets = {
+        "UVR-MDX-NET-Voc_FT.onnx": {
+            "mdx_dim_f_set": 2048,
+            "mdx_dim_t_set": 3,
+            "mdx_n_fft_scale_set": 6144,
+            "primary_stem": "Vocals",
+            "compensate": 1.035
+        },
+        "UVR_MDXNET_KARA_2.onnx": {
+            "mdx_dim_f_set": 1024,
+            "mdx_dim_t_set": 3,
+            "mdx_n_fft_scale_set": 4096,
+            "primary_stem": "Main",
+            "compensate": 1.035
+        },
+        "UVR-MDX-NET-Inst_HQ_4.onnx": {
+            "mdx_dim_f_set": 2048,
+            "mdx_dim_t_set": 3,
+            "mdx_n_fft_scale_set": 6144,
+            "primary_stem": "Instrumental",
+            "compensate": 1.0
+        },
+        "Reverb_HQ_By_FoxJoy.onnx": {
+            "mdx_dim_f_set": 2048,
+            "mdx_dim_t_set": 3,
+            "mdx_n_fft_scale_set": 6144,
+            "primary_stem": "Vocals",
+            "compensate": 1.035
+        },
+    }
 
-def remove_directory_contents(path):
-    if not os.path.exists(path):
-        return
-    for filename in os.listdir(path):
-        file_path = os.path.join(path, filename)
-        if os.path.isfile(file_path):
-            os.remove(file_path)
-        elif os.path.isdir(file_path):
-            import shutil
-            shutil.rmtree(file_path)
+    data = {}
+    for filename, params in presets.items():
+        path = mdxnet_models_dir / filename
+        if path.exists():
+            model_hash = MDX.get_hash(str(path))
+            data[model_hash] = params
 
+    with open(mdxnet_models_dir / "data.json", "w") as f:
+        json.dump(data, f, indent=2)
+    print("✅ Auto-generated data.json with", len(data), "entries.")
 
-def download_manager(url, output_dir):
-    filename = os.path.basename(url)
-    output_path = os.path.join(output_dir, filename)
-    if os.path.exists(output_path):
-        return output_path
 
-    logger.info(f"Downloading: {filename}")
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)  # Ensure directory exists
+# Check for data.json and create if missing
+json_path = mdxnet_models_dir / "data.json"
+print("🔍 Checking for data.json at:", json_path)
+print("📁 mdx_models contents:", list(mdxnet_models_dir.iterdir()))
+print("📂 Current working directory:", os.getcwd())
 
-    with requests.get(url, stream=True) as r:
-        r.raise_for_status()
-        with open(output_path, 'wb') as f:
-            for chunk in r.iter_content(chunk_size=8192):
-                f.write(chunk)
+if not json_path.exists():
+    print("⚠️ data.json missing, generating fallback.")
+    generate_data_json(mdxnet_models_dir)
+else:
+    print("✅ data.json found.")
 
-    logger.info(f"Downloaded: {filename}")
-    return output_path
+
+def download_models():
+    for model in UVR_MODELS:
+        url = os.path.join(MDX_DOWNLOAD_LINK, model)
+        download_manager(url, str(mdxnet_models_dir))
+
+
+if __name__ == "__main__":
+    download_models()
+    theme = "NoCrypt/miku"
+    app = get_gui(theme)
+    app.queue(default_concurrency_limit=40)
+    app.launch(
+        max_threads=40,
+        share=False,
+        show_error=True,
+        quiet=False,
+        debug=False,
+    )