feat!: Add WAV file upload functionality and remove YouTube download option

app.py

@@ -23,25 +23,20 @@ def delete_input_files(input_dir):
         wav_file.unlink()
         print(f"Deleted {wav_file}")
 
-def process_uploaded_audio(file, state=True):
-    if state:
-        delete_input_files(INPUT_FOLDER)
-
-    sanitized_filename = sanitize_filename(file.name)
-    input_path = Path(INPUT_FOLDER) / "wav" / sanitized_filename
-    input_path.parent.mkdir(parents=True, exist_ok=True)
-
-    # Check if the file object has a `read` method, otherwise handle it differently.
-    if hasattr(file, 'read'):
-        file_content = file.read()
-    else:
-        # Assuming `file` is a NamedString or similar, with the content directly as a string
-        file_content = file.value  # or file['value'] if it's a dictionary-like object
+def handle_file_upload(file):
+    if file is None:
+        return None, "No file uploaded"
 
-    with open(input_path, 'wb') as f:
-        f.write(file_content)
-
-    return str(input_path)
+    filename = Path(file.name).stem
+    formatted_title = sanitize_filename(filename)
+    
+    input_path = os.path.join(INPUT_FOLDER, "wav", f"{formatted_title}.wav")
+    os.makedirs(os.path.dirname(input_path), exist_ok=True)
+    
+    with open(input_path, "wb") as f:
+        f.write(file.read())
+    
+    return input_path, formatted_title
 
 def run_inference(model_type, config_path, start_check_point, input_dir, output_dir, device_ids="0"):
     command = [
@@ -97,6 +92,8 @@ def move_stems_to_parent(input_dir):
                 new_instrumental_path = os.path.join(parent_dir, "instrumental.wav")
                 print(f"Moving {instrumental_path} to {new_instrumental_path}")
                 shutil.move(instrumental_path, new_instrumental_path)
+            else:
+                print(f"Instrumental file not found: {instrumental_path}")
 
 def combine_stems_for_all(input_dir):
     for subdir, _, _ in os.walk(input_dir):
@@ -155,12 +152,14 @@ def delete_folders_and_files(input_dir):
 
 def process_audio(uploaded_file):
     try:
-        yield "Processing audio file...", None
+        yield "Processing audio...", None
         
-        if uploaded_file is None:
+        if uploaded_file:
+            input_path, formatted_title = handle_file_upload(uploaded_file)
+            if input_path is None:
+                raise ValueError("File upload failed.")
+        else:
             raise ValueError("Please upload a WAV file.")
-        
-        file_path = process_uploaded_audio(uploaded_file, False)
 
         yield "Starting SCNet inference...", None
         proc_folder_direct("scnet", "configs/config_scnet_other.yaml", "results/model_scnet_other.ckpt", f"{INPUT_FOLDER}/wav", OUTPUT_FOLDER)
@@ -171,13 +170,13 @@ def process_audio(uploaded_file):
         yield "Starting HTDemucs inference...", None
         proc_folder_direct("htdemucs", "configs/config_htdemucs_bass.yaml", "results/model_htdemucs_bass.th", f"{INPUT_FOLDER}/wav", OUTPUT_FOLDER)
 
-        source_path = f'{OUTPUT_FOLDER}/{file_path.stem}/mel_band_roformer/{file_path.stem}_instrumental.wav'
-        destination_path = f'{OUTPUT_FOLDER}/{file_path.stem}/mel_band_roformer/{file_path.stem}.wav'
-    
+        source_path = f'{OUTPUT_FOLDER}{formatted_title}/mel_band_roformer/{formatted_title}_instrumental.wav'
+        destination_path = f'{OUTPUT_FOLDER}{formatted_title}/mel_band_roformer/{formatted_title}.wav'
+
         os.rename(source_path, destination_path)
-        
+
         yield "Starting BS Roformer inference...", None
-        proc_folder_direct("bs_roformer", "configs/config_bs_roformer_instrumental.yaml", "results/model_bs_roformer_instrumental.ckpt", f'{OUTPUT_FOLDER}/{file_path.stem}/mel_band_roformer', OUTPUT_FOLDER)
+        proc_folder_direct("bs_roformer", "configs/config_bs_roformer_instrumental.yaml", "results/model_bs_roformer_instrumental.ckpt", f'{OUTPUT_FOLDER}{formatted_title}/mel_band_roformer', OUTPUT_FOLDER)
 
         yield "Moving input files...", None
         delete_input_files(INPUT_FOLDER)
@@ -191,27 +190,35 @@ def process_audio(uploaded_file):
         yield "Cleaning up...", None
         delete_folders_and_files(OUTPUT_FOLDER)
 
-        yield f"Audio processing completed successfully.", f'{OUTPUT_FOLDER}/{file_path.stem}/{file_path.stem}.MDS.wav'
+        yield f"Audio processing completed successfully.", f'{OUTPUT_FOLDER}{formatted_title}/{formatted_title}.MDS.wav'
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}\n{traceback.format_exc()}"
         logging.error(error_msg)
         yield error_msg, None
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Music Drum Separation")
+    gr.Markdown("# Music Player and Processor")
+
+    file_upload = gr.File(label="Upload WAV file", file_types=[".wav"])
+
+    audio_output = gr.Audio(label="Audio Player")
 
-    with gr.Row():
-        file_input = gr.File(label="Upload WAV File", file_types=['wav'])
-        process_button = gr.Button("Process Audio")
+    process_button = gr.Button("Process Audio")
 
     log_output = gr.Textbox(label="Processing Log", interactive=False)
     processed_audio_output = gr.Audio(label="Processed Audio")
 
+    file_upload.change(
+        fn=lambda file: file.name if file else None,
+        inputs=file_upload,
+        outputs=audio_output
+    )
+
     process_button.click(
         fn=process_audio,
-        inputs=file_input,
+        inputs=file_upload,
         outputs=[log_output, processed_audio_output],
         show_progress=True
     )
-
+        
 demo.launch()