Update app.py

app.py

@@ -1,15 +1,16 @@
 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
 
+# ==========================================
+# 1. Theme Definition (Orange Red)
+# ==========================================
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -30,7 +31,7 @@ class OrangeRedTheme(Soft):
         self,
         *,
         primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red, # Use the new color
+        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] = (
@@ -54,41 +55,28 @@ class OrangeRedTheme(Soft):
             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()
 
+# ==========================================
+# 2. Model Loading
+# ==========================================
 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}")
+    print(f"Warning: 'sam_audio' library not found. 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
@@ -101,173 +89,67 @@ try:
 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
-
+# ==========================================
+# 3. Processing Function
+# ==========================================
 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()
+        # torchaudio expects [channels, time]
         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()):
+def process_audio(file_path, text_prompt, rerank, 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...")
+        return None, None, "❌ Model not loaded correctly."
 
     if not file_path:
-        return None, None, "❌ Please upload an audio or video file."
+        return None, None, "❌ Please upload an audio 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}'"
+        progress(0.2, desc="Processing audio...")
+        
+        # Prepare inputs
+        inputs = processor(audios=[file_path], descriptions=[text_prompt.strip()]).to(device)
+
+        progress(0.5, desc="Separating sound...")
+        with torch.inference_mode():
+            # Run separation
+            # Using reranking improves quality but adds latency
+            candidates = int(rerank) if rerank else 1
+            result = model.separate(inputs, predict_spans=True, reranking_candidates=candidates)
+
+        progress(0.9, desc="Saving results...")
+        sr = processor.audio_sampling_rate
+        
+        # Save Target
+        target_path = save_audio(result.target[0], sr)
+        
+        # Save Residual (Background)
+        residual_path = save_audio(result.residual[0], sr)
+
+        progress(1.0, desc="Done!")
+        return target_path, residual_path, f"✅ Successfully isolated '{text_prompt}'"
 
     except Exception as e:
         import traceback
         traceback.print_exc()
         return None, None, f"❌ Error: {str(e)}"
 
+# ==========================================
+# 4. Gradio Interface
+# ==========================================
 css = """
 #main-title h1 {font-size: 2.4em}
+#col-container {max-width: 1000px; margin: 0 auto;}
 """
 
 with gr.Blocks() as demo:
@@ -276,37 +158,40 @@ with gr.Blocks() as demo:
 
     with gr.Column(elem_id="col-container"):
         with gr.Row():
+            # Left Column: Inputs
             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."
+                    rerank_slider = gr.Slider(
+                        minimum=1, maximum=8, value=3, step=1, 
+                        label="Reranking Candidates", 
+                        info="Higher values improve quality but take longer."
                     )
-
+                
                 run_btn = gr.Button("Segment Audio", variant="primary")
 
+            # Right Column: Outputs
             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)
+                status_out = gr.Textbox(label="Status", interactive=False, show_label=True, lines=2)
 
+        # Examples
         gr.Examples(
             examples=[
-                ["example_audio/speech.mp3", "Music", 30],
-                ["example_audio/song.mp3", "Drum", 30],
-                ["example_audio/song2.mp3", "Music", 30],
+                ["example_audio/speech.mp3", "Music"],
+                ["example_audio/song.mp3", "Drum"],
+                ["example_audio/song2.mp3", "Vocals"],
             ],
-            inputs=[input_file, text_prompt, chunk_duration_slider],
+            inputs=[input_file, text_prompt],
             label="Audio Examples"
         )
     
     run_btn.click(
         fn=process_audio,
-        inputs=[input_file, text_prompt, chunk_duration_slider],
+        inputs=[input_file, text_prompt, rerank_slider],
         outputs=[output_target, output_residual, status_out]
     )