from audio_denoiser.AudioDenoiser import AudioDenoiser
from timeit import default_timer as timer
from datetime import datetime
import gradio as gr
from torchaudio import AudioMetaData
import torch
import tempfile
import os,pytz

tz = pytz.timezone('Asia/Singapore')
theme='remilia/Ghostly'
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')


def denoiser(win, auto_scale):
    if win is None:
      gr.Warning('Audio does not exist. Please ensure that the audio has been successfully uploaded.')
      return None,None
    
    startTime=timer()
    denoiser = AudioDenoiser(device=device)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
    wout = temp_file.name
    temp_file.close() 
    try:
      denoiser.process_audio_file(win, wout, auto_scale)
    except RuntimeError as e:
      gr.Warning(str(e))
      return None,None
    
    endTime=timer()
    info=(f'🆗Completion time: {round(endTime-startTime,4)}s')
    now=datetime.now(tz).strftime('%H:%M:%S')
    print(f'{now}-{info}')
    return wout,info

examples = [
    ["sample/exp1.wav", "sample/exp1_d.wav"],
    ["sample/exp2.wav", "sample/exp2_d.wav"]
]
with gr.Blocks(theme=theme) as app:
    gr.HTML('''
    <h1 style="font-size: 25px;">Audio Denoiser</h1>
    <p style="margin-bottom: 10px; font-size: 100%">
 
Originating from the project: <a href='https://github.com/jose-solorzano/audio-denoiser'>audio-denoiser</a><br>
model: <a href='https://huggingface.co/jose-h-solorzano/audio-denoiser-512-32-v1'>jose-h-solorzano/audio-denoiser-512-32-v1</a><br>

    </p>''')
    audio_in = gr.Audio(type="filepath", label="Upload Audio")
    btn=gr.Button(value='💥Remove Noise',variant="primary")
    scale = gr.Checkbox(
                label="Auto Scale",
                info='Recommended for low-volume input audio',
                value=True,
            )
    audio_out = gr.Audio(type="filepath", label="Denoised Audio")
    info=gr.Textbox(label='info')

    btn.click(
      denoiser,
      inputs=[audio_in, scale],
      outputs=[audio_out,info])

    examples = gr.Examples(examples=examples, inputs=audio_in, outputs=audio_out, fn=denoiser)

     

app.launch(share=True)