Update app.py

app.py

@@ -10,6 +10,10 @@ os.system("aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://hugging
 os.system("aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/sail-rvc/yoimiya-jp/resolve/main/model.pth -d ./weights/yoimiya -o yoimiya.pth")
 os.system("aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/sail-rvc/yoimiya-jp/resolve/main/model.index -d ./weights/yoimiya -o yoimiya.index")
 
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+rvc_models_dir = os.path.join(BASE_DIR, 'weights')
+
 
 
 import edge_tts
@@ -232,6 +236,68 @@ def tts(
         return info, None, None
 
 
+
+def extract_zip(extraction_folder, zip_name):
+    os.makedirs(extraction_folder)
+    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 gr.Error(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')
+
+    # move model and index file to 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 any 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_online_model(url, dir_name, progress=gr.Progress()):
+    try:
+        progress(0, desc=f'[~] Downloading voice model with name {dir_name}...')
+        zip_name = url.split('/')[-1]
+        extraction_folder = os.path.join(rvc_models_dir, dir_name)
+        if os.path.exists(extraction_folder):
+            raise gr.Error(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')
+
+        if 'huggingface.co' in url:
+            urllib.request.urlretrieve(url, zip_name)
+
+        if 'pixeldrain.com' in url:
+            zip_name = dir_name + '.zip'
+            url = f'https://pixeldrain.com/api/file/{zip_name}'
+            urllib.request.urlretrieve(url, zip_name)
+
+        elif 'drive.google.com' in url:
+            # Extract the Google Drive file ID
+            zip_name = dir_name + '.zip'
+            file_id = url.split('/')[-2]
+            output = os.path.join('.', f'{dir_name}.zip')  # Adjust the output path if needed
+            gdown.download(id=file_id, output=output, quiet=False)
+
+        progress(0.5, desc='[~] Extracting zip...')
+        extract_zip(extraction_folder, zip_name)
+        return f'[+] {dir_name} Model successfully downloaded!'
+
+    except Exception as e:
+        raise gr.Error(str(e))
+
+  
+
 initial_md = """
 # RVC text-to-speech webui
 
@@ -326,5 +392,49 @@ with app:
             inputs=[tts_text, tts_voice],
         )
 
+        with gr.Tab('Download model'):
+
+            with gr.Accordion('From HuggingFace/Pixeldrain URL', open=True):
+                with gr.Row():
+                    model_zip_link = gr.Text(label='Download link to model', info='Should be a zip file containing a .pth model file and an optional .index file.')
+                    model_name = gr.Text(label='Name your model', info='Give your new model a unique name from your other voice models.')
+
+                with gr.Row():
+                    download_btn = gr.Button('Download', variant='primary', scale=19)
+                    dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)
+
+                download_btn.click(download_online_model, inputs=[model_zip_link, model_name], outputs=dl_output_message)
+
+                gr.Markdown('## Input Examples',)
+                gr.Examples(
+                    [
+                        ['https://huggingface.co/phant0m4r/LiSA/resolve/main/LiSA.zip', 'Lisa'],
+                        ['https://huggingface.co/Hev832/rvc/resolve/main/Sonic.zip?download=true', 'Sonic'],
+                        ['https://huggingface.co/jkhgf/SLWooly/resolve/main/Jax.zip', 'Jax']
+                    ],
+                    [model_zip_link, model_name],
+                    [],
+                    download_online_model,
+                )
+
+            with gr.Accordion('From Public Index', open=False):
+
+                gr.Markdown('## How to use')
+                gr.Markdown('- Click Initialize public models table')
+                gr.Markdown('- Filter models using tags or search bar')
+                gr.Markdown('- Select a row to autofill the download link and model name')
+                gr.Markdown('- Click Download')
+
+                with gr.Row():
+                    pub_zip_link = gr.Text(label='Download link to model')
+                    pub_model_name = gr.Text(label='Model name')
+
+                with gr.Row():
+                    download_pub_btn = gr.Button('Download', variant='primary', scale=19)
+                    pub_dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)
+
+
+
+
 
 app.launch()