Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import os
 import tempfile
 import spaces
+import pandas as pd
 from typing import Optional
 
 # ---------------------------------------------------------
@@ -17,7 +18,6 @@ try:
     from sam3.model_builder import build_sam3_video_predictor
 except ImportError as e:
     print(f"Warning: Essential libraries (sam_audio, sam3, torchcodec) not found. {e}")
-    # Define dummy placeholders to prevent immediate crash if libs are missing during check
     SAMAudio = None
     SAMAudioProcessor = None
     build_sam3_video_predictor = None
@@ -54,8 +54,6 @@ except Exception as e:
 print("Loading SAM3 Video Predictor...")
 try:
     video_predictor = build_sam3_video_predictor()
-    # If SAM3 needs specific device movement, it's usually handled internally or via .to(device) 
-    # dependent on implementation, but standard build function usually sets it up.
     print("✅ SAM3 loaded successfully.")
 except Exception as e:
     print(f"❌ Error loading SAM3: {e}")
@@ -157,9 +155,8 @@ def save_audio_temp(tensor, sample_rate):
 # ---------------------------------------------------------
 # Tab 1: Text Prompting (with Chunking)
 # ---------------------------------------------------------
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=300)
 def process_text_prompting(file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_DURATION, progress=gr.Progress()):
-    # Global models are already loaded
     global model, processor
     
     progress(0.05, desc="Checking inputs...")
@@ -192,17 +189,14 @@ def process_text_prompting(file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_
                 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:
@@ -217,10 +211,9 @@ def process_text_prompting(file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_
             target_path = save_audio_temp(target_merged, sample_rate)
             residual_path = save_audio_temp(residual_merged, sample_rate)
             
-            progress(1.0, desc="Done!")
-            return target_path, residual_path, f"✅ Isolated '{text_prompt}' ({num_chunks} chunks processed)"
+            return target_path, residual_path, f"✅ Isolated '{text_prompt}' ({num_chunks} chunks)"
         else:
-            # Process without chunking (faster for short files)
+            # Process without chunking
             progress(0.3, desc="Processing audio...")
             inputs = processor(audios=[file_path], descriptions=[text_prompt.strip()]).to(device)
             
@@ -233,7 +226,6 @@ def process_text_prompting(file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_
             target_path = save_audio_temp(result.target[0].unsqueeze(0).cpu(), sr)
             residual_path = save_audio_temp(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:
@@ -244,9 +236,8 @@ def process_text_prompting(file_path, text_prompt, chunk_duration=DEFAULT_CHUNK_
 # ---------------------------------------------------------
 # Tab 2: Visual Prompting
 # ---------------------------------------------------------
-@spaces.GPU(duration=180)
+@spaces.GPU
 def process_visual_prompting(video_file, visual_prompt_text, progress=gr.Progress()):
-    # Global models are already loaded
     global model, processor, video_predictor
 
     if model is None or video_predictor is None:
@@ -257,14 +248,10 @@ def process_visual_prompting(video_file, visual_prompt_text, progress=gr.Progres
     
     try:
         progress(0.1, desc="Initializing Video Decoder...")
-        # 1. Initialize Video Decoder
         decoder = VideoDecoder(video_file)
-        frames = decoder[:] # Get all frames (careful with memory on long videos)
+        frames = decoder[:]
         
-        # 2. Generate Masks using SAM3
         progress(0.2, desc="Starting SAM3 Session...")
-        
-        # Start SAM3 Session
         response = video_predictor.handle_request({
             "type": "start_session",
             "resource_path": video_file,
@@ -273,9 +260,7 @@ def process_visual_prompting(video_file, visual_prompt_text, progress=gr.Progres
 
         progress(0.3, desc=f"Generating masks for {len(decoder)} frames...")
         masks = []
-        
-        # Frame-by-frame masking
-        step = 1 # Process every frame
+        step = 1 
         total_frames = len(decoder)
         
         for frame_index in range(0, total_frames, step):
@@ -292,16 +277,14 @@ def process_visual_prompting(video_file, visual_prompt_text, progress=gr.Progres
             
             if mask.shape[0] == 0:
                 mask = np.zeros_like(frames[0, [0]], dtype=bool)
-            masks.append(mask[:1]) # Take first mask
+            masks.append(mask[:1]) 
 
-        # Concatenate masks: (Time, 1, H, W)
         final_mask = torch.from_numpy(np.concatenate(masks)).unsqueeze(1)
 
         progress(0.8, desc="Separating Audio with SAM-Audio...")
-        # 3. Process with SAM-Audio using the visual mask
         inputs = processor(
             audios=[video_file],
-            descriptions=[""], # Empty description as we use masks
+            descriptions=[""], 
             masked_videos=processor.mask_videos([frames], [final_mask]),
         ).to(device)
 
@@ -320,35 +303,108 @@ def process_visual_prompting(video_file, visual_prompt_text, progress=gr.Progres
         traceback.print_exc()
         return None, None, f"❌ Error: {str(e)}"
 
+# ---------------------------------------------------------
+# Tab 3: Span Prompting (Temporal Anchors)
+# ---------------------------------------------------------
+@spaces.GPU
+def process_span_prompting(audio_file, description, anchors_df, progress=gr.Progress()):
+    global model, processor
+    
+    if model is None:
+        return None, None, "❌ Model not loaded."
+    if audio_file is None:
+        return None, None, "❌ Please upload an audio file."
+    
+    try:
+        progress(0.1, desc="Parsing anchors...")
+        # Parse Dataframe to List: [["+", 6.3, 7.0], ...]
+        formatted_anchors = []
+        
+        # Validating anchors_df (it usually comes as a list of lists from Gradio Dataframe)
+        if anchors_df is not None:
+            # Handle standard list of lists
+            if isinstance(anchors_df, list):
+                iterable = anchors_df
+            # Handle pandas DataFrame if Gradio returns that (depends on config)
+            elif hasattr(anchors_df, "values"):
+                iterable = anchors_df.values.tolist()
+            else:
+                iterable = []
+
+            for row in iterable:
+                # row structure: [Type, Start, End]
+                # Ensure we have data
+                if len(row) >= 3 and row[0] and row[1] is not None and row[2] is not None:
+                    # Validate Type
+                    anc_type = str(row[0]).strip()
+                    if anc_type not in ["+", "-"]:
+                        continue # Skip invalid rows
+                    
+                    try:
+                        start_t = float(row[1])
+                        end_t = float(row[2])
+                        formatted_anchors.append([anc_type, start_t, end_t])
+                    except ValueError:
+                        continue # Skip if numbers aren't valid
+        
+        if not formatted_anchors:
+            return None, None, "❌ No valid anchors provided. Use '+' or '-' and valid time numbers."
+        
+        print(f"Anchors: {formatted_anchors}")
+
+        progress(0.3, desc="Processing with Anchors...")
+        
+        # Note: We do NOT use chunking here because anchors rely on absolute timestamps 
+        # of the original file.
+        inputs = processor(
+            audios=[audio_file],
+            descriptions=[description if description else "Sound"], # Desc can be generic if anchors are strong
+            anchors=[formatted_anchors],
+        ).to(device)
+
+        with torch.inference_mode():
+            result = model.separate(inputs)
+
+        progress(0.9, desc="Saving results...")
+        sr = processor.audio_sampling_rate
+        target_path = save_audio_temp(result.target[0].unsqueeze(0).cpu(), sr)
+        residual_path = save_audio_temp(result.residual[0].unsqueeze(0).cpu(), sr)
+
+        return target_path, residual_path, f"✅ Processed with {len(formatted_anchors)} anchors."
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, None, f"❌ Error: {str(e)}"
+
 # ---------------------------------------------------------
 # Gradio Interface
 # ---------------------------------------------------------
 css = """
 #main-title h1 {font-size: 2.3em !important; text-align: center;}
-.gradio-container {max-width: 1400px !important; margin: auto;}
+.gradio-container {max-width: 1000px !important; margin: auto;}
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# **SAM-Audio-Demo**", elem_id="main-title")
-    gr.Markdown("Segment and isolate sounds using **Text** or **Visual** prompts.")
+    gr.Markdown("# **SAM-Audio** 🔊", elem_id="main-title")
+    gr.Markdown("Segment and isolate sounds using **Text**, **Visual**, or **Time-based** prompts.")
 
     with gr.Tabs():
         # ================= TAB 1: TEXT =================
         with gr.Tab("Text Prompting"):
             gr.Markdown("### Isolate sound using a text description")
-            gr.Markdown("Supports long audio files via auto-chunking.")
             
             with gr.Row():
                 with gr.Column():
                     t1_input = gr.Audio(label="Input Audio/Video", type="filepath")
                     t1_desc = gr.Textbox(label="Description", placeholder="e.g., 'A man speaking', 'Glass breaking'")
-                    t1_chunk = gr.Slider(minimum=10, maximum=60, value=30, step=5, label="Chunk Duration (seconds)", info="Processing window size. Smaller is faster but loses context.")
+                    t1_chunk = gr.Slider(minimum=10, maximum=60, value=30, step=5, label="Chunk Duration (s)", info="Split long audio into chunks.")
                     t1_btn = gr.Button("Separate Audio", variant="primary")
                 
                 with gr.Column():
                     t1_status = gr.Textbox(label="Status", interactive=False)
-                    t1_target = gr.Audio(label="Target Audio (Isolated)", type="filepath")
-                    t1_residual = gr.Audio(label="Residual Audio (Background)", type="filepath")
+                    t1_target = gr.Audio(label="Target Audio", type="filepath")
+                    t1_residual = gr.Audio(label="Residual Audio", type="filepath")
             
             t1_btn.click(
                 fn=process_text_prompting,
@@ -358,18 +414,17 @@ with gr.Blocks(css=css) as demo:
 
         # ================= TAB 2: VISUAL =================
         with gr.Tab("Visual Prompting"):
-            gr.Markdown("### Isolate sound corresponding to a visual object in a video")
-            gr.Markdown("*Note: Uses SAM3 for object detection. Processing is heavier.*")
+            gr.Markdown("### Isolate sound corresponding to a visual object")
             
             with gr.Row():
                 with gr.Column():
                     t2_input = gr.Video(label="Input Video", format="mp4")
-                    t2_visual_desc = gr.Textbox(label="Visual Object Description", placeholder="e.g., 'The person on the left', 'The red car'")
+                    t2_visual_desc = gr.Textbox(label="Visual Object Description", placeholder="e.g., 'The person on the left'")
                     t2_btn = gr.Button("Generate Mask & Separate", variant="primary")
                 
                 with gr.Column():
                     t2_status = gr.Textbox(label="Status", interactive=False)
-                    t2_target = gr.Audio(label="Target Audio (From Object)", type="filepath")
+                    t2_target = gr.Audio(label="Target Audio", type="filepath")
                     t2_residual = gr.Audio(label="Residual Audio", type="filepath")
 
             t2_btn.click(
@@ -378,6 +433,37 @@ with gr.Blocks(css=css) as demo:
                 outputs=[t2_target, t2_residual, t2_status]
             )
 
+        # ================= TAB 3: SPAN =================
+        with gr.Tab("Span Prompting"):
+            gr.Markdown("### Isolate sound using Temporal Anchors (Time Ranges)")
+            gr.Markdown("Specify when a sound **IS (+)** or **IS NOT (-)** present.")
+            
+            with gr.Row():
+                with gr.Column():
+                    t3_input = gr.Audio(label="Input Audio", type="filepath")
+                    t3_desc = gr.Textbox(label="Description (Optional)", placeholder="e.g. 'A horn honking'")
+                    
+                    t3_anchors = gr.Dataframe(
+                        headers=["Type (+/-)", "Start (s)", "End (s)"],
+                        datatype=["str", "number", "number"],
+                        row_count=3,
+                        col_count=(3, "fixed"),
+                        label="Temporal Anchors",
+                        value=[["+", 0.0, 5.0], ["-", 5.0, 10.0], ["", None, None]]
+                    )
+                    t3_btn = gr.Button("Separate with Anchors", variant="primary")
+                
+                with gr.Column():
+                    t3_status = gr.Textbox(label="Status", interactive=False)
+                    t3_target = gr.Audio(label="Target Audio", type="filepath")
+                    t3_residual = gr.Audio(label="Residual Audio", type="filepath")
+
+            t3_btn.click(
+                fn=process_span_prompting,
+                inputs=[t3_input, t3_desc, t3_anchors],
+                outputs=[t3_target, t3_residual, t3_status]
+            )
+
 if __name__ == "__main__":
     demo.launch(theme=gr.themes.Soft(
             primary_hue="blue",