Spaces:
Configuration error
Configuration error
| from __future__ import annotations | |
| import os | |
| import shutil | |
| import subprocess | |
| import tempfile | |
| import threading | |
| import uuid | |
| from pathlib import Path | |
| from typing import Optional | |
| import gradio as gr | |
| import numpy as np | |
| import soundfile as sf | |
| APP_NAME = "RSR AI" | |
| MAX_VIDEO_SECONDS = 180 | |
| AUDIO_SAMPLE_RATE = 48_000 | |
| # The model is downloaded and initialized lazily on the first enhancement request. | |
| _enhancer = None | |
| _model_lock = threading.Lock() | |
| def _run(command: list[str], *, timeout: int = 1200) -> subprocess.CompletedProcess: | |
| """Run an external media command and include useful stderr on failure.""" | |
| try: | |
| result = subprocess.run( | |
| command, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True, | |
| timeout=timeout, | |
| check=False, | |
| ) | |
| except FileNotFoundError as exc: | |
| raise RuntimeError( | |
| "FFmpeg/FFprobe is missing. This Space should install it from packages.txt; " | |
| "check the build logs and restart the Space." | |
| ) from exc | |
| except subprocess.TimeoutExpired as exc: | |
| raise RuntimeError("The media operation timed out. Try a shorter video.") from exc | |
| if result.returncode != 0: | |
| details = (result.stderr or result.stdout or "No details returned.")[-3500:] | |
| raise RuntimeError(f"Media processing failed. Details:\n{details}") | |
| return result | |
| def _probe_video(path: str) -> tuple[float, bool]: | |
| """Return container duration and whether the source has an audio stream.""" | |
| result = _run( | |
| [ | |
| "ffprobe", "-v", "error", "-show_entries", "format=duration", | |
| "-of", "default=noprint_wrappers=1:nokey=1", path, | |
| ], | |
| timeout=60, | |
| ) | |
| try: | |
| duration = float(result.stdout.strip()) | |
| except (TypeError, ValueError) as exc: | |
| raise RuntimeError("Could not read the video's duration.") from exc | |
| audio_result = _run( | |
| [ | |
| "ffprobe", "-v", "error", "-select_streams", "a:0", | |
| "-show_entries", "stream=index", "-of", "csv=p=0", path, | |
| ], | |
| timeout=60, | |
| ) | |
| return duration, bool(audio_result.stdout.strip()) | |
| def _get_enhancer(): | |
| global _enhancer | |
| if _enhancer is None: | |
| with _model_lock: | |
| if _enhancer is None: | |
| try: | |
| from clearvoice import ClearVoice | |
| except Exception as exc: | |
| raise RuntimeError( | |
| "Could not import ClearVoice. Check the Space build logs for " | |
| "dependency installation errors." | |
| ) from exc | |
| try: | |
| _enhancer = ClearVoice( | |
| task="speech_enhancement", | |
| model_names=["MossFormer2_SE_48K"], | |
| ) | |
| except Exception as exc: | |
| raise RuntimeError( | |
| "The open-source speech-enhancement model could not be loaded. " | |
| "The first run downloads its model weights, so check Space logs " | |
| "and confirm the Space has internet access and enough disk/RAM." | |
| ) from exc | |
| return _enhancer | |
| def _adjust_channels(audio: np.ndarray, target_channels: int) -> np.ndarray: | |
| """Adapt a model's output to the channel count used for the extracted source.""" | |
| if audio.ndim == 1: | |
| audio = audio[:, None] | |
| current_channels = audio.shape[1] | |
| if current_channels == target_channels: | |
| return audio | |
| if current_channels == 1 and target_channels > 1: | |
| return np.repeat(audio, target_channels, axis=1) | |
| if target_channels == 1: | |
| return np.mean(audio, axis=1, keepdims=True) | |
| if current_channels > target_channels: | |
| return audio[:, :target_channels] | |
| # Repeat the last channel only for unusual channel layouts. | |
| extra = np.repeat(audio[:, -1:], target_channels - current_channels, axis=1) | |
| return np.concatenate([audio, extra], axis=1) | |
| def _match_length(audio: np.ndarray, target_frames: int) -> np.ndarray: | |
| """Fit enhanced audio to the source audio's exact sample count.""" | |
| if audio.ndim == 1: | |
| audio = audio[:, None] | |
| if len(audio) == target_frames: | |
| return audio | |
| if len(audio) == 0: | |
| return np.zeros((target_frames, audio.shape[1]), dtype=np.float32) | |
| if len(audio) > target_frames: | |
| return audio[:target_frames] | |
| # Pad short model outputs without changing the original clip's timeline. | |
| missing = target_frames - len(audio) | |
| return np.pad(audio, ((0, missing), (0, 0)), mode="constant") | |
| def _safe_name(filename: str) -> str: | |
| stem = Path(filename).stem or "video" | |
| cleaned = "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in stem) | |
| return cleaned[:70] or "video" | |
| def enhance_video(video_file: Optional[str], strength: float, mode: str, keep_timing: bool): | |
| """Enhance speech audio and remux it with the original, un-reencoded video stream.""" | |
| if not video_file: | |
| raise gr.Error("Upload a video first.") | |
| source = Path(video_file) | |
| if not source.is_file(): | |
| raise gr.Error("The uploaded video could not be found. Please upload it again.") | |
| try: | |
| duration, has_audio = _probe_video(str(source)) | |
| except Exception as exc: | |
| raise gr.Error(str(exc)) from exc | |
| if not has_audio: | |
| raise gr.Error("This video has no audio track to repair.") | |
| if duration <= 0: | |
| raise gr.Error("The video's duration is invalid.") | |
| if duration > MAX_VIDEO_SECONDS: | |
| raise gr.Error( | |
| f"This demo is limited to {MAX_VIDEO_SECONDS} seconds per video to fit free Space resources. " | |
| "Trim the clip and try again." | |
| ) | |
| # Keep job artifacts alive for Gradio's preview and download components. | |
| job_dir = Path(tempfile.gettempdir()) / "rsr-ai-jobs" / uuid.uuid4().hex | |
| job_dir.mkdir(parents=True, exist_ok=True) | |
| source_audio = job_dir / "source_48k.wav" | |
| raw_enhanced = job_dir / "enhanced_model.wav" | |
| mixed_audio = job_dir / "enhanced_audio.wav" | |
| try: | |
| # The model's full-band enhancement checkpoint is trained for 48 kHz audio. | |
| _run( | |
| [ | |
| "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", | |
| "-i", str(source), "-map", "0:a:0", "-vn", | |
| "-ar", str(AUDIO_SAMPLE_RATE), "-ac", "2", | |
| "-c:a", "pcm_s16le", str(source_audio), | |
| ] | |
| ) | |
| if float(strength) <= 0: | |
| shutil.copyfile(source_audio, mixed_audio) | |
| else: | |
| enhancer = _get_enhancer() | |
| try: | |
| result = enhancer(input_path=str(source_audio), online_write=False) | |
| enhancer.write(result, output_path=str(raw_enhanced)) | |
| except Exception as exc: | |
| raise RuntimeError( | |
| "ClearVoice could not enhance this clip. Try a shorter/common MP4 file " | |
| "or inspect the Space logs for model errors." | |
| ) from exc | |
| if not raw_enhanced.is_file(): | |
| raise RuntimeError( | |
| "The enhancement model finished without producing a WAV file. " | |
| "Check the ClearVoice output in the Space logs." | |
| ) | |
| original, original_rate = sf.read(str(source_audio), dtype="float32", always_2d=True) | |
| enhanced, enhanced_rate = sf.read(str(raw_enhanced), dtype="float32", always_2d=True) | |
| if original_rate != enhanced_rate: | |
| raise RuntimeError( | |
| f"The model returned {enhanced_rate} Hz audio but {original_rate} Hz was expected." | |
| ) | |
| enhanced = _adjust_channels(enhanced, original.shape[1]) | |
| enhanced = _match_length(enhanced, len(original)) | |
| blend = min(1.0, max(0.0, float(strength) / 100.0)) | |
| output_audio = (1.0 - blend) * original + blend * enhanced | |
| output_audio = np.nan_to_num(output_audio, nan=0.0, posinf=0.0, neginf=0.0) | |
| output_audio = np.clip(output_audio, -1.0, 1.0) | |
| sf.write(str(mixed_audio), output_audio, original_rate, subtype="PCM_16") | |
| ext = source.suffix.lower() | |
| basename = _safe_name(source.name) | |
| if ext == ".webm": | |
| output_video = job_dir / f"{basename}_rsr.webm" | |
| audio_codec = ["-c:a", "libopus", "-b:a", "160k"] | |
| container_options: list[str] = [] | |
| elif ext in {".mp4", ".m4v"}: | |
| output_video = job_dir / f"{basename}_rsr.mp4" | |
| audio_codec = ["-c:a", "aac", "-b:a", "192k"] | |
| container_options = ["-movflags", "+faststart"] | |
| elif ext == ".mov": | |
| output_video = job_dir / f"{basename}_rsr.mov" | |
| audio_codec = ["-c:a", "aac", "-b:a", "192k"] | |
| container_options = [] | |
| else: | |
| # Matroska can carry a broad range of original video codecs without re-encoding. | |
| output_video = job_dir / f"{basename}_rsr.mkv" | |
| audio_codec = ["-c:a", "aac", "-b:a", "192k"] | |
| container_options = [] | |
| remux_command = [ | |
| "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", | |
| "-i", str(source), "-i", str(mixed_audio), | |
| "-map", "0:v:0", "-map", "1:a:0", "-map_metadata", "0", | |
| "-c:v", "copy", *audio_codec, | |
| ] | |
| if keep_timing: | |
| # apad + shortest stops at the end of the unchanged source video stream. | |
| remux_command += ["-af", "apad", "-shortest"] | |
| remux_command += container_options + [str(output_video)] | |
| _run(remux_command) | |
| if not output_video.is_file() or output_video.stat().st_size == 0: | |
| raise RuntimeError("FFmpeg did not create the enhanced video file.") | |
| timing_note = ( | |
| "Original video stream copied without re-encoding; frame pixels and source timestamps are retained." | |
| if keep_timing | |
| else "Original video stream copied without re-encoding." | |
| ) | |
| status = ( | |
| f"### ✅ Audio enhancement complete\n\n" | |
| f"- **Preset:** {mode}\n" | |
| f"- **Removal strength:** {int(round(float(strength)))}%\n" | |
| f"- **Clip duration:** {duration:.2f} seconds\n" | |
| f"- **Video:** {timing_note}\n\n" | |
| "The model targets speech enhancement and noise reduction; severe vocoder/robotic artifacts may remain." | |
| ) | |
| return str(output_video), str(mixed_audio), str(output_video), status | |
| except gr.Error: | |
| raise | |
| except Exception as exc: | |
| raise gr.Error(str(exc)) from exc | |
| CSS = """ | |
| .gradio-container {max-width: 1120px !important; margin: 0 auto !important;} | |
| #rsr-hero {padding: 24px 26px; border-radius: 22px; background: linear-gradient(135deg, #0c1830 0%, #122d4b 55%, #174766 100%); color: #f5fbff; border: 1px solid #28536d;} | |
| #rsr-hero h1 {margin: 0 0 6px 0; font-size: clamp(28px, 4vw, 42px); letter-spacing: -0.8px;} | |
| #rsr-hero p {margin: 0; color: #d4eaf8; max-width: 780px; line-height: 1.55;} | |
| .rsr-subtle {color: var(--body-text-color-subdued); font-size: 0.92rem;} | |
| """ | |
| with gr.Blocks(title="RSR AI — Robotic Sounds Remover AI") as demo: | |
| gr.HTML( | |
| """<section id="rsr-hero"> | |
| <h1>🔊 RSR AI</h1> | |
| <p><strong>Robotic Sounds Remover AI</strong> — enhance speech in a video with an open-source AI audio-restoration model, then attach the improved audio to the original video stream.</p> | |
| </section>""" | |
| ) | |
| gr.Markdown( | |
| "Upload a video with speech. RSR AI processes its audio track and remuxes it with the original video stream. " | |
| "The model loads on first use, so the first job may take longer." | |
| ) | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=1): | |
| video_in = gr.File( | |
| label="Upload video", | |
| file_types=[".mp4", ".m4v", ".mov", ".webm", ".mkv", ".avi"], | |
| type="filepath", | |
| ) | |
| strength = gr.Slider( | |
| minimum=0, | |
| maximum=100, | |
| value=100, | |
| step=5, | |
| label="Robotic / artifact reduction strength", | |
| info="100% uses the full enhanced signal from the model.", | |
| ) | |
| mode = gr.Dropdown( | |
| choices=["Balanced"], | |
| value="Balanced", | |
| label="Enhancement mode", | |
| info="Balanced speech enhancement with MossFormer2_SE_48K.", | |
| interactive=False, | |
| ) | |
| keep_timing = gr.Checkbox( | |
| value=True, | |
| label="Keep original video frames and timing", | |
| info="The video stream is copied without re-encoding; only audio is replaced.", | |
| ) | |
| process_btn = gr.Button("✨ Enhance audio", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| video_out = gr.Video(label="Enhanced video preview", height=320) | |
| audio_out = gr.Audio(label="Enhanced audio preview", type="filepath") | |
| download_out = gr.File(label="Download enhanced video") | |
| status = gr.Markdown() | |
| process_btn.click( | |
| fn=enhance_video, | |
| inputs=[video_in, strength, mode, keep_timing], | |
| outputs=[video_out, audio_out, download_out, status], | |
| ) | |
| with gr.Accordion("What RSR AI does and its limits", open=False): | |
| gr.Markdown( | |
| "- **Model:** [ClearVoice / MossFormer2_SE_48K](https://huggingface.co/alibabasglab/MossFormer2_SE_48K), an open-source 48 kHz speech-enhancement model.\n" | |
| "- **Video preservation:** FFmpeg copies the source video stream (`-c:v copy`); it does not re-render or regenerate video frames.\n" | |
| "- **Audio:** The source track is extracted at 48 kHz, enhanced, fitted back to the source audio's sample count, then muxed into the video.\n" | |
| "- **Important limit:** This model is designed for speech enhancement/noise reduction, not specifically trained to reverse every synthetic voice or vocoder artifact. Results vary with the clip and it can alter music, sound effects, or speech. Keep a copy of your original.\n" | |
| f"- **Free-resource limit:** Videos are limited to {MAX_VIDEO_SECONDS} seconds per request. CPU processing may be slow." | |
| ) | |
| if __name__ == "__main__": | |
| # The queue prevents simultaneous model calls from exhausting a free Space instance. | |
| demo.queue(default_concurrency_limit=1).launch( | |
| max_file_size="500mb", theme=gr.themes.Soft(), css=CSS | |
| ) | |