from __future__ import annotations

import gradio as gr
try:
    import spaces
except ImportError:  # 'spaces' is only provided by Hugging Face Spaces
    import types as _types

    def _gpu(*args, **kwargs):
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]

        def _decorator(func):
            return func

        return _decorator

    spaces = _types.SimpleNamespace(GPU=_gpu)

from pyharp import *


import tempfile
from voicefixer import VoiceFixer

# Initialize VoiceFixer. It handles downloading checkpoints and setting device.
# The VoiceFixer class automatically detects and uses CUDA if available.
voicefixer_model = VoiceFixer()


model_card = ModelCard(
    name="VoiceFixer",
    description="VoiceFixer aims to restore human speech regardless how serious its degraded. It can handle noise, reverberation, low resolution (2kHz~44.1kHz) and clipping (0.1-1.0 threshold) effect within one model.",
    author="haoheliu",
    tags=["declipping", "denoise", "dereverberation", "mel", "speech", "speech-analysis", "speech-enhancement", "speech-processing", "speech-synthesis", "super-resolution", "tts", "vocoder"],
)


@spaces.GPU
def process_fn(input_audio, mode):
    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    voicefixer_model.restore(input_audio, output_file, mode=int(mode))
    return output_file


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info("Upload an audio file to be processed by VoiceFixer."),
        gr.Dropdown(choices=["0", "1", "2"], value="0", label="Processing Mode", info="Select the VoiceFixer processing mode:\n0: Original Model (suggested by default)\n1: Add preprocessing module (remove higher frequency)\n2: Train mode (might work sometimes on seriously degraded real speech)"),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Fixed Audio").set_info("The enhanced audio output from VoiceFixer."),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch(share=True, show_error=False, pwa=True)