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 from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn from rich.text import Text import time import shutil # 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): opacity = i / 10 yield f"
{text}
" time.sleep(duration / 10) @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 = [] def stream_log(message, style=""): formatted_message = f"[{model_name}] {message}" log_messages.append(formatted_message) for frame in fade_text(f"
{formatted_message}
"): yield None, None, gr.HTML(frame) yield from stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;") time.sleep(1) # Simulate initialization time yield from stream_log("Loading audio file...", "color: #2196F3;") time.sleep(0.5) # Simulate loading time if audio_file is None: yield from stream_log("Error: No audio file provided", "color: #F44336;") raise gr.Error("Please upload an audio file") # Use absolute paths base_output_dir = os.path.abspath("separated") 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) # 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 ] yield from stream_log("Preparing separation process...", "color: #FF9800;") 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, env=env) # 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(bar_width=progress_width), 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 from 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 from stream_log(f"Unexpected error: {str(e)}", "color: #F44336;") raise gr.Error(f"An unexpected error occurred: {str(e)}") yield from stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;") time.sleep(0.5) # Pause for effect yield from 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]) yield from stream_log(f"Searching for stems in: {stem_search_dir}") stems: Dict[str, str] = {} for stem in ["vocals", "drums", "bass", "other"]: stem_path = os.path.join(stem_search_dir, f"{stem}.wav") yield from stream_log(f"Checking for {stem} stem at: {stem_path}") if os.path.exists(stem_path): stems[stem] = stem_path yield from stream_log(f"Found {stem} stem") else: yield from stream_log(f"Warning: {stem} stem not found") if not stems: yield from stream_log("Error: No stems found. Checking alternative directory...") stem_search_dir = os.path.join(base_output_dir, model_name) for stem in ["vocals", "drums", "bass", "other"]: stem_path = os.path.join(stem_search_dir, f"{stem}.wav") yield from stream_log(f"Checking for {stem} stem at: {stem_path}") if os.path.exists(stem_path): stems[stem] = stem_path yield from stream_log(f"Found {stem} stem") else: yield from stream_log(f"Warning: {stem} stem not found") yield from stream_log(f"All found stems: {list(stems.keys())}") selected_stems: List[str] = [] for stem, selected in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]): if selected: yield from stream_log(f"{stem} is selected by user") if stem in stems: selected_stems.append(stems[stem]) yield from stream_log(f"Selected {stem} stem for mixing") else: yield from stream_log(f"Warning: {stem} was selected but not found") yield from stream_log(f"Final selected stems: {selected_stems}") if not selected_stems: yield from 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 from stream_log("Mixing selected stems...", "color: #FF5722;") time.sleep(0.5) # Simulate mixing time 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)) if len(mixed_audio) < max_length: mixed_audio += AudioSegment.silent(duration=max_length - len(mixed_audio)) if len(overlay_audio) < max_length: overlay_audio += AudioSegment.silent(duration=max_length - len(overlay_audio)) # 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;") 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 from stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;") yield output_file, list(stems.values()), gr.HTML("
Separation and mixing completed successfully!
") # Define the Gradio interface with gr.Blocks() as iface: gr.Markdown("# Demucs Music Source Separation and Mixing") gr.Markdown("Separate vocals, drums, bass, and other instruments from your music using Demucs and mix the selected stems.") with gr.Row(): with gr.Column(scale=1): audio_input = gr.Audio(type="filepath", label="Upload Audio File") model_dropdown = gr.Dropdown( ["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"], label="Model Name", value="htdemucs_ft" ) with gr.Row(): vocals_checkbox = gr.Checkbox(label="Vocals", value=True) drums_checkbox = gr.Checkbox(label="Drums", value=True) with gr.Row(): bass_checkbox = gr.Checkbox(label="Bass", value=True) other_checkbox = gr.Checkbox(label="Other", value=True) mp3_checkbox = gr.Checkbox(label="Save as MP3", value=False) mp3_bitrate = gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False) submit_btn = gr.Button("Process", variant="primary") with gr.Column(scale=1): output_audio = gr.Audio(type="filepath", label="Processed Audio (Mixed)") stems_output = gr.File(label="Individual Stems", file_count="multiple") separation_log = gr.HTML() submit_btn.click( fn=inference, inputs=[audio_input, model_dropdown, vocals_checkbox, drums_checkbox, bass_checkbox, other_checkbox, mp3_checkbox, mp3_bitrate], outputs=[output_audio, stems_output, separation_log] ) mp3_checkbox.change( fn=lambda mp3: gr.update(visible=mp3), inputs=mp3_checkbox, outputs=mp3_bitrate ) # Launch the Gradio interface iface.launch()