Update app.py

app.py

@@ -1,321 +1,76 @@
-import os
 import gradio as gr
-import spaces
-from infer_rvc_python import BaseLoader
-import random
-import logging
-import time
-import soundfile as sf
-from infer_rvc_python.main import download_manager
-import zipfile
-import librosa
-import traceback
-import soundfile as sf
-from pedalboard import Pedalboard, Reverb, Compressor, HighpassFilter
-from pedalboard.io import AudioFile
-from pydub import AudioSegment
-import noisereduce as nr
+import os
 import numpy as np
-import urllib.request
-import shutil
-import threading
-
-logging.getLogger("infer_rvc_python").setLevel(logging.ERROR)
-
-# Ensure the correct path to the models directory
-model_dir = os.path.join(os.path.dirname(__file__), "models")
+import torch
+from infer_rvc_python import BaseLoader
 
+# Initialize the audio transformation model, but don't configure it yet
 converter = BaseLoader(only_cpu=False, hubert_path=None, rmvpe_path=None)
 
-title = "<center><strong><font size='7'>Vodex AI</font></strong></center>"
-theme = "aliabid94/new-theme"
-
-def find_files(directory):
-    file_paths = []
-    for filename in os.listdir(directory):
-        if filename.endswith('.pth') or filename.endswith('.zip') or filename.endswith('.index'):
-            file_paths.append(os.path.join(directory, filename))
-    return file_paths
-
-def unzip_in_folder(my_zip, my_dir):
-    with zipfile.ZipFile(my_zip) as zip:
-        for zip_info in zip.infolist():
-            if zip_info.is_dir():
-                continue
-            zip_info.filename = os.path.basename(zip_info.filename)
-            zip.extract(zip_info, my_dir)
-
-def find_my_model(a_, b_):
-    if a_ is None or a_.endswith(".pth"):
-        return a_, b_
-
-    txt_files = []
-    for base_file in [a_, b_]:
-        if base_file is not None and base_file.endswith(".txt"):
-            txt_files.append(base_file)
-
-    directory = os.path.dirname(a_)
-
-    for txt in txt_files:
-        with open(txt, 'r') as file:
-            first_line = file.readline()
-
-        download_manager(
-            url=first_line.strip(),
-            path=directory,
-            extension="",
-        )
-
-    for f in find_files(directory):
-        if f.endswith(".zip"):
-            unzip_in_folder(f, directory)
-
-    model = None
-    index = None
-    end_files = find_files(directory)
-
-    for ff in end_files:
-        if ff.endswith(".pth"):
-            model = os.path.join(directory, ff)
-            gr.Info(f"Model found: {ff}")
-        if ff.endswith(".index"):
-            index = os.path.join(directory, ff)
-            gr.Info(f"Index found: {ff}")
-
-    if not model:
-        gr.Error(f"Model not found in: {end_files}")
-
-    if not index:
-        gr.Warning("Index not found")
-
-    return model, index
-
-def get_file_size(url):
-    if "huggingface" not in url:
-        raise ValueError("Only downloads from Hugging Face are allowed")
-
-    try:
-        with urllib.request.urlopen(url) as response:
-            info = response.info()
-            content_length = info.get("Content-Length")
-
-        file_size = int(content_length)
-        if file_size > 500000000:
-            raise ValueError("The file is too large. You can only download files up to 500 MB in size.")
-
-    except Exception as e:
-        raise e
-
-def clear_files(directory):
-    time.sleep(15)
-    print(f"Clearing files: {directory}.")
-    shutil.rmtree(directory)
-
-def get_my_model(url_data):
-    if not url_data:
-        return None, None
-
-    if "," in url_data:
-        a_, b_ = url_data.split()
-        a_, b_ = a_.strip().replace("/blob/", "/resolve/"), b_.strip().replace("/blob/", "/resolve/")
-    else:
-        a_, b_ = url_data.strip().replace("/blob/", "/resolve/"), None
-
-    out_dir = "downloads"
-    folder_download = str(random.randint(1000, 9999))
-    directory = os.path.join(out_dir, folder_download)
-    os.makedirs(directory, exist_ok=True)
-
-    try:
-        get_file_size(a_)
-        if b_:
-            get_file_size(b_)
-
-        valid_url = [a_] if not b_ else [a_, b_]
-        for link in valid_url:
-            download_manager(
-                url=link,
-                path=directory,
-                extension="",
-            )
-
-        for f in find_files(directory):
-            if f.endswith(".zip"):
-                unzip_in_folder(f, directory)
-
-        model = None
-        index = None
-        end_files = find_files(directory)
-
-        for ff in end_files:
-            if ff.endswith(".pth"):
-                model = ff
-                gr.Info(f"Model found: {ff}")
-            if ff.endswith(".index"):
-                index = ff
-                gr.Info(f"Index found: {ff}")
-
-        if not model:
-            raise ValueError(f"Model not found in: {end_files}")
-
-        if not index:
-            gr.Warning("Index not found")
-        else:
-            index = os.path.abspath(index)
-
-        return os.path.abspath(model), index
-
-    except Exception as e:
-        raise e
-    finally:
-        t = threading.Thread(target=clear_files, args=(directory,))
-        t.start()
-
-def convert_now(audio_files, random_tag, converter):
-    return converter(
-        audio_files,
-        random_tag,
-        overwrite=False,
-        parallel_workers=8
-    )
-
-def apply_noisereduce(audio_list):
-    print("Applying noise reduction")
-
-    result = []
-    for audio_path in audio_list:
-        out_path = f'{os.path.splitext(audio_path)[0]}_noisereduce.wav'
-
-        try:
-            # Load audio file
-            audio = AudioSegment.from_file(audio_path)
-
-            # Convert audio to numpy array
-            samples = np.array(audio.get_array_of_samples())
-
-            reduced_noise = nr.reduce_noise(samples, sr=audio.frame_rate, prop_decrease=0.6)
-
-            reduced_audio = AudioSegment(
-                reduced_noise.tobytes(),
-                frame_rate=audio.frame_rate,
-                sample_width=audio.sample_width,
-                channels=audio.channels
-            )
-
-            reduced_audio.export(out_path, format="wav")
-            result.append(out_path)
-
-        except Exception as e:
-            traceback.print_exc()
-            print(f"Error in noise reduction: {str(e)}")
-            result.append(audio_path)
-
-    return result
-
-def run(audio_files, file_m, file_index):
-    if not audio_files:
-        raise ValueError("Please provide an audio file.")
-
-    if isinstance(audio_files, str):
-        audio_files = [audio_files]
-
-    try:
-        duration_base = librosa.get_duration(filename=audio_files[0])
-        print("Duration:", duration_base)
-    except Exception as e:
-        print(e)
-
-    file_m = os.path.join(model_dir, file_m)
-    file_index = os.path.join(model_dir, file_index) if file_index else None
-
-    random_tag = "USER_" + str(random.randint(10000000, 99999999))
+# Path to the models directory
+model_dir = "./models"
 
+# Function to configure the model based on user selection
+def configure_model(file_model, file_index):
+    model_path = os.path.join(model_dir, file_model)
+    index_path = os.path.join(model_dir, file_index) if file_index else None
+    
     converter.apply_conf(
-        tag=random_tag,
-        file_model=file_m,
+        tag="live_transform",
+        file_model=model_path,
         pitch_algo="rmvpe+",
         pitch_lvl=0,
-        file_index=file_index,
+        file_index=index_path,
         index_influence=0.75,
         respiration_median_filtering=3,
         envelope_ratio=0.25,
         consonant_breath_protection=0.5,
-        resample_sr=44100 if audio_files[0].endswith('.mp3') else 0,
+        resample_sr=44100
     )
-    time.sleep(0.1)
-
-    result = convert_now(audio_files, random_tag, converter)
-    result = apply_noisereduce(result)
-
-    return result[0]  # Return the first (and only) file path
-
+    return "Model configured successfully."
+
+# Function to process each audio chunk
+def transform_audio_chunk(audio, instream):
+    if audio is None:
+        return None, instream
+    
+    audio_data = torch.tensor(audio[1], dtype=torch.float32).unsqueeze(0)  # Prepare audio for processing
+    with torch.no_grad():
+        transformed_audio, _ = converter.generate_from_cache(
+            audio_data=(audio[0], audio_data.numpy()),
+            tag="live_transform",
+        )
 
-def process_audio(audio_file, file_m, file_index):
-    if audio_file is not None:
-        result = run([audio_file], file_m, file_index)
+    # Update the stream by concatenating the new transformed audio chunk
+    if instream is None:
+        return (audio[0], transformed_audio.squeeze(0).numpy()), (audio[0], transformed_audio.squeeze(0).numpy())
     else:
-        raise ValueError("No audio recorded.")
-
-    # Return the processed audio file for playback and download
-    return gr.update(value=result, visible=True), result
+        new_stream = np.concatenate((instream[1], transformed_audio.squeeze(0).numpy()))
+        return (audio[0], new_stream), (audio[0], new_stream)
 
-
-def model_conf():
+# Gradio interface setup
+with gr.Blocks() as demo:
+    # Get the list of available model and index files
     model_files = [f for f in os.listdir(model_dir) if f.endswith(".pth")]
-    return gr.Dropdown(
-        label="Select Model File",
-        choices=model_files,
-        value=model_files[0] if model_files else None,
-        interactive=True,
-    )
-
-def index_conf():
     index_files = [f for f in os.listdir(model_dir) if f.endswith(".index")]
-    return gr.Dropdown(
-        label="Select Index File",
-        choices=index_files,
-        value=index_files[0] if index_files else None,
-        interactive=True,
-    )
-
-def button_conf():
-    return gr.Button(
-        "Inference",
-        variant="primary",
-    )
-
-
-def get_gui(theme):
-    with gr.Blocks(theme=theme, delete_cache=(3200, 3200)) as app:
-        gr.Markdown(title)
-
-        # Only keep the microphone input option
-        audio = gr.Audio(label="Record Audio", sources="microphone", type="filepath", visible=True)
-        
-        model = model_conf()
-        indx = index_conf()
-        button_base = button_conf()
-
-        output_audio = gr.Audio(label="Play Processed Audio", visible=False, show_share_button=False)
-        output_files = gr.File(label="Download Processed Audio", interactive=False)
-
-        button_base.click(
-            process_audio,
-            inputs=[audio, model, indx],  # Removed None, keeping only valid inputs
-            outputs=[output_audio, output_files],  # Only need to output the audio file for playback and download
-        )
-
-    return app
-
+    
+    # Dropdowns for model and index file selection
+    model_file = gr.Dropdown(choices=model_files, label="Select Model File")
+    index_file = gr.Dropdown(choices=index_files, label="Select Index File")
+    configure_button = gr.Button("Configure Model")
+    
+    # Audio input component with streaming enabled
+    inp = gr.Audio(source="microphone", streaming=True, type="numpy")
+    # Audio output component to play back the transformed audio
+    out = gr.Audio(streaming=True)
+    # State to manage the ongoing audio stream
+    stream = gr.State()
+
+    # Link the input to the processing function and output
+    inp.stream(transform_audio_chunk, [inp, stream], [out, stream])
+
+    # Link the model configuration button to the configure_model function
+    configure_button.click(configure_model, inputs=[model_file, index_file], outputs=[])
 
 if __name__ == "__main__":
-    app = get_gui(theme)
-    app.queue(default_concurrency_limit=40)
-    app.launch(
-        max_threads=40,
-        share=False,
-        show_error=True,
-        quiet=False,
-        debug=False,
-        allowed_paths=["./downloads/"],
-    )
+    demo.launch()