import os import time from typing import Optional, Any import gradio as gr from loguru import logger from constants import ( MINUTES_KEEP, ) from sdk import SDKParams, process_file_sdk from audio_tools import spec_image, mix_at_snr import shutil import tempfile # =============================== # Temporary File & Cache Management # =============================== def cleanup_tmp(minutes_keep: int = MINUTES_KEEP, filter: list[str] = []): skipped = 0 removed = 0 if os.path.exists("/tmp"): for root, _, files in os.walk("/tmp"): for name in files: f = os.path.join(root, name) is_old = (time.time() - os.path.getmtime(f)) / 60 > minutes_keep filtered = any(filt in f for filt in filter) if filtered: skipped += 1 continue if not is_old: skipped += 1 continue try: os.remove(f) removed += 1 except Exception as e: logger.warning(f"Failed to remove file {f}: {e}") logger.info(f"Cleanup tmp complete. Removed {removed} files, skipped {skipped} files.") # =============================== # Interface Logic # =============================== def denoise_audio( sample_path: str, enhancement_level: float = 50.0, ) -> tuple[Optional[str], Optional[str], Optional[str], Optional[str]]: gr.Info( "Processing started. This may take a moment. Please do not refresh or close the window." ) base, ext = os.path.splitext(sample_path) enhanced_path = f"{base}_enhanced{ext}" noisy_path = f"{base}_noisy{ext}" noisy_spec_path = f"{base}_noisy_spectrogram.png" enhanced_spec_path = f"{base}_enhanced_spectrogram.png" noisy_path = sample_path try: sdk_key = os.getenv("SECRET_SDK_KEY") sdk_params = SDKParams( enhancement_level=enhancement_level / 100, sdk_key=sdk_key, ) process_file_sdk(noisy_path, enhanced_path, sdk_params) except Exception as e: gr.Warning(f"{e}") delete_related_files(sample_path) return None, None, None, None noisy_im = spec_image(noisy_path) noisy_im.save(noisy_spec_path) enhanced_im = spec_image(enhanced_path) enhanced_im.save(enhanced_spec_path) print(f"Enhancement complete. id: {base}") return noisy_path, noisy_spec_path, enhanced_path, enhanced_spec_path def toggle_SNR(choice: str): if choice == "None": return gr.update(visible=False, value="None") else: return gr.update(visible=True, value="10") def delete_related_files(path: str): """ Deletes all files in /tmp containing the base filename of the given path. """ filename_no_ext = os.path.splitext(os.path.basename(path))[0] base_dir = "/tmp" deleted = 0 for root, _, files in os.walk(base_dir): for f in files: if filename_no_ext in f: full_path = os.path.join(root, f) try: os.remove(full_path) deleted += 1 except Exception as e: logger.warning(f"Failed to delete file {full_path}: {e}") if deleted == 0: logger.info(f"No files found to delete containing '{filename_no_ext}' in {base_dir}") logger.info(f"Deleted {deleted} files related to the last enhancement '{filename_no_ext}'.") def cleanup(last_enhancement: str, last_audio_file: str = "", new_audio_file: str = ""): """ Deletes - all enhancement files (usually .png & .wav) in the /tmp directory from the last enhancement - the last uploaded audio file if a new one is uploaded - other files in /tmp older than 2 hours to help manage disk space. Note: Consider adding a flag to files to indicate whether files can be deleted, if disk usage becomes an issue. """ # Delete the last uploaded audio file if a new one is uploaded if last_audio_file and last_audio_file != new_audio_file and os.path.exists(last_audio_file): try: os.remove(last_audio_file) logger.info(f"Deleted last uploaded audio file: {last_audio_file}") except Exception as e: logger.warning(f"Failed to delete last uploaded audio file {last_audio_file}: {e}") if last_enhancement: delete_related_files(last_enhancement) cleanup_tmp(minutes_keep=120) def start_processing(sample_path: str) -> tuple[Any, str, str]: success = True if not sample_path or not os.path.exists(sample_path): gr.Warning("Please upload an audio sample or use the microphone input.") success = False if not os.getenv("SECRET_SDK_KEY"): gr.Warning("No SDK key provided. Please contact us at https://ai-coustics.com/contact/.") success = False if not success: raise ValueError("Missing audio sample or API/SDK key. Processing cannot start.") # Generate a new /tmp path with a unique filename using tempfile ext = os.path.splitext(sample_path)[1] with tempfile.NamedTemporaryFile(delete=False, suffix=ext, dir="/tmp") as tmp_file: shutil.copy(sample_path, tmp_file.name) input_enhancement_path = tmp_file.name return ( gr.update(interactive=False), input_enhancement_path, sample_path, ) def enable_new_input(): return gr.update(visible=True) # =============================== # Gradio UI Layout # =============================== with gr.Blocks() as demo: input_enhancement = gr.State() last_audio_file = gr.State() with gr.Row(): gr.HTML( '' 'ai-coustics Logo' "" ) with gr.Row(): gr.Markdown(open("intro.md").read()) with gr.Row(): with gr.Column(): audio_file = gr.Audio( type="filepath", label="Input", visible=True, sources=["upload", "microphone"] ) percent_slider = gr.Slider( minimum=1, maximum=100, value=100, step=1, label="Enhancement Level (%)", info=( "Set how much enhancement to apply. " "Lower values are more subtle, higher values are stronger." ), ) btn = gr.Button("Enhance") with gr.Column(): noisy_audio = gr.Audio(type="filepath", label="Noisy audio") noisy_image = gr.Image(label="Noisy spectrogram", format="png", type="filepath") enhanced_audio = gr.Audio(type="filepath", label="Enhanced audio") enhanced_image = gr.Image(label="Enhanced spectrogram", format="png", type="filepath") btn.click(cleanup, [input_enhancement, last_audio_file, audio_file], None).then( start_processing, inputs=audio_file, outputs=[btn, input_enhancement, last_audio_file], ).success( denoise_audio, inputs=[ input_enhancement, percent_slider, ], outputs=[noisy_audio, noisy_image, enhanced_audio, enhanced_image], ).then(enable_new_input, None, btn) cleanup_tmp(minutes_keep=0, filter=[]) demo.launch(allowed_paths=["/tmp", "/"])