import gradio as gr
import torchaudio
import resampy
import numpy as np
import soundfile as sf
import torch

def upsample_audio(audio_path, backend="CPU"):
    # Read audio file
    audio, orig_sr = sf.read(audio_path)
    
    # Target sample rate (2x upsampling)
    target_sr = orig_sr * 2
    
    if backend == "CPU":
        # Resampy CPU processing
        upsampled = resampy.resample(audio, orig_sr, target_sr, filter="kaiser_best")
    else:
        # Torchaudio GPU processing
        device = "cuda" if torch.cuda.is_available() else "cpu"
        tensor_audio = torch.from_numpy(audio).float().to(device)
        # If mono, add batch dimension
        if tensor_audio.ndim == 1:
            tensor_audio = tensor_audio.unsqueeze(0)
        resampler = torchaudio.transforms.Resample(
            orig_freq=orig_sr,
            new_freq=target_sr,
            resampling_method="kaiser_window",
            lowpass_filter_width=64
        ).to(device)
        upsampled = resampler(tensor_audio).cpu().numpy()
        # Remove batch dimension if mono
        if upsampled.shape[0] == 1:
            upsampled = upsampled[0]

    # Save temporary output
    output_path = "upsampled_audio.wav"
    sf.write(output_path, upsampled, target_sr)
    
    return output_path

DESCRIPTION = """
# Audio Upsampler (2x Sample Rate)

**Easily upsample your audio files to double their original sample rate using high-quality algorithms.**

Choose between:
- **CPU (resampy):** High-quality upsampling on any computer.
- **GPU (torchaudio):** Fast upsampling if you have a compatible graphics card.

**How to Use:**
1. Upload or record an audio file.
2. Select your processing backend.
3. Click submit to receive an upsampled audio file.

The output will always be a WAV file with double the original sample rate.
"""

ARTICLE = """
## How This App Works

This app increases the sample rate of your audio files (upsampling) for improved quality or compatibility with other tools.  
- **CPU backend** uses the `resampy` library for high-quality, anti-aliased upsampling.
- **GPU backend** uses PyTorch and `torchaudio` to leverage CUDA acceleration (if available).

**Supported Input File Types:**
- WAV
- FLAC
- MP3
- OGG
- AAC
- M4A

The app automatically detects your audio's sample rate and outputs an upsampled version in standard WAV format.

**Maximum File Length:**  
There is no strict limit, but very large files may take longer to process, especially on CPU.

**No AudioSR:**  
This app does not use the AudioSR library; it performs real sample rate conversion.

**Example Usage:**
- Upload a 44.1kHz WAV file, select "GPU" if you have a compatible graphics card, and get an 88.2kHz upsampled output.
- Record a voice memo, select "CPU", and receive a higher-sample-rate version for further editing.

If you have questions about supported formats or encounter issues, please let us know!
"""

interface = gr.Interface(
    fn=upsample_audio,
    inputs=[
        gr.Audio(label="Input Audio", type="filepath"),
        gr.Radio(["CPU", "GPU"], label="Processing Backend", value="CPU")
    ],
    outputs=gr.Audio(label="Upsampled Audio"),
    title="Audio Upsampler (2x Sample Rate)",
    description=DESCRIPTION,
    article=ARTICLE,
    examples=[
        ["sample1.wav", "CPU"],
        ["sample2.flac", "GPU"]
    ]
)

if __name__ == "__main__":
    interface.launch()