Update app.py

app.py

@@ -4,20 +4,31 @@ import subprocess
 import spaces
 from typing import Tuple, List, Dict
 from pydub import AudioSegment
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn
+from rich.text import Text
+import time
+
+console = Console()
 
 @spaces.GPU
 def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, gr.HTML]:
     log_messages = []
 
-    def stream_log(message):
-        log_messages.append(f"[{model_name}] {message}")
-        return gr.HTML("<pre style='margin-bottom: 0;'>" + "<br>".join(log_messages) + "</pre>")
+    def stream_log(message, style=""):
+        formatted_message = f"[{model_name}] {message}"
+        log_messages.append(formatted_message)
+        return gr.HTML(f"<pre style='margin-bottom: 0;{style}'>{formatted_message}</pre>")
+
+    yield None, stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;")
+    time.sleep(1)  # Simulate initialization time
 
-    yield None, stream_log("Starting separation process...")
-    yield None, stream_log(f"Loading audio file: {audio_file}")
+    yield None, stream_log("Loading audio file...", "color: #2196F3;")
+    time.sleep(0.5)  # Simulate loading time
     
     if audio_file is None:
-        yield None, stream_log("Error: No audio file provided")
+        yield None, stream_log("Error: No audio file provided", "color: #F44336;")
         raise gr.Error("Please upload an audio file")
 
     # Use absolute paths
@@ -33,30 +44,42 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
         audio_file
     ]
 
-    yield None, stream_log(f"Running Demucs command: {' '.join(cmd)}")
+    yield None, stream_log("Preparing separation process...", "color: #FF9800;")
+    time.sleep(0.5)  # Simulate preparation time
 
     try:
         # Run the Demucs command
         process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
         
-        # Stream the output
-        for line in process.stdout:
-            yield None, stream_log(line.strip())
-        
-        # Wait for the process to complete
-        process.wait()
+        # Simulate a loading animation
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TimeRemainingColumn(),
+            console=console
+        ) as progress:
+            task = progress.add_task("[cyan]Separating stems...", total=100)
+            while process.poll() is None:
+                progress.update(task, advance=1)
+                time.sleep(0.1)
+            progress.update(task, completed=100)
 
         if process.returncode != 0:
             error_output = process.stderr.read()
-            yield None, stream_log(f"Error: Demucs command failed with return code {process.returncode}")
-            yield None, stream_log(f"Error output: {error_output}")
+            yield None, stream_log(f"Error: Separation failed", "color: #F44336;")
             raise gr.Error(f"Demucs separation failed. Check the logs for details.")
 
     except Exception as e:
-        yield None, stream_log(f"Unexpected error: {str(e)}")
+        yield None, stream_log(f"Unexpected error: {str(e)}", "color: #F44336;")
         raise gr.Error(f"An unexpected error occurred: {str(e)}")
 
-    yield None, stream_log("Separation completed. Processing stems...")
+    yield None, stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;")
+    time.sleep(0.5)  # Pause for effect
+
+    yield None, stream_log("Processing stems...", "color: #9C27B0;")
+    time.sleep(0.5)  # Simulate processing time
 
     # Change the stem search directory using full path
     stem_search_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
@@ -99,14 +122,16 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     yield None, stream_log(f"Final selected stems: {selected_stems}")
 
     if not selected_stems:
-        yield None, stream_log("Error: No stems selected for mixing")
+        yield None, stream_log("Error: No stems selected for mixing", "color: #F44336;")
         raise gr.Error("Please select at least one stem to mix and ensure it was successfully separated.")
 
     output_file: str = os.path.join(output_dir, "mixed.wav")
-    yield None, stream_log("Mixing selected stems...")
+    yield None, stream_log("Mixing selected stems...", "color: #FF5722;")
+    time.sleep(0.5)  # Simulate mixing time
     
     if len(selected_stems) == 1:
-        os.rename(selected_stems[0], output_file)
+        mixed_audio = AudioSegment.from_wav(selected_stems[0])
+        mixed_audio.export(output_file, format="wav")
     else:
         mixed_audio: AudioSegment = AudioSegment.empty()
         for stem_path in selected_stems:
@@ -114,13 +139,20 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
         mixed_audio.export(output_file, format="wav")
 
     if mp3:
-        yield None, stream_log(f"Converting to MP3 (bitrate: {mp3_bitrate}k)...")
+        yield None, stream_log(f"Converting to MP3...", "color: #795548;")
+        time.sleep(0.5)  # Simulate conversion time
         mp3_output_file: str = os.path.splitext(output_file)[0] + ".mp3"
         mixed_audio.export(mp3_output_file, format="mp3", bitrate=str(mp3_bitrate) + "k")
         output_file = mp3_output_file
 
-    yield None, stream_log("Process completed successfully!")
-    yield output_file, gr.HTML("<pre style='color: green;'>Separation and mixing completed successfully!</pre>")
+    yield None, stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;")
+    yield output_file, gr.HTML(
+        Panel.fit(
+            Text("Separation and mixing completed successfully!", style="bold green"),
+            title="Demucs Result",
+            border_style="green"
+        ).render()
+    )
 
 # Define the Gradio interface
 with gr.Blocks() as iface: