Update app.py

app.py

@@ -1,70 +1,83 @@
 import gradio as gr
-import os
-import numpy as np
 import torch
+import numpy as np
+import datetime
 from infer_rvc_python import BaseLoader
+import torchaudio
 
-# Initialize the audio transformation model
-converter = BaseLoader(only_cpu=False, hubert_path=None, rmvpe_path=None)
+# Initialize the model
+now = datetime.datetime.now()
+timestamp = now.strftime("%Y-%m-%d_%H-%M-%S")
+random_tag = "USER_" + str(timestamp)
 
-# Path to the models directory
-model_dir = "./models"
+converter = BaseLoader(only_cpu=False, hubert_path='./hubert_base.pt', rmvpe_path='./rmvpe.pt')
+converter.apply_conf(
+    tag=random_tag,
+    file_model="./model.pth",
+    pitch_algo="rmvpe+",
+    pitch_lvl=0,
+    file_index="./model.index",
+    index_influence=0.80,
+    respiration_median_filtering=3,
+    envelope_ratio=0.25,
+    consonant_breath_protection=0.5,
+    resample_sr=0, 
+)
 
-# 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="live_transform",
-        file_model=model_path,
-        pitch_algo="rmvpe+",
-        pitch_lvl=0,
-        file_index=index_path,
-        index_influence=0.75,
-        respiration_median_filtering=3,
-        envelope_ratio=0.25,
-        consonant_breath_protection=0.5,
-        resample_sr=44100
-    )
-    return "Model configured successfully."
+chunk_sec = 0.1
+sr = 16000
+chunk_len = int(sr * chunk_sec)
+L = 16
+stop_recording = False
+first_output_latency = 0
+
+# Initialize global audio buffer
+audio_buffer = torch.zeros(0, dtype=torch.float32)
 
-# Function to process each audio chunk
-def transform_audio_chunk(audio):
+# Function to process audio chunks
+def process_audio_chunk(audio, buffer_state):
+    global first_output_latency
+    
+    if buffer_state is None:
+        buffer_state = torch.zeros(0, dtype=torch.float32)
+    
     if audio is None:
-        return None
+        return None, buffer_state
+    
+    # Convert input audio to tensor
+    audio_data = torch.tensor(audio[1], dtype=torch.float32)
+    buffer_state = torch.cat((buffer_state, audio_data))
+    
+    if len(buffer_state) < chunk_len:
+        return None, buffer_state
     
-    audio_data = torch.tensor(audio[1], dtype=torch.float32).unsqueeze(0)  # Prepare audio for processing
-    with torch.no_grad():
-        # Ensure that source_sr and base_sr are scalar values
-        transformed_audio, _ = converter.generate_from_cache(
-            audio_data=(audio[0], audio_data.numpy()),
-            tag="live_transform",
+    # Process the chunk
+    previous_chunk = buffer_state[:chunk_len]
+    buffer_state = buffer_state[chunk_len:]
+    
+    input_chunk = torch.cat([torch.zeros(L * 2, dtype=torch.float32), previous_chunk])
+    
+    with torch.inference_mode():
+        data = (input_chunk.numpy().astype(np.int16), sr)
+        result_array, sample_rate = converter.generate_from_cache(
+            audio_data=data,
+            tag=random_tag,
         )
-
-    return audio[0], transformed_audio.squeeze(0).numpy()
+        
+        if first_output_latency == 0:
+            first_output_latency = time.time()
+        
+        output = torch.tensor(result_array, dtype=torch.float32)
+        output = output.squeeze(0).numpy()
+    
+    return (audio[0], output), buffer_state
 
 # 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")]
-    index_files = [f for f in os.listdir(model_dir) if f.endswith(".index")]
+    audio_input = gr.Audio(sources="microphone", streaming=True, type="numpy", label="Input Audio")
+    audio_output = gr.Audio(label="Output Audio")
+    buffer_state = gr.State()
     
-    # 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(sources="microphone", streaming=True, type="numpy")
-    # Audio output component to play back the transformed audio
-    out = gr.Audio()
-    
-    # Link the input to the processing function and output
-    inp.stream(transform_audio_chunk, inputs=[inp], outputs=[out])
-
-    # Link the model configuration button to the configure_model function
-    configure_button.click(configure_model, inputs=[model_file, index_file], outputs=[])
+    audio_input.stream(process_audio_chunk, inputs=[audio_input, buffer_state], outputs=[audio_output, buffer_state])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()