app.py

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

# --- Custom Theme Configuration ---
class MidnightTheme(Soft):
    def __init__(self):
        super().__init__(
            # Using your specific text and button colors for the palettes
            primary_hue=colors.Color(
                name="brand",
                c50="#eef2ff", c100="#e0e7ff", c200="#c7d2fe", c300="#a5b4fc",
                c400="#818cf8", c500="#5248e9", c600="#4f46e5", c700="#4338ca",
                c800="#3730a3", c900="#312e81", c950="#1e1b4b"
            ),
            neutral_hue=colors.Color(
                name="dark_slate",
                c50="#f8fafc", c100="#f1f5f9", c200="#e2e8f0", c300="#cbd5e1",
                c400="#94a3b8", c500="#64748b", c600="#51748c", c700="#334155", # c600 is your secondary text
                c800="#20293c", c900="#10172b", c950="#030617" # c800-950 are your BG/Button darks
            ),
            font=(fonts.GoogleFont("Outfit"), "Arial", "sans-serif"),
        )
        super().set(
            # Backgrounds
            body_background_fill="#030617",
            block_background_fill="#10172b",
            block_border_color="#20293c",
            
            # Text Colors
            body_text_color="#cdd6e2",
            block_label_text_color="#51748c",
            block_title_text_color="#cdd6e2",
            
            # Buttons
            button_primary_background_fill="#5248e9",
            button_primary_text_color="white",
            button_secondary_background_fill="#20293c",
            button_secondary_text_color="#cdd6e2",
            
            # Inputs
            input_background_fill="#030617",
            input_border_color="#20293c",
        )

midnight_theme = MidnightTheme()

# --- CSS for Layout Polish ---
css = """
#container { max-width: 1000px; margin: auto; padding-top: 2rem; }
#title-area { text-align: center; margin-bottom: 2rem; }
.gradio-container { background-color: #030617 !important; }
.output-audio { background-color: #030617 !important; }
"""

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}...")

try:
    model = SAMAudio.from_pretrained(MODEL_ID,token=os.environ.get("HF_TOKEN")).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)}"

def dummy_process(file, text, duration): # Placeholder for structure
    return None, None, "Processing..."

with gr.Blocks(theme=midnight_theme, css=css) as demo:
    with gr.Column(elem_id="container"):
        # Header Section
        gr.Markdown(
            """
            # 🎙️ SAM-Audio Segmenter
            ### Isolate specific sounds using natural language descriptions.
            """, 
            elem_id="title-area"
        )

        with gr.Row(equal_height=True):
            # Left Side: Inputs
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### 1. Upload & Describe")
                    input_file = gr.Audio(label="Input Audio Source", type="filepath")
                    text_prompt = gr.Textbox(
                        label="Target Sound", 
                        placeholder="e.g. 'electric guitar solo' or 'birds chirping'",
                        info="What sound should we isolate from the background?"
                    )
                
                with gr.Accordion("Advanced Processing Settings", open=False):
                    chunk_duration_slider = gr.Slider(
                        minimum=10, maximum=60, value=30, step=5,
                        label="Chunk Duration (s)",
                        info="Shorter chunks save memory for long files."
                    )
                
                run_btn = gr.Button("🚀 Start Separation", variant="primary")

            # Right Side: Outputs
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### 2. Results")
                    output_target = gr.Audio(label="Isolated Result", type="filepath")
                    output_residual = gr.Audio(label="Background / Remainder", type="filepath")
                    status_out = gr.Textbox(label="Status Log", interactive=False, lines=2)

        # Examples Section at Bottom
        gr.Markdown("---")
        gr.Examples(
            examples=[
                ["example_audio/speech.mp3", "Music", 30],
                ["example_audio/song.mp3", "Drum", 30]
            ],
            inputs=[input_file, text_prompt, chunk_duration_slider],
            label="Try an Example"
        )

    # Event Binding
    run_btn.click(
        fn=process_audio, # Use your real function here
        inputs=[input_file, text_prompt, chunk_duration_slider],
        outputs=[output_target, output_residual, status_out]
    )

if __name__ == "__main__":
    demo.launch()