app.py

"""Universal speech enhancement with USES.

Published as https://huggingface.co/spaces/espnet/universal-se; the source
lives in espnet, at egs2/universal_se_v1/enh1/demo. The recipe that trained
the checkpoint is the directory above this one.

One of the five demo Spaces this repository maintains; the others are
egs2/owsm_ctc_v4/s2t1/demo, egs2/owsm_v4/s2t1/demo, egs2/ljspeech/tts1/demo
and egs2/voxceleb/spk1/demo.
"""

# The ZeroGPU package patches torch as it is imported, so it has to come
# first - before torch, and before anything that imports torch.
try:  # only Hugging Face's runners have it
    import spaces  # isort: skip
except ImportError:  # running elsewhere: the decorator does nothing

    class spaces:  # noqa: N801 - stands in for the module
        @staticmethod
        def GPU(func=None, **kwargs):
            return func if func is not None else (lambda f: f)


import os  # noqa: E402
import re  # noqa: E402

import gradio as gr  # noqa: E402
import librosa  # noqa: E402
import torch  # noqa: E402

from espnet2.bin.enh_inference import SeparateSpeech  # noqa: E402

# USES runs a recurrent block over the whole recording rather than over a
# fixed window, so the cost is linear in the input and half a minute is what
# comfortably fits the GPU slice asked for below. Longer audio is trimmed
# rather than refused.
MAX_SECS = 30
GPU_SECONDS = 120
MODEL_TAG = os.environ.get(
    "ENH_MODEL_TAG", "espnet/Wangyou_Zhang_universal_train_enh_uses_refch0_2mem_raw"
)
# ZeroGPU attaches the GPU only while a @spaces.GPU function runs, so
# torch.cuda.is_available() is False here and asking it would pin the models to
# the CPU on the very hardware bought to run them. SPACES_ZERO_GPU is the
# runtime's own marker; `spaces` being importable is not, since anyone can
# install it.
ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))
if os.environ.get("DEVICE"):
    DEVICE = os.environ["DEVICE"]
elif ZERO_GPU or torch.cuda.is_available():
    DEVICE = "cuda"
else:
    DEVICE = "cpu"

KEEP = "Same as the file"
EXAMPLE_WAV = (
    "https://github.com/espnet/espnet/raw/master/test_utils/ctc_align_test.wav"
)


def _categories(model):
    """The training conditions a checkpoint names, in a stable order.

    ESPnet stores them as {index: name} so that a category id in the data
    means the same thing at inference; a checkpoint trained without them has
    none, and every menu below is then empty.
    """
    categories = getattr(model, "categories", None) or {}
    if isinstance(categories, dict):
        return [categories[key] for key in sorted(categories)]
    return list(categories)


def _rates(categories):
    """The sample rates a category list names, highest first.

    A category reads like "2ch_16k" or "1ch_16000Hz": two channels at
    16 kHz, one channel at 16 kHz. Recipes spell the rate both ways - USES
    one way, the URGENT 2025 recipe the other - and a checkpoint reached
    through ENH_MODEL_TAG comes from whichever recipe trained it, so both
    spellings fill the menu.
    """
    found = set()
    for category in categories:
        match = re.search(r"(\d+)k(?:Hz)?\b", category)
        if match:
            found.add(int(match.group(1)) * 1000)
            continue
        match = re.search(r"(\d+)\s*Hz", category)
        if match:
            found.add(int(match.group(1)))
    return sorted(found, reverse=True)


def _channels(categories):
    """The largest channel count a category list names, at least one."""
    found = {1}
    for category in categories:
        match = re.match(r"(\d+)ch", category)
        if match:
            found.add(int(match.group(1)))
    return max(found)


TITLE = "Universal speech enhancement"
DESCRIPTION = """# Universal speech enhancement

USES is a single speech enhancement model for every condition the task
usually splits into separate ones: noise, reverberation, one microphone or
several, and any sample rate. It was trained on the
[`egs2/universal_se_v1`](https://github.com/espnet/espnet/tree/master/egs2/universal_se_v1)
data with [ESPnet](https://github.com/espnet/espnet), and it is the model
`espnet enhance` uses by default.

Upload something noisy. Multi-channel files keep their channels on the way
in, because a model that can use them should be given them; what comes back
is one channel.
"""
ARTICLE = """Model:
[`espnet/Wangyou_Zhang_universal_train_enh_uses_refch0_2mem_raw`](https://huggingface.co/espnet/Wangyou_Zhang_universal_train_enh_uses_refch0_2mem_raw).
Source of this Space:
[`egs2/universal_se_v1/enh1/demo`](https://github.com/espnet/espnet/tree/master/egs2/universal_se_v1/enh1/demo).

```bibtex
@inproceedings{zhang2023uses,
  title={Toward Universal Speech Enhancement for Diverse Input Conditions},
  author={Zhang, Wangyou and Saijo, Kohei and Wang, Zhong-Qiu and
          Watanabe, Shinji and Qian, Yanmin},
  booktitle={Proc. ASRU},
  year={2023}
}
```"""


enh = SeparateSpeech.from_pretrained(MODEL_TAG, device=DEVICE)

# The menu comes from the checkpoint. USES is sampling-frequency independent,
# so any rate runs; the rates it was trained on are the ones worth offering,
# and the recipe records them in the category names. The same names say how
# many channels it has ever seen, which is what the warning below compares
# against.
CATEGORIES = _categories(enh.enh_model)
RATES = _rates(CATEGORIES)
MAX_CHANNELS = _channels(CATEGORIES)
RATE_LABELS = [f"{rate // 1000} kHz" for rate in RATES]
RATE_OF_LABEL = dict(zip(RATE_LABELS, RATES))
# One output per speaker the checkpoint separates: this one enhances, so it
# has a single output, and a separation checkpoint reached through
# ENH_MODEL_TAG grows the column instead of dropping speakers on the floor.
OUTPUT_LABELS = (
    ["Enhanced"]
    if enh.num_spk == 1
    else [f"Speaker {i}" for i in range(1, enh.num_spk + 1)]
)


def _read(path):
    """The recording as (samples, channels) or (samples,), at its own rate."""
    if path is None:
        raise gr.Error("Record or upload some audio first.")
    # The header says how long the file is, so only the part that is enhanced
    # is decoded: an hour of audio uploaded to a 30 s demo is not read in full
    # first, and the warning is still about the file rather than about what
    # came back.
    if librosa.get_duration(path=path) > MAX_SECS:
        gr.Warning(
            f"Only the first {MAX_SECS} s were enhanced. "
            "Run the app yourself for longer audio - the model has no such "
            "limit."
        )
    # sr=None keeps the file's rate: the model is free to work at it, and
    # resampling is the user's choice below rather than a silent one here.
    speech, rate = librosa.load(path, sr=None, mono=False, duration=MAX_SECS)
    if speech.ndim > 1:  # librosa hands back (channels, samples)
        speech = speech.T
    channels = speech.shape[1] if speech.ndim > 1 else 1
    if channels > MAX_CHANNELS:
        gr.Warning(
            f"This file has {channels} channels and the model was trained on "
            f"at most {MAX_CHANNELS}; it will run, but off what it knows."
        )
    return speech, rate


@spaces.GPU(duration=GPU_SECONDS)
def predict(audio_path, rate_label):
    speech, rate = _read(audio_path)
    target = RATE_OF_LABEL.get(rate_label, rate)
    if target != rate:
        # axis=0 is the sample axis for both shapes _read returns
        speech = librosa.resample(speech, orig_sr=rate, target_sr=target, axis=0)
        rate = target
    waves = enh(speech[None, ...], fs=rate)
    return [(rate, wave[0]) for wave in waves]


EXAMPLES = [[EXAMPLE_WAV, KEEP]]


with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            audio = gr.Audio(
                sources=["microphone", "upload"], type="filepath", label="Noisy audio"
            )
            rate = gr.Dropdown(
                [KEEP] + RATE_LABELS,
                value=KEEP,
                label="Process at",
                info="The rates this model was trained on",
            )
            button = gr.Button("Enhance", variant="primary")
        with gr.Column():
            outputs = [gr.Audio(label=label) for label in OUTPUT_LABELS]
    button.click(predict, [audio, rate], outputs)
    gr.Examples(
        EXAMPLES,
        inputs=[audio, rate],
        outputs=outputs,
        fn=predict,
        cache_examples=False,
    )
    gr.Markdown(ARTICLE)


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