From da3c5ab43888d5b2413b40d40631351d1c714f48 Mon Sep 17 00:00:00 2001
From: Neo <neoforevershog@gmail.com>
Date: Wed, 18 Dec 2024 08:49:12 +0700
Subject: [PATCH] Update infer.py

---
 rvc_inferpy/infer.py | 59 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/rvc_inferpy/infer.py b/rvc_inferpy/infer.py
index ea9356d..8d0c164 100644
--- a/rvc_inferpy/infer.py
+++ b/rvc_inferpy/infer.py
@@ -11,6 +11,13 @@
 )
 from pathlib import Path
 import requests
+import os
+import zipfile
+import shutil
+import urllib.request
+import gdown
+
+models_dir = "models"
 
 
 class Configs:
@@ -114,6 +121,58 @@ def dl_model(link, model_name, dir_name):
     pass
 
 
+
+def extract_zip(extraction_folder, zip_name):
+    os.makedirs(extraction_folder, exist_ok=True)
+    with zipfile.ZipFile(zip_name, 'r') as zip_ref:
+        zip_ref.extractall(extraction_folder)
+    os.remove(zip_name)
+
+    index_filepath, model_filepath = None, None
+    for root, dirs, files in os.walk(extraction_folder):
+        for name in files:
+            if name.endswith('.index') and os.stat(os.path.join(root, name)).st_size > 1024 * 100:
+                index_filepath = os.path.join(root, name)
+
+            if name.endswith('.pth') and os.stat(os.path.join(root, name)).st_size > 1024 * 1024 * 40:
+                model_filepath = os.path.join(root, name)
+
+    if not model_filepath:
+        raise Exception(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')
+
+    os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
+    if index_filepath:
+        os.rename(index_filepath, os.path.join(extraction_folder, os.path.basename(index_filepath)))
+
+    # Remove unnecessary nested folders
+    for filepath in os.listdir(extraction_folder):
+        if os.path.isdir(os.path.join(extraction_folder, filepath)):
+            shutil.rmtree(os.path.join(extraction_folder, filepath))
+
+def download_rvc_model(url, dir_name):
+    try:
+        print(f'[~] Downloading voice model with name {dir_name}...')
+        zip_name = url.split('/')[-1]
+        extraction_folder = os.path.join(models_dir, dir_name)
+        if os.path.exists(extraction_folder):
+            raise Exception(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')
+
+        if 'pixeldrain.com' in url:
+            url = f'https://pixeldrain.com/api/file/{zip_name}'
+        if 'drive.google.com' in url:
+            zip_name = dir_name + ".zip"
+            gdown.download(url, output=zip_name, use_cookies=True, quiet=True, fuzzy=True)
+        else:
+            urllib.request.urlretrieve(url, zip_name)
+
+        print(f'[~] Extracting zip file...')
+        extract_zip(extraction_folder, zip_name)
+        print(f'[+] {dir_name} Model successfully downloaded!')
+
+    except Exception as e:
+        raise Exception(str(e))
+
+
 def infer_audio(
     model_name,
     audio_path,