app.py

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(
            '<a href="https://ai-coustics.com/" target="_blank">'
            '<img src="https://mintcdn.com/ai-coustics/Sxcrv8jVSE2qWMR1/logo/dark.svg?fit=max&auto=format&n=Sxcrv8jVSE2qWMR1&q=85&s=7f26caaf21e963912961cbd8541e6d84" alt="ai-coustics Logo" width="400" style="display: block; margin: 0 auto;">'
            "</a>"
        )
    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", "/"])