From f2011b44086337e0669e365b93aacd21ccb61058 Mon Sep 17 00:00:00 2001
From: Joel Mathew Thomas <90510078+joelmathewthomas@users.noreply.github.com>
Date: Wed, 26 Feb 2025 14:23:08 +0530
Subject: [PATCH] endpoint : /api/separate

- Define new endpoint /api/separate, to separate music using demucs, params: file_uuid
- Replace original file with vocals.wav while retaining original filename
- Move all other files to file_path/sources/
---
 api/api/tasks.py    | 58 ++++++++++++++++++++++++++++++++++++++++++++-
 api/api/views.py    | 17 +++++++++++++
 api/backend/urls.py |  4 +++-
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/api/api/tasks.py b/api/api/tasks.py
index 40abe8b..ba7ac41 100644
--- a/api/api/tasks.py
+++ b/api/api/tasks.py
@@ -1,3 +1,6 @@
+import os
+import shutil
+from pathlib import Path
 from celery import shared_task
 from freqsplit.input.file_reader import read_audio
 from freqsplit.preprocessing.classify import classify_audio
@@ -5,6 +8,7 @@ from freqsplit.preprocessing.normalize import normalize_audio
 from freqsplit.preprocessing.trim import trim_audio 
 from freqsplit.preprocessing.resample import resample
 from freqsplit.postprocessing.audio_writer import export_audio
+from freqsplit.separation.demucs_wrapper import separate_audio_with_demucs
 
 @shared_task
 def save_and_classify(file_path, file_content):
@@ -55,4 +59,56 @@ def resample_audio_task(file_path, sr):
          return True
     except Exception as e:
         raise RuntimeError(f"RuntimeError: {e}")
-        return False
\ No newline at end of file
+        return False
+
+@shared_task
+def music_separation_task(file_path):
+    """Celery task to separate music audio into sources"""
+    file_path = Path(file_path)
+    print("File path is ", file_path)
+    
+    # Determine the base directory (output path)
+    output_path = file_path.parent
+    
+    # Run Demucs separation
+    separate_audio_with_demucs(str(file_path), str(output_path))
+    
+    # Define expected output dir
+    demucs_dir = output_path / 'htdemucs'
+    file_folder = demucs_dir / file_path.stem 
+    
+    if not file_folder.exists():
+        raise RuntimeError(f"Demucs output folder not found: {file_folder}")
+
+    # Expected output files
+    expected_files = ["bass.wav", "drums.wav", "other.wav", "vocals.wav"]
+
+    # Create "sources" directory to store separated components
+    sources_dir = output_path / "sources"
+    sources_dir.mkdir(exist_ok=True)
+    
+    # Move separate files to output_path and replace original file with vocals.wav, and move other files into sources/
+    try:
+        vocals_path = file_folder / "vocals.wav"
+        if not vocals_path.exists():
+            raise RuntimeError("Vocals file not found in Demucs output")
+        
+        # Replace original file with vocals.wav while keeping original name
+        shutil.move(str(vocals_path), str(file_path))
+        
+        # Move other separated files to the "sources" directory
+        for expected_file in expected_files:
+            src_file = file_folder / expected_file
+            if src_file.exists() and expected_file != "vocals.wav":
+                shutil.move(str(src_file), str(sources_dir / expected_file))
+            
+        # Cleanup: Remove htedemucs directory
+        shutil.rmtree(str(demucs_dir))
+        
+        return True
+    
+    except Exception as e:
+        raise RuntimeError(f"Music source separation task failed: {e}")
+    
+    return False 
+   
\ No newline at end of file
diff --git a/api/api/views.py b/api/api/views.py
index 88ef97a..29e18e9 100644
--- a/api/api/views.py
+++ b/api/api/views.py
@@ -8,6 +8,7 @@ from .tasks import save_and_classify
 from .tasks import normalize_audio_task
 from .tasks import trim_audio_task
 from .tasks import resample_audio_task
+from .tasks import music_separation_task
 from freqsplit.input.format_checker import is_supported_format
 
 UPLOAD_DIR = "/tmp/freqsplit"
@@ -105,3 +106,19 @@ def resample_audio(request):
         return Response({"message": f"Audio resampled to {sr} successfully"}, status=status.HTTP_200_OK)
     else:
         return Response({"error": "Failed to resample audio"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+# Endpoint to separate music into sources
+@api_view(['POST'])
+def separate_music(request):
+    """Handles the separatioo of music audio into source components"""
+    stat, result, status_code = get_audio_file_path(request, UPLOAD_DIR)
+    if stat == False:
+        return Response({"error": result}, status=status_code)  
+
+    # Call Celery task synchronously
+    task = music_separation_task.apply(args=(result,))
+
+    if task.get():
+        return Response({"message": f"Audio separated into sources successfully"}, status=status.HTTP_200_OK)
+    else:
+        return Response({"error": "Failed to source separate audio"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
diff --git a/api/backend/urls.py b/api/backend/urls.py
index 961df0e..e8cd6c5 100644
--- a/api/backend/urls.py
+++ b/api/backend/urls.py
@@ -20,11 +20,13 @@ from api.views import upload_audio
 from api.views import normalize_audio
 from api.views import trim_audio
 from api.views import resample_audio
+from api.views import separate_music
 
 urlpatterns = [
     path('admin/', admin.site.urls),
     path('api/upload', upload_audio, name='upload_audio'),
     path('api/normalize', normalize_audio, name="normalize_audio"),
     path('api/trim', trim_audio, name='trim_audio'),
-    path('api/resample', resample_audio, name='resample_audio')
+    path('api/resample', resample_audio, name='resample_audio'),
+    path('api/separate', separate_music, name="separate_music")
 ]