feat: To implement the audio chunking function with overlapping and cross-fading for processing long audio files, and to add chunk duration control to the UI.

app.py

@@ -23,6 +23,11 @@ DEFAULT_MODEL = "sam-audio-small"
 EXAMPLES_DIR = "examples"
 EXAMPLE_FILE = os.path.join(EXAMPLES_DIR, "office.mp4")
 
+# Chunk processing settings
+DEFAULT_CHUNK_DURATION = 30  # seconds per chunk
+OVERLAP_DURATION = 2  # seconds of overlap between chunks
+MAX_DURATION_WITHOUT_CHUNKING = 60  # auto-chunk if longer than this
+
 # Global model cache
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 current_model_name = None
@@ -42,16 +47,77 @@ def load_model(model_name):
 
 load_model(DEFAULT_MODEL)
 
+def load_audio(file_path):
+    """Load audio from file (supports both audio and video files)."""
+    waveform, sample_rate = torchaudio.load(file_path)
+    # Convert to mono if stereo
+    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]
+    
+    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:
+        return chunks[0]
+    
+    overlap_samples = int(overlap_duration * sample_rate)
+    result = chunks[0]
+    
+    for i in range(1, len(chunks)):
+        prev_chunk = result
+        next_chunk = chunks[i]
+        
+        # Create fade curves
+        fade_out = torch.linspace(1.0, 0.0, overlap_samples).to(prev_chunk.device)
+        fade_in = torch.linspace(0.0, 1.0, overlap_samples).to(next_chunk.device)
+        
+        # Get overlapping regions
+        prev_overlap = prev_chunk[:, -overlap_samples:]
+        next_overlap = next_chunk[:, :overlap_samples]
+        
+        # Crossfade mix
+        crossfaded = prev_overlap * fade_out + next_overlap * fade_in
+        
+        # Concatenate: non-overlap of prev + crossfaded + non-overlap of next
+        result = torch.cat([
+            prev_chunk[:, :-overlap_samples],
+            crossfaded,
+            next_chunk[:, overlap_samples:]
+        ], dim=1)
+    
+    return result
+
 def save_audio(tensor, sample_rate):
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
         torchaudio.save(tmp.name, tensor, sample_rate)
         return tmp.name
 
 @spaces.GPU(duration=300)
-def separate_audio(model_name, file_path, text_prompt, progress=gr.Progress()):
+def separate_audio(model_name, file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_DURATION, progress=gr.Progress()):
     global model, processor
     
-    progress(0.1, desc="Checking inputs...")
+    progress(0.05, desc="Checking inputs...")
     
     if not file_path:
         return None, None, "❌ Please upload an audio or video file."
@@ -59,23 +125,70 @@ def separate_audio(model_name, file_path, text_prompt, progress=gr.Progress()):
         return None, None, "❌ Please enter a text prompt."
     
     try:
-        progress(0.2, desc="Loading model...")
+        progress(0.1, desc="Loading model...")
         load_model(model_name)
         
-        progress(0.4, 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.15, desc="Loading audio...")
+        waveform, sample_rate = load_audio(file_path)
+        duration = waveform.shape[1] / sample_rate
         
-        progress(0.8, desc="Saving results...")
-        sample_rate = processor.audio_sampling_rate
-        target_path = save_audio(result.target[0].unsqueeze(0).cpu(), sample_rate)
-        residual_path = save_audio(result.residual[0].unsqueeze(0).cpu(), sample_rate)
+        # Decide whether to use chunking
+        use_chunking = duration > MAX_DURATION_WITHOUT_CHUNKING
         
-        progress(1.0, desc="Done!")
-        return target_path, residual_path, f"✅ Isolated '{text_prompt}' using {model_name}"
+        if use_chunking:
+            progress(0.2, desc=f"Audio is {duration:.1f}s, splitting into chunks...")
+            chunks = split_audio_into_chunks(waveform, sample_rate, chunk_duration, 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}...")
+                
+                # Save chunk to temp file for processor
+                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].cpu())
+                    residual_chunks.append(result.residual[0].cpu())
+                finally:
+                    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.unsqueeze(0), sample_rate)
+            residual_path = save_audio(residual_merged.unsqueeze(0), sample_rate)
+            
+            progress(1.0, desc="Done!")
+            return target_path, residual_path, f"✅ Isolated '{text_prompt}' using {model_name} ({num_chunks} chunks)"
+        else:
+            # Process without chunking
+            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...")
+            sample_rate = processor.audio_sampling_rate
+            target_path = save_audio(result.target[0].unsqueeze(0).cpu(), sample_rate)
+            residual_path = save_audio(result.residual[0].unsqueeze(0).cpu(), sample_rate)
+            
+            progress(1.0, desc="Done!")
+            return target_path, residual_path, f"✅ Isolated '{text_prompt}' using {model_name}"
     except Exception as e:
         import traceback
         traceback.print_exc()
@@ -98,6 +211,16 @@ with gr.Blocks(title="SAM-Audio Demo") as demo:
                 label="Model"
             )
             
+            with gr.Accordion("⚙️ Advanced Options", open=False):
+                chunk_duration_slider = gr.Slider(
+                    minimum=10,
+                    maximum=60,
+                    value=DEFAULT_CHUNK_DURATION,
+                    step=5,
+                    label="Chunk Duration (seconds)",
+                    info=f"Audio longer than {MAX_DURATION_WITHOUT_CHUNKING}s will be automatically split"
+                )
+            
             gr.Markdown("#### Upload Audio")
             input_audio = gr.Audio(label="Audio File", type="filepath")
             
@@ -128,13 +251,13 @@ with gr.Blocks(title="SAM-Audio Demo") as demo:
             example_btn3 = gr.Button("🎵 Background Music")
     
     # Main process button
-    def process(model_name, audio_path, video_path, prompt, progress=gr.Progress()):
+    def process(model_name, audio_path, video_path, prompt, chunk_duration, progress=gr.Progress()):
         file_path = video_path if video_path else audio_path
-        return separate_audio(model_name, file_path, prompt, progress)
+        return separate_audio(model_name, file_path, prompt, chunk_duration, progress)
     
     run_btn.click(
         fn=process,
-        inputs=[model_selector, input_audio, input_video, text_prompt],
+        inputs=[model_selector, input_audio, input_video, text_prompt, chunk_duration_slider],
         outputs=[output_target, output_residual, status_output]
     )