Update app.py

app.py

@@ -1,7 +1,73 @@
 import gradio as gr
+import torch
+import torchaudio
+import os
+import tempfile
+from sam_audio import SAMAudio, SAMAudioProcessor
 
-def greet(name):
-    return "Hello " + name + "!!"
+# --- Initialization ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Load model and processor once when the app starts
+model = SAMAudio.from_pretrained("facebook/sam-audio-large").to(device).eval()
+processor = SAMAudioProcessor.from_pretrained("facebook/sam-audio-large")
+
+def separate_audio(audio_path, description, reranking_candidates):
+    if audio_path is None or not description:
+        return None, None
+
+    # Process inputs
+    inputs = processor(audios=[audio_path], descriptions=[description]).to(device)
+    
+    with torch.inference_mode():
+        # Using reranking if candidates > 1
+        result = model.separate(
+            inputs, 
+            predict_spans=True, 
+            reranking_candidates=int(reranking_candidates)
+        )
+
+    # Use temporary files to store the results for Gradio
+    target_path = os.path.join(tempfile.gettempdir(), "target.wav")
+    residual_path = os.path.join(tempfile.gettempdir(), "residual.wav")
+
+    # Save target and residual
+    torchaudio.save(target_path, result.target[0].unsqueeze(0).cpu(), processor.audio_sampling_rate)
+    torchaudio.save(residual_path, result.residual[0].unsqueeze(0).cpu(), processor.audio_sampling_rate)
+
+    return target_path, residual_path
+
+# --- UI Design ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎵 SAM-Audio Separation")
+    gr.Markdown("Upload an audio file and describe the specific sound you want to isolate (e.g., 'A dog barking' or 'A man speaking').")
+    
+    with gr.Row():
+        with gr.Column():
+            input_audio = gr.Audio(label="Input Audio", type="filepath")
+            description = gr.Textbox(
+                label="What do you want to isolate?", 
+                placeholder="e.g. A person laughing"
+            )
+            rerank_slider = gr.Slider(
+                minimum=1, 
+                maximum=16, 
+                value=1, 
+                step=1, 
+                label="Reranking Candidates",
+                info="Higher values improve quality but increase processing time."
+            )
+            btn = gr.Button("Separate Sound", variant="primary")
+            
+        with gr.Column():
+            output_target = gr.Audio(label="Isolated (Target) Audio")
+            output_residual = gr.Audio(label="Residual Audio")
+
+    btn.click(
+        fn=separate_audio,
+        inputs=[input_audio, description, rerank_slider],
+        outputs=[output_target, output_residual]
+    )
+
+if __name__ == "__main__":
+    demo.launch()