Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import os
 import subprocess
 import spaces
+import torch
 from typing import Tuple, List, Dict
 from pydub import AudioSegment
 from rich.console import Console
@@ -9,8 +10,13 @@ from rich.panel import Panel
 from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn
 from rich.text import Text
 import time
+import shutil
 
-console = Console()
+# Get the terminal width, or use a default if not available
+terminal_width = shutil.get_terminal_size((80, 20)).columns
+
+# Create a console with a specific width
+console = Console(width=min(terminal_width, 100))  # Limit to 100 columns max
 
 def fade_text(text, duration=0.5):
     for i in range(10):
@@ -21,8 +27,7 @@ def fade_text(text, duration=0.5):
 @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, List[str], gr.HTML]:
     log_messages = []
-    console = Console(record=True, width=80)
-
+    
     def stream_log(message, style=""):
         formatted_message = f"[{model_name}] {message}"
         log_messages.append(formatted_message)
@@ -44,11 +49,17 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     output_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
     os.makedirs(output_dir, exist_ok=True)
 
-    # Construct the Demucs command with full paths
+    # Check if CUDA is available
+    cuda_available = torch.cuda.is_available()
+    device = "cuda" if cuda_available else "cpu"
+    yield from stream_log(f"Using device: {device}", "color: #4CAF50; font-weight: bold;")
+
+    # Construct the Demucs command with full paths and GPU flag
     cmd = [
         "python", "-m", "demucs", 
         "--out", base_output_dir,
         "-n", model_name,
+        "--device", device,
         audio_file
     ]
 
@@ -56,14 +67,20 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     time.sleep(0.5)  # Simulate preparation time
 
     try:
+        # Set CUDA_VISIBLE_DEVICES environment variable
+        env = os.environ.copy()
+        if cuda_available:
+            env["CUDA_VISIBLE_DEVICES"] = "0"  # Use the first GPU
+        
         # Run the Demucs command
-        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, env=env)
         
-        # Simulate a loading animation
+        # Simulate a loading animation with adjusted width
+        progress_width = min(terminal_width - 20, 60)  # Adjust the width of the progress bar
         with Progress(
             SpinnerColumn(),
             TextColumn("[progress.description]{task.description}"),
-            BarColumn(),
+            BarColumn(bar_width=progress_width),
             TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
             TimeRemainingColumn(),
             console=console
@@ -137,10 +154,27 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     yield from stream_log("Mixing selected stems...", "color: #FF5722;")
     time.sleep(0.5)  # Simulate mixing time
     
-    mixed_audio: AudioSegment = AudioSegment.empty()
-    for stem_path in selected_stems:
-        mixed_audio += AudioSegment.from_wav(stem_path)
-    mixed_audio.export(output_file, format="wav")
+    if selected_stems:
+        # Load the first stem as the base
+        mixed_audio: AudioSegment = AudioSegment.from_wav(selected_stems[0])
+        
+        # Overlay the remaining stems
+        for stem_path in selected_stems[1:]:
+            overlay_audio = AudioSegment.from_wav(stem_path)
+            
+            # Ensure both segments have the same duration
+            max_length = max(len(mixed_audio), len(overlay_audio))
+            mixed_audio = mixed_audio.pad_to_width(max_length)
+            overlay_audio = overlay_audio.pad_to_width(max_length)
+            
+            # Overlay the audio
+            mixed_audio = mixed_audio.overlay(overlay_audio)
+        
+        # Export the mixed audio
+        mixed_audio.export(output_file, format="wav")
+    else:
+        yield from stream_log("Error: No stems to mix", "color: #F44336;")
+        raise gr.Error("No stems were selected or found for mixing.")
 
     if mp3:
         yield from stream_log(f"Converting to MP3...", "color: #795548;")
@@ -153,7 +187,8 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     final_message = Panel(
         Text("Separation and mixing completed successfully!", style="bold green"),
         title="Demucs Result",
-        border_style="green"
+        border_style="green",
+        width=min(terminal_width - 2, 98)  # Adjust panel width
     )
     console.print(final_message)
     yield output_file, list(stems.values()), gr.HTML(console.export_html())