| import gradio as gr |
| import torch |
| import torchaudio |
| import numpy as np |
| import os |
| import tempfile |
| import spaces |
|
|
| from typing import Iterable |
| from gradio.themes import Soft |
| from gradio.themes.utils import colors, fonts, sizes |
|
|
| colors.orange_red = colors.Color( |
| name="orange_red", |
| c50="#FFF0E5", |
| c100="#FFE0CC", |
| c200="#FFC299", |
| c300="#FFA366", |
| c400="#FF8533", |
| c500="#FF4500", |
| c600="#E63E00", |
| c700="#CC3700", |
| c800="#B33000", |
| c900="#992900", |
| c950="#802200", |
| ) |
|
|
| class OrangeRedTheme(Soft): |
| def __init__( |
| self, |
| *, |
| primary_hue: colors.Color | str = colors.gray, |
| secondary_hue: colors.Color | str = colors.orange_red, |
| neutral_hue: colors.Color | str = colors.slate, |
| text_size: sizes.Size | str = sizes.text_lg, |
| font: fonts.Font | str | Iterable[fonts.Font | str] = ( |
| fonts.GoogleFont("Outfit"), "Arial", "sans-serif", |
| ), |
| font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( |
| fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", |
| ), |
| ): |
| super().__init__( |
| primary_hue=primary_hue, |
| secondary_hue=secondary_hue, |
| neutral_hue=neutral_hue, |
| text_size=text_size, |
| font=font, |
| font_mono=font_mono, |
| ) |
| super().set( |
| background_fill_primary="*primary_50", |
| background_fill_primary_dark="*primary_900", |
| body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", |
| body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", |
| button_primary_text_color="white", |
| button_primary_text_color_hover="white", |
| button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", |
| button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", |
| button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", |
| button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", |
| button_secondary_text_color="black", |
| button_secondary_text_color_hover="white", |
| button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)", |
| button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)", |
| button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)", |
| button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)", |
| slider_color="*secondary_500", |
| slider_color_dark="*secondary_600", |
| block_title_text_weight="600", |
| block_border_width="3px", |
| block_shadow="*shadow_drop_lg", |
| button_primary_shadow="*shadow_drop_lg", |
| button_large_padding="11px", |
| color_accent_soft="*primary_100", |
| block_label_background_fill="*primary_200", |
| ) |
|
|
| orange_red_theme = OrangeRedTheme() |
|
|
| try: |
| from sam_audio import SAMAudio, SAMAudioProcessor |
| except ImportError as e: |
| print(f"Warning: 'sam_audio' library not found. Please install it to use this app. Error: {e}") |
|
|
| MODEL_ID = "facebook/sam-audio-large" |
| DEFAULT_CHUNK_DURATION = 30.0 |
| OVERLAP_DURATION = 2.0 |
| MAX_DURATION_WITHOUT_CHUNKING = 30.0 |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"Loading {MODEL_ID} on {device}...") |
|
|
| model = None |
| processor = None |
|
|
| try: |
| model = SAMAudio.from_pretrained(MODEL_ID).to(device).eval() |
| processor = SAMAudioProcessor.from_pretrained(MODEL_ID) |
| print("✅ SAM-Audio loaded successfully.") |
| except Exception as e: |
| print(f"❌ Error loading SAM-Audio: {e}") |
|
|
| def load_audio(file_path): |
| """Load audio from file (supports both audio and video files).""" |
| waveform, sample_rate = torchaudio.load(file_path) |
| if waveform.shape[0] > 1: |
| waveform = waveform.mean(dim=0, keepdim=True) |
| return waveform, sample_rate |
|
|
| def split_audio_into_chunks(waveform, sample_rate, chunk_duration, overlap_duration): |
| """Split audio waveform into overlapping chunks.""" |
| chunk_samples = int(chunk_duration * sample_rate) |
| overlap_samples = int(overlap_duration * sample_rate) |
| stride = chunk_samples - overlap_samples |
|
|
| chunks = [] |
| total_samples = waveform.shape[1] |
|
|
| if total_samples <= chunk_samples: |
| return [waveform] |
|
|
| start = 0 |
| while start < total_samples: |
| end = min(start + chunk_samples, total_samples) |
| chunk = waveform[:, start:end] |
| chunks.append(chunk) |
| if end >= total_samples: |
| break |
| start += stride |
|
|
| return chunks |
|
|
| def merge_chunks_with_crossfade(chunks, sample_rate, overlap_duration): |
| """Merge audio chunks with crossfade on overlapping regions.""" |
| if len(chunks) == 1: |
| chunk = chunks[0] |
| if chunk.dim() == 1: |
| chunk = chunk.unsqueeze(0) |
| return chunk |
|
|
| overlap_samples = int(overlap_duration * sample_rate) |
|
|
| processed_chunks = [] |
| for chunk in chunks: |
| if chunk.dim() == 1: |
| chunk = chunk.unsqueeze(0) |
| processed_chunks.append(chunk) |
|
|
| result = processed_chunks[0] |
|
|
| for i in range(1, len(processed_chunks)): |
| prev_chunk = result |
| next_chunk = processed_chunks[i] |
|
|
| actual_overlap = min(overlap_samples, prev_chunk.shape[1], next_chunk.shape[1]) |
|
|
| if actual_overlap <= 0: |
| result = torch.cat([prev_chunk, next_chunk], dim=1) |
| continue |
|
|
| fade_out = torch.linspace(1.0, 0.0, actual_overlap).to(prev_chunk.device) |
| fade_in = torch.linspace(0.0, 1.0, actual_overlap).to(next_chunk.device) |
|
|
| prev_overlap = prev_chunk[:, -actual_overlap:] |
| next_overlap = next_chunk[:, :actual_overlap] |
|
|
| crossfaded = prev_overlap * fade_out + next_overlap * fade_in |
|
|
| result = torch.cat([ |
| prev_chunk[:, :-actual_overlap], |
| crossfaded, |
| next_chunk[:, actual_overlap:] |
| ], dim=1) |
|
|
| return result |
|
|
| def save_audio(tensor, sample_rate): |
| """Saves a tensor to a temporary WAV file and returns path.""" |
| with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: |
| tensor = tensor.cpu() |
| if tensor.dim() == 1: |
| tensor = tensor.unsqueeze(0) |
| torchaudio.save(tmp.name, tensor, sample_rate) |
| return tmp.name |
|
|
| @spaces.GPU(duration=120) |
| def process_audio(file_path, text_prompt, chunk_duration_val, progress=gr.Progress()): |
| global model, processor |
|
|
| if model is None or processor is None: |
| return None, None, "❌ Model not loaded correctly. Check logs." |
|
|
| progress(0.05, desc="Checking inputs...") |
|
|
| if not file_path: |
| return None, None, "❌ Please upload an audio or video file." |
| if not text_prompt or not text_prompt.strip(): |
| return None, None, "❌ Please enter a text prompt." |
|
|
| try: |
| progress(0.15, desc="Loading audio...") |
| waveform, sample_rate = load_audio(file_path) |
| duration = waveform.shape[1] / sample_rate |
|
|
| c_dur = chunk_duration_val if chunk_duration_val else DEFAULT_CHUNK_DURATION |
| use_chunking = duration > MAX_DURATION_WITHOUT_CHUNKING |
|
|
| if use_chunking: |
| progress(0.2, desc=f"Audio is {duration:.1f}s, splitting into chunks...") |
| chunks = split_audio_into_chunks(waveform, sample_rate, c_dur, OVERLAP_DURATION) |
| num_chunks = len(chunks) |
|
|
| target_chunks = [] |
| residual_chunks = [] |
|
|
| for i, chunk in enumerate(chunks): |
| chunk_progress = 0.2 + (i / num_chunks) * 0.6 |
| progress(chunk_progress, desc=f"Processing chunk {i+1}/{num_chunks}...") |
|
|
| with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: |
| torchaudio.save(tmp.name, chunk, sample_rate) |
| chunk_path = tmp.name |
|
|
| try: |
| inputs = processor(audios=[chunk_path], descriptions=[text_prompt.strip()]).to(device) |
|
|
| with torch.inference_mode(): |
| result = model.separate(inputs, predict_spans=False, reranking_candidates=1) |
|
|
| target_chunks.append(result.target[0].detach().cpu()) |
| residual_chunks.append(result.residual[0].detach().cpu()) |
| finally: |
| if os.path.exists(chunk_path): |
| os.unlink(chunk_path) |
|
|
| progress(0.85, desc="Merging chunks...") |
| target_merged = merge_chunks_with_crossfade(target_chunks, sample_rate, OVERLAP_DURATION) |
| residual_merged = merge_chunks_with_crossfade(residual_chunks, sample_rate, OVERLAP_DURATION) |
|
|
| progress(0.95, desc="Saving results...") |
| target_path = save_audio(target_merged, sample_rate) |
| residual_path = save_audio(residual_merged, sample_rate) |
|
|
| progress(1.0, desc="Done!") |
| return target_path, residual_path, f"✅ Isolated '{text_prompt}' ({num_chunks} chunks)" |
|
|
| else: |
| progress(0.3, desc="Processing audio...") |
| inputs = processor(audios=[file_path], descriptions=[text_prompt.strip()]).to(device) |
|
|
| progress(0.6, desc="Separating sounds...") |
| with torch.inference_mode(): |
| result = model.separate(inputs, predict_spans=False, reranking_candidates=1) |
|
|
| progress(0.9, desc="Saving results...") |
| sr = processor.audio_sampling_rate |
| target_path = save_audio(result.target[0].unsqueeze(0).cpu(), sr) |
| residual_path = save_audio(result.residual[0].unsqueeze(0).cpu(), sr) |
|
|
| progress(1.0, desc="Done!") |
| return target_path, residual_path, f"✅ Isolated '{text_prompt}'" |
|
|
| except Exception as e: |
| import traceback |
| traceback.print_exc() |
| return None, None, f"❌ Error: {str(e)}" |
|
|
| css = """ |
| #main-title h1 {font-size: 2.4em} |
| """ |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("# **SAM-Audio-Demo**", elem_id="main-title") |
| gr.Markdown("Segment and isolate specific music/sounds from audio files using natural language descriptions, powered by [SAM-Audio-Large](https://huggingface.co/facebook/sam-audio-large).") |
|
|
| with gr.Column(elem_id="col-container"): |
| with gr.Row(): |
| with gr.Column(scale=1): |
| input_file = gr.Audio(label="Input Audio", type="filepath") |
| text_prompt = gr.Textbox(label="Sound to Isolate", placeholder="e.g., 'A man speaking', 'Bird chirping'") |
|
|
| with gr.Accordion("Advanced Settings", open=False): |
| chunk_duration_slider = gr.Slider( |
| minimum=10, maximum=60, value=30, step=5, |
| label="Chunk Duration (seconds)", |
| info="Processing long audio in chunks prevents out-of-memory errors." |
| ) |
|
|
| run_btn = gr.Button("Segment Audio", variant="primary") |
|
|
| with gr.Column(scale=1): |
| output_target = gr.Audio(label="Isolated Sound (Target)", type="filepath") |
| output_residual = gr.Audio(label="Background (Residual)", type="filepath") |
| status_out = gr.Textbox(label="Status", interactive=False, show_label=True, lines=6) |
|
|
| gr.Examples( |
| examples=[ |
| ["example_audio/speech.mp3", "Music", 30], |
| ["example_audio/song.mp3", "Drum", 30], |
| ["example_audio/song2.mp3", "Music", 30], |
| ], |
| inputs=[input_file, text_prompt, chunk_duration_slider], |
| label="Audio Examples" |
| ) |
| |
| run_btn.click( |
| fn=process_audio, |
| inputs=[input_file, text_prompt, chunk_duration_slider], |
| outputs=[output_target, output_residual, status_out] |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False) |
