Spaces:
Paused
Paused
| import gradio as gr | |
| import subprocess | |
| import os | |
| import numpy as np | |
| import librosa | |
| import soundfile as sf | |
| import matplotlib.pyplot as plt | |
| import librosa.display | |
| import gc | |
| import torch | |
| import noisereduce as nr | |
| import warnings | |
| from scipy import signal | |
| from scipy.stats import kurtosis, skew | |
| import spaces | |
| import urllib.request | |
| warnings.filterwarnings("ignore") | |
| # Set environment variables and torch settings | |
| os.environ["TOKENIZERS_PARALLELISM"] = "true" | |
| torch.set_float32_matmul_precision("high") | |
| # Check for GPU availability | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") | |
| # Create output directory | |
| output_folder = "output_file" | |
| os.makedirs(output_folder, exist_ok=True) | |
| print(f"Output folder ready: {output_folder}") | |
| # Setup function to download models and configs using wget or urllib | |
| def setup(): | |
| # Create Apollo directory structure | |
| os.makedirs("Apollo/model", exist_ok=True) | |
| os.makedirs("Apollo/configs", exist_ok=True) | |
| # Define files and their URLs | |
| files_to_download = { | |
| "Apollo/inference.py": "https://raw.githubusercontent.com/jarredou/Apollo-Colab-Inference/main/inference.py", | |
| "Apollo/model/pytorch_model.bin": "https://huggingface.co/JusperLee/Apollo/resolve/main/pytorch_model.bin", | |
| "Apollo/model/apollo_model.ckpt": "https://huggingface.co/jarredou/lew_apollo_vocal_enhancer/resolve/main/apollo_model.ckpt", | |
| "Apollo/model/apollo_model_v2.ckpt": "https://huggingface.co/jarredou/lew_apollo_vocal_enhancer/resolve/main/apollo_model_v2.ckpt", | |
| "Apollo/model/apollo_universal_model.ckpt": "https://huggingface.co/ASesYusuf1/Apollo_universal_model/resolve/main/apollo_universal_model.ckpt", | |
| "Apollo/configs/config_apollo_vocal.yaml": "https://huggingface.co/jarredou/lew_apollo_vocal_enhancer/resolve/main/config_apollo_vocal.yaml", | |
| "Apollo/configs/config_apollo.yaml": "https://huggingface.co/ASesYusuf1/Apollo_universal_model/resolve/main/config_apollo.yaml", | |
| "Apollo/configs/apollo.yaml": "https://huggingface.co/JusperLee/Apollo/resolve/main/apollo.yaml", | |
| } | |
| # Download files if they don't exist | |
| for file_path, url in files_to_download.items(): | |
| if not os.path.exists(file_path): | |
| print(f"Downloading {file_path}...") | |
| try: | |
| # Try wget first | |
| subprocess.run(["wget", "-O", file_path, url], check=True, capture_output=True, text=True) | |
| print(f"Downloaded {file_path} with wget") | |
| except (subprocess.CalledProcessError, FileNotFoundError) as e: | |
| print(f"wget failed for {file_path}: {e}. Falling back to urllib...") | |
| try: | |
| urllib.request.urlretrieve(url, file_path) | |
| print(f"Downloaded {file_path} with urllib") | |
| except Exception as e: | |
| print(f"Failed to download {file_path}: {e}") | |
| raise Exception(f"Failed to download {file_path}") | |
| # Run setup | |
| try: | |
| setup() | |
| except Exception as e: | |
| print(f"Setup failed: {e}") | |
| raise | |
| # Processing function | |
| def process_audio(input_file, model, chunk_size, overlap): | |
| if not input_file: | |
| return "No file uploaded.", None, None | |
| input_file_path = input_file | |
| original_file_name = os.path.splitext(os.path.basename(input_file_path))[0] | |
| output_file_path = f'{output_folder}/{original_file_name}_processed.wav' | |
| # Model selection | |
| model_paths = { | |
| 'MP3 Enhancer': ('Apollo/model/pytorch_model.bin', 'Apollo/configs/apollo.yaml'), | |
| 'Lew Vocal Enhancer': ('Apollo/model/apollo_model.ckpt', 'Apollo/configs/apollo.yaml'), | |
| 'Lew Vocal Enhancer v2 (beta)': ('Apollo/model/apollo_model_v2.ckpt', 'Apollo/configs/config_apollo_vocal.yaml'), | |
| 'Apollo Universal Model': ('Apollo/model/apollo_universal_model.ckpt', 'Apollo/configs/config_apollo.yaml') | |
| } | |
| if model not in model_paths: | |
| return "Invalid model selected.", None, None | |
| ckpt, config = model_paths[model] | |
| if not os.path.exists(ckpt) or not os.path.exists(config): | |
| return f"Model files not found: {ckpt} or {config}", None, None | |
| print(f"Model selected: {model}") | |
| print("Processing started. Please wait...") | |
| command = f"python Apollo/inference.py --in_wav '{input_file_path}' --out_wav '{output_file_path}' --chunk_size {chunk_size} --overlap {overlap} --ckpt '{ckpt}' --config '{config}'" | |
| try: | |
| process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) | |
| for line in process.stdout: | |
| print(f"Processing: {line.strip()}") | |
| process.stdout.close() | |
| process.wait() | |
| if process.returncode != 0: | |
| return "An error occurred while processing the audio.", None, None | |
| print("Processing completed.") | |
| return output_file_path, input_file_path, None | |
| except Exception as e: | |
| return f"Error: {str(e)}", None, None | |
| # Mid/Side separation and combination | |
| def mid_side_separation(audio_file): | |
| try: | |
| y, sr = librosa.load(audio_file, sr=None, mono=False) | |
| if y.ndim == 1: | |
| raise ValueError("Stereo audio file required!") | |
| left, right = y[0], y[1] | |
| mid = (left + right) / 2 | |
| side = (left - right) / 2 | |
| mid_path = os.path.join(output_folder, "mid.wav") | |
| side_path = os.path.join(output_folder, "side.wav") | |
| sf.write(mid_path, mid, sr) | |
| sf.write(side_path, side, sr) | |
| return mid_path, side_path, sr | |
| except Exception as e: | |
| raise ValueError(f"Error in mid/side separation: {str(e)}") | |
| def mid_side_combine(mid_file, side_file, output_path): | |
| try: | |
| mid_data, sr_mid = librosa.load(mid_file, sr=None, mono=True) | |
| side_data, sr_side = librosa.load(side_file, sr=None, mono=True) | |
| if sr_mid != sr_side: | |
| raise ValueError("Mid and Side sample rates do not match!") | |
| left = mid_data + side_data | |
| right = mid_data - side_data | |
| stereo = np.stack([left, right], axis=0) | |
| sf.write(output_path, stereo.T, sr_mid) | |
| return output_path | |
| except Exception as e: | |
| raise ValueError(f"Error in mid/side combination: {str(e)}") | |
| def process_mid_side_upscale(input_file, model, chunk_size, overlap): | |
| if not input_file: | |
| return "No file uploaded.", None, None | |
| try: | |
| print("Separating Mid and Side channels...") | |
| mid_path, side_path, sr = mid_side_separation(input_file) | |
| print("Processing Mid channel...") | |
| mid_restored, _, _ = process_audio(mid_path, model, chunk_size, overlap) | |
| if not mid_restored.endswith(".wav"): | |
| return mid_restored, None, None | |
| print("Processing Side channel...") | |
| side_restored, _, _ = process_audio(side_path, model, chunk_size, overlap) | |
| if not side_restored.endswith(".wav"): | |
| return side_restored, None, None | |
| original_file_name = os.path.splitext(os.path.basename(input_file))[0] | |
| final_output_path = os.path.join(output_folder, f"{original_file_name}_upscaled.wav") | |
| print("Combining processed Mid and Side channels...") | |
| final_audio = mid_side_combine(mid_restored, side_restored, final_output_path) | |
| print("Mid/Side upscaling completed.") | |
| return final_audio, input_file, None | |
| except Exception as e: | |
| return f"Error: {str(e)}", None, None | |
| # Spectrum analysis | |
| def spectrum(audio_file): | |
| if not audio_file: | |
| return None, "No file selected" | |
| try: | |
| chunk_duration = 30 # Reduced for ZeroGPU memory constraints | |
| hop_length = 512 | |
| n_fft = 2048 # Reduced for ZeroGPU memory constraints | |
| with sf.SoundFile(audio_file) as sf_desc: | |
| duration = len(sf_desc) / sf_desc.samplerate | |
| num_chunks = int(np.ceil(duration / chunk_duration)) | |
| freqs = librosa.fft_frequencies(sr=sf_desc.samplerate, n_fft=n_fft) | |
| total_frames = int(np.ceil(duration * sf_desc.samplerate / hop_length)) | |
| S_db_full = np.zeros((len(freqs), total_frames)) | |
| for chunk_idx in range(num_chunks): | |
| start_time = chunk_idx * chunk_duration | |
| y, sr = librosa.load(audio_file, offset=start_time, duration=chunk_duration, sr=None) | |
| S_chunk = np.abs(librosa.stft(y, n_fft=n_fft, hop_length=hop_length)) | |
| S_db_chunk = librosa.amplitude_to_db(S_chunk, ref=np.max) | |
| start_frame = int(start_time * sr / hop_length) | |
| end_frame = start_frame + S_db_chunk.shape[1] | |
| S_db_full[:, start_frame:end_frame] = S_db_chunk | |
| del S_chunk, S_db_chunk | |
| gc.collect() | |
| downsample_factor = 4 | |
| S_db_downsampled = S_db_full[:, ::downsample_factor] | |
| threshold = np.max(S_db_downsampled) - 60 | |
| significant_freqs = freqs[np.any(S_db_downsampled > threshold, axis=1)] | |
| max_freq = np.max(significant_freqs) if len(significant_freqs) > 0 else sr / 2 | |
| plt.figure(figsize=(15, 8)) # Reduced size for ZeroGPU | |
| display_hop = 4 | |
| librosa.display.specshow( | |
| S_db_full[:, ::display_hop], | |
| sr=sr, | |
| hop_length=hop_length * display_hop, | |
| x_axis='time', | |
| y_axis='hz', | |
| cmap='magma' | |
| ) | |
| freq_ticks = [2000, 4000, 6000, 8000, 10000, 12000, 14000, 16000, 18000, 20000] | |
| plt.yticks(freq_ticks, [f"{f/1000:.0f}" for f in freq_ticks]) | |
| plt.colorbar(format='%+2.0f dB') | |
| plt.title('Frequency Spectrum', fontsize=16) | |
| plt.xlabel('Time (seconds)', fontsize=12) | |
| plt.ylabel('Frequency (kHz)', fontsize=12) | |
| output_image_path = os.path.join(output_folder, 'spectrum.png') | |
| plt.savefig(output_image_path, bbox_inches='tight', dpi=150) # Reduced DPI for ZeroGPU | |
| plt.close() | |
| del S_db_full, S_db_downsampled | |
| gc.collect() | |
| closest_freq = min(freq_ticks, key=lambda x: abs(x - max_freq)) | |
| return output_image_path, f"Maximum Frequency {int(closest_freq)} Hz" | |
| except Exception as e: | |
| return None, f"Error: {str(e)}" | |
| # Credits | |
| def show_credits(): | |
| return """This Web UI was created using AI tools and written by U.Z.S. | |
| **Apollo-Colab-Inference** (https://github.com/jarredou/Apollo-Colab-Inference): | |
| This project was developed by Jarred Ou and provides a colab-based inference implementation of the Apollo model for audio enhancement. | |
| **Apollo** (https://github.com/JusperLee/Apollo): | |
| Created by Jusper Lee, Apollo is a deep learning-based model aimed at improving vocal clarity and overall audio quality in recordings. | |
| """ | |
| # Gradio Interface | |
| app = gr.Blocks(css=""" | |
| .gradio-container { background-color: black; color: white; font-family: Arial, sans-serif; } | |
| .footer { position: absolute; bottom: 10px; right: 10px; font-size: 12px; color: white; } | |
| .gradio-button { background-color: #6a0dad; color: white; border: 1px solid #5a0b8a; border-radius: 5px; } | |
| .gradio-button:hover { background-color: #5a0b8a; } | |
| .gradio-input { background-color: rgba(106, 13, 173, 0.8); border: 1px solid #5a0b8a; color: white; border-radius: 5px; } | |
| .gradio-input:focus { border-color: #ffffff; box-shadow: 0 0 5px rgba(255, 255, 255, 0.5); } | |
| .gradio-slider { background-color: rgba(106, 13, 173, 0.8); color: white; } | |
| .gradio-label { color: white; } | |
| .gradio-tabs { background-color: rgba(106, 13, 173, 0.8); color: white; } | |
| @media (max-width: 600px) { | |
| .gradio-button { width: 100%; font-size: 16px; } | |
| .gradio-input { width: 100%; font-size: 16px; } | |
| .gradio-slider { width: 100%; } | |
| .gradio-label { font-size: 14px; } | |
| } | |
| """) | |
| with app: | |
| with gr.Tab("Home"): | |
| gr.Markdown("# Apollo Audio Enhancement") | |
| with gr.Row(): | |
| audio_input = gr.File(label="Select Audio File", file_types=["audio"]) | |
| model = gr.Radio( | |
| ["MP3 Enhancer", "Lew Vocal Enhancer", "Lew Vocal Enhancer v2 (beta)", "Apollo Universal Model"], | |
| label="Select Model" | |
| ) | |
| gr.Markdown("**For Universal model, please set Chunk_Size to 19**", elem_classes="model-note") | |
| chunk_size = gr.Slider(minimum=3, maximum=25, step=1, value=25, label="Chunk Size") | |
| overlap = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Overlap") | |
| output_audio = gr.Audio(label="Processed Audio") | |
| original_audio = gr.Audio(label="Original Audio") | |
| error_message = gr.Textbox(label="Status") | |
| process_button = gr.Button("Process Audio") | |
| process_button.click(process_audio, inputs=[audio_input, model, chunk_size, overlap], outputs=[output_audio, original_audio, error_message]) | |
| with gr.Tab("Spectrum"): | |
| gr.Markdown("# Spectrum Analysis") | |
| spectrogram_input = gr.File(label="Select Audio File for Spectrum", file_types=["audio"]) | |
| output_spectrum = gr.Image(label="Frequency Spectrum") | |
| max_freq_info = gr.Textbox(label="Maximum Frequency Information") | |
| spectrum_button = gr.Button("Show Spectrum") | |
| spectrum_button.click(spectrum, inputs=[spectrogram_input], outputs=[output_spectrum, max_freq_info]) | |
| with gr.Tab("Mid/Side Upscale"): | |
| gr.Markdown("# 🎚️ Mid/Side Audio Upscaling") | |
| gr.Markdown("Upload a stereo audio file to separate, enhance, and recombine its Mid and Side channels using Apollo.") | |
| with gr.Row(): | |
| ms_input = gr.File(label="Select Stereo Audio File", file_types=["audio"]) | |
| ms_model = gr.Radio( | |
| ["MP3 Enhancer", "Lew Vocal Enhancer", "Lew Vocal Enhancer v2 (beta)", "Apollo Universal Model"], | |
| label="Select Model", | |
| value="Apollo Universal Model" | |
| ) | |
| ms_chunk_size = gr.Slider(minimum=3, maximum=25, step=1, value=18, label="Chunk Size") | |
| ms_overlap = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Overlap") | |
| ms_output = gr.Audio(label="Upscaled Audio") | |
| ms_original = gr.Audio(label="Original Audio") | |
| ms_error_message = gr.Textbox(label="Status") | |
| ms_process_button = gr.Button("Process Mid/Side Upscale") | |
| ms_process_button.click( | |
| process_mid_side_upscale, | |
| inputs=[ms_input, ms_model, ms_chunk_size, ms_overlap], | |
| outputs=[ms_output, ms_original, ms_error_message] | |
| ) | |
| with gr.Tab("Credits"): | |
| gr.Markdown("## Credits") | |
| gr.Markdown(show_credits()) | |
| gr.Markdown("Developed by U.Z.S using ChatGPT.", elem_classes="footer") | |
| app.launch() |