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( """

🔊 RSR AI

Robotic Sounds Remover AI — enhance speech in a video with an open-source AI audio-restoration model, then attach the improved audio to the original video stream.

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