Upload 3 files

app.py

@@ -0,0 +1,225 @@
+import gradio as gr
+import torch
+
+try:
+    import spaces
+    ZERO_GPU = True
+except ImportError:
+    ZERO_GPU = False
+import numpy as np
+from transformers import ASTForAudioClassification, AutoFeatureExtractor
+from pydub import AudioSegment
+import tempfile
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+MODEL_NAME = "Vyvo-Research/AST-Music-Classifier-1K"
+DETECTION_THRESHOLD = 0.50
+WINDOW_SIZE = 5.0
+HOP_SIZE = 5.0
+
+logger.info("Loading model on CPU...")
+model = ASTForAudioClassification.from_pretrained(MODEL_NAME)
+feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
+model.eval()
+logger.info("Model loaded")
+
+
+def load_audio(file_path: str, target_sr: int):
+    audio = AudioSegment.from_file(file_path)
+    audio = audio.set_channels(1).set_frame_rate(target_sr)
+    samples = np.array(audio.get_array_of_samples()).astype(np.float32)
+    samples = samples / np.iinfo(audio.array_type).max
+    return samples, audio
+
+
+@torch.no_grad()
+def detect_music(audio_array, sample_rate):
+    window_samples = int(WINDOW_SIZE * sample_rate)
+    hop_samples = int(HOP_SIZE * sample_rate)
+    total_samples = len(audio_array)
+
+    music_segments = []
+    last_was_music = False
+    device = next(model.parameters()).device
+    use_half = device.type == "cuda"
+
+    for start in range(0, total_samples, hop_samples):
+        end = min(start + window_samples, total_samples)
+        segment = audio_array[start:end]
+
+        if len(segment) < sample_rate:
+            continue
+
+        needs_padding = len(segment) < window_samples
+        if needs_padding:
+            segment = np.pad(segment, (0, window_samples - len(segment)), mode='constant')
+
+        inputs = feature_extractor(
+            segment,
+            sampling_rate=sample_rate,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=1024
+        )
+
+        if use_half:
+            inputs = {k: v.to(device).half() for k, v in inputs.items()}
+        else:
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=-1)
+
+        pred_idx = torch.argmax(probs[0]).item()
+        pred_label = model.config.id2label.get(pred_idx, "")
+        pred_score = probs[0][pred_idx].item()
+
+        is_music = "music" in pred_label.lower()
+        is_uncertain = 0.40 <= pred_score <= 0.60
+
+        start_sec = start / sample_rate
+        end_sec = end / sample_rate
+
+        if is_uncertain and needs_padding:
+            if last_was_music:
+                music_segments.append((int(start_sec * 1000), int(end_sec * 1000), pred_score))
+        elif is_music and pred_score >= DETECTION_THRESHOLD:
+            music_segments.append((int(start_sec * 1000), int(end_sec * 1000), pred_score))
+            last_was_music = True
+        else:
+            last_was_music = False
+
+    return music_segments
+
+
+def merge_segments(segments):
+    if not segments:
+        return []
+
+    segments = sorted(segments, key=lambda x: x[0])
+    merged = [segments[0]]
+
+    for current in segments[1:]:
+        last = merged[-1]
+        if current[0] <= last[1]:
+            merged[-1] = (last[0], max(last[1], current[1]), max(last[2], current[2]))
+        else:
+            merged.append(current)
+
+    return merged
+
+
+def remove_music_segments(audio, segments):
+    if not segments:
+        return audio
+
+    clean_parts = []
+    last_end = 0
+
+    for start_ms, end_ms, _ in segments:
+        if start_ms > last_end:
+            clean_parts.append(audio[last_end:start_ms])
+        last_end = end_ms
+
+    if last_end < len(audio):
+        clean_parts.append(audio[last_end:])
+
+    if not clean_parts:
+        return AudioSegment.silent(duration=0)
+
+    return sum(clean_parts)
+
+
+def build_report(original_dur, clean_dur, segments):
+    removed = original_dur - clean_dur
+    pct = (removed / original_dur) * 100 if original_dur > 0 else 0
+
+    report = f"""## Processing Report
+
+| Metric | Value |
+|--------|-------|
+| Original Duration | {original_dur:.2f}s |
+| Clean Duration | {clean_dur:.2f}s |
+| Removed | {removed:.2f}s ({pct:.1f}%) |
+| Segments Found | {len(segments)} |
+| Output Format | WAV |
+"""
+
+    if segments:
+        report += "\n### Detected Music Segments\n| # | Start | End | Confidence |\n|---|-------|-----|------------|\n"
+        for i, (start_ms, end_ms, score) in enumerate(segments, 1):
+            report += f"| {i} | {start_ms/1000:.1f}s | {end_ms/1000:.1f}s | {score:.0%} |\n"
+    else:
+        report += "\n*No music detected in this audio.*\n"
+
+    return report
+
+
+@spaces.GPU if ZERO_GPU else lambda f: f
+def process_audio(audio_file, progress=gr.Progress()):
+    if audio_file is None:
+        return None, "Please upload an audio file."
+
+    try:
+        progress(0.1, desc="Preparing model...")
+        if torch.cuda.is_available():
+            model.to("cuda").half()
+            torch.backends.cudnn.benchmark = True
+
+        progress(0.2, desc="Loading audio...")
+        sample_rate = feature_extractor.sampling_rate
+        audio_array, audio = load_audio(audio_file, sample_rate)
+        original_duration = len(audio) / 1000
+
+        progress(0.4, desc="Detecting music...")
+        segments = detect_music(audio_array, sample_rate)
+        segments = merge_segments(segments)
+
+        progress(0.7, desc="Processing...")
+        clean_audio = remove_music_segments(audio, segments)
+        clean_duration = len(clean_audio) / 1000
+
+        progress(0.9, desc="Exporting...")
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+            clean_audio.export(f.name, format="wav")
+            output_path = f.name
+
+        progress(1.0, desc="Done")
+        report = build_report(original_duration, clean_duration, segments)
+
+        return output_path, report
+
+    except Exception as e:
+        logger.exception("Processing failed")
+        return None, f"Error: {str(e)}"
+
+
+with gr.Blocks(title="CleanSpeech AI") as demo:
+    gr.Markdown("""
+# CleanSpeech AI
+### Remove Background Music from Audio
+
+Upload your audio file to automatically detect and remove background music.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            audio_input = gr.Audio(label="Upload Audio", type="filepath")
+            process_btn = gr.Button("Remove Music", variant="primary", size="lg")
+
+        with gr.Column(scale=1):
+            audio_output = gr.Audio(label="Cleaned Audio")
+            report_output = gr.Markdown()
+
+    process_btn.click(
+        fn=process_audio,
+        inputs=[audio_input],
+        outputs=[audio_output, report_output]
+    )
+
+demo.queue()
+demo.launch()

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,8 @@
+transformers
+torch
+torchaudio
+gradio
+librosa
+soundfile
+numpy
+pydub