Create infer-web.py

infer-web.py

@@ -0,0 +1,137 @@
+import os
+import shutil
+import subprocess
+from mega import Mega
+import gradio as gr
+
+def download_from_url(url, model):
+    try:
+        model = model.replace('.pth', '').replace('.index', '').replace('.zip', '')
+        url = url.replace('/blob/main/', '/resolve/main/').strip()
+
+        for directory in ["downloads", "unzips", "zip"]:
+            os.makedirs(directory, exist_ok=True)
+
+        if url.endswith('.pth'):
+            subprocess.run(["wget", url, "-O", f'assets/weights/{model}.pth'])
+        elif url.endswith('.index'):
+            os.makedirs(f'logs/{model}', exist_ok=True)
+            subprocess.run(["wget", url, "-O", f'logs/{model}/added_{model}.index'])
+        elif url.endswith('.zip'):
+            subprocess.run(["wget", url, "-O", f'downloads/{model}.zip'])
+        else:
+            if "drive.google.com" in url:
+                url = url.split('/')[0]
+                subprocess.run(["gdown", url, "--fuzzy", "-O", f'downloads/{model}'])
+            elif "mega.nz" in url:
+                Mega().download_url(url, 'downloads')
+            else:
+                subprocess.run(["wget", url, "-O", f'downloads/{model}'])
+
+        downloaded_file = next((f for f in os.listdir("downloads")), None)
+        if downloaded_file:
+            if downloaded_file.endswith(".zip"):
+                shutil.unpack_archive(f'downloads/{downloaded_file}', "unzips", 'zip')
+                for root, _, files in os.walk('unzips'):
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        if file.endswith(".index"):
+                            os.makedirs(f'logs/{model}', exist_ok=True)
+                            shutil.copy2(file_path, f'logs/{model}')
+                        elif file.endswith(".pth") and "G_" not in file and "D_" not in file:
+                            shutil.copy(file_path, f'assets/weights/{model}.pth')
+            elif downloaded_file.endswith(".pth"):
+                shutil.copy(f'downloads/{downloaded_file}', f'assets/weights/{model}.pth')
+            elif downloaded_file.endswith(".index"):
+                os.makedirs(f'logs/{model}', exist_ok=True)
+                shutil.copy(f'downloads/{downloaded_file}', f'logs/{model}/added_{model}.index')
+            else:
+                return "Failed to download file"
+        return f"Successfully downloaded {model} voice models"
+    except Exception as e:
+        return f"Error: {str(e)}"
+    finally:
+        shutil.rmtree("downloads", ignore_errors=True)
+        shutil.rmtree("unzips", ignore_errors=True)
+        shutil.rmtree("zip", ignore_errors=True)
+
+def listen_to_model(model_path, index_path, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection):
+    if not os.path.exists(model_path):
+        return "Model path not found"
+    if not os.path.exists(index_path):
+        return f"{index_path} was not found"
+    if not os.path.exists(input_path):
+        return f"{input_path} was not found"
+    
+    os.environ['index_root'] = os.path.dirname(index_path)
+    index_path = os.path.basename(index_path)
+    model_name = os.path.basename(model_path)
+    os.environ['weight_root'] = os.path.dirname(model_path)
+
+    try:
+        command = [
+            "python", "tools/infer_cli.py",
+            "--f0up_key", str(pitch),
+            "--input_path", input_path,
+            "--index_path", index_path,
+            "--f0method", f0_method,
+            "--opt_path", save_as,
+            "--model_name", model_name,
+            "--index_rate", str(index_rate),
+            "--device", "cuda:0",
+            "--is_half", "True",
+            "--filter_radius", "3",
+            "--resample_sr", "0",
+            "--rms_mix_rate", str(volume_normalization),
+            "--protect", str(consonant_protection)
+        ]
+        subprocess.run(command, check=True)
+        return save_as
+    except subprocess.CalledProcessError as e:
+        return f"Error: {str(e)}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# RVC V2 Web UI")
+
+    with gr.Tabs():
+        with gr.TabItem("Download Model"):
+            gr.Markdown("### Download RVC Model")
+            url_input = gr.Textbox(label="Model URL", placeholder="Enter the model URL here")
+            model_input = gr.Textbox(label="Model Name", placeholder="Enter the model name here")
+            download_button = gr.Button("Download")
+            download_output = gr.Textbox(label="Download Status")
+            download_button.click(download_from_url, inputs=[url_input, model_input], outputs=download_output)
+
+        with gr.TabItem("Listen to Model"):
+            gr.Markdown("### Listen to Your Model")
+            model_path_input = gr.Textbox(label="Model Path", value="/content/RVC/assets/weights/Sonic.pth")
+            index_path_input = gr.Textbox(label="Index Path", value="/content/RVC/logs/Sonic/added_IVF905_Flat_nprobe_1_Sonic_v2.index")
+            input_path_input = gr.Textbox(label="Input Audio Path", value="/content/RVC/audios/astronauts.mp3")
+            save_as_input = gr.Textbox(label="Save Output As", value="/content/RVC/audios/cli_output.wav")
+            f0_method_input = gr.Radio(label="F0 Method", choices=["rmvpe", "pm", "harvest"], value="rmvpe")
+
+            with gr.Row():
+                pitch_input = gr.Slider(label="Pitch", minimum=-12, maximum=12, step=1, value=0)
+                index_rate_input = gr.Slider(label="Index Rate", minimum=0, maximum=1, step=0.01, value=0.5)
+                volume_normalization_input = gr.Slider(label="Volume Normalization", minimum=0, maximum=1, step=0.01, value=0)
+                consonant_protection_input = gr.Slider(label="Consonant Protection", minimum=0, maximum=1, step=0.01, value=0.5)
+
+            listen_button = gr.Button("Generate and Listen")
+            audio_output = gr.Audio(label="Output Audio")
+            listen_button.click(
+                listen_to_model, 
+                inputs=[
+                    model_path_input, 
+                    index_path_input, 
+                    pitch_input, 
+                    input_path_input, 
+                    f0_method_input, 
+                    save_as_input, 
+                    index_rate_input, 
+                    volume_normalization_input, 
+                    consonant_protection_input
+                ], 
+                outputs=audio_output
+            )
+
+demo.launch()