app.py

import matplotlib
matplotlib.use("Agg")

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import librosa
import gradio as gr
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from scipy.ndimage import uniform_filter1d
from transformers import Wav2Vec2Processor, pipeline
from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2Model, Wav2Vec2PreTrainedModel
from silero_vad import load_silero_vad, get_speech_timestamps

try:
    import spaces                      # available on ZeroGPU Spaces
    gpu = spaces.GPU(duration=60)
except Exception:                      # running locally / on CPU Space
    def gpu(fn):
        return fn

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SR = 16000


# ---------------- age + gender model (custom head from the model card) ----------------
class ModelHead(nn.Module):
    def __init__(self, config, n):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.dropout = nn.Dropout(config.final_dropout)
        self.out_proj = nn.Linear(config.hidden_size, n)

    def forward(self, x):
        x = self.dropout(x)
        x = torch.tanh(self.dense(x))
        x = self.dropout(x)
        return self.out_proj(x)


class AgeGenderModel(Wav2Vec2PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.wav2vec2 = Wav2Vec2Model(config)
        self.age = ModelHead(config, 1)
        self.gender = ModelHead(config, 3)
        self.post_init()

    def forward(self, x):
        h = self.wav2vec2(x)[0].mean(1)
        return self.age(h), torch.softmax(self.gender(h), dim=1)


AG_ID = "audeering/wav2vec2-large-robust-24-ft-age-gender"
ag_proc = Wav2Vec2Processor.from_pretrained(AG_ID)
ag_model = AgeGenderModel.from_pretrained(AG_ID).to(DEVICE).eval()

dev_idx = 0 if DEVICE == "cuda" else -1
emo_pipe = pipeline("audio-classification", model="superb/wav2vec2-base-superb-er", device=dev_idx)
lid_pipe = pipeline("audio-classification", model="facebook/mms-lid-126", device=dev_idx)

vad_model = load_silero_vad()

EMO = {"neu": "neutral", "hap": "happy", "ang": "angry", "sad": "sad"}


def vad_segments(y):
    ts = get_speech_timestamps(torch.from_numpy(np.asarray(y, np.float32)), vad_model,
                               sampling_rate=SR, return_seconds=True)
    return [(float(t["start"]), float(t["end"])) for t in ts]


def age_bucket(a):
    for hi, name in [(12, "child"), (19, "teen"), (30, "young adult"), (50, "adult"), (65, "middle-aged")]:
        if a < hi:
            return name
    return "senior"


# ---------------- GPU part: only model inference ----------------
@gpu
@torch.no_grad()
def gpu_infer(speech):
    x = ag_proc(speech, sampling_rate=SR)["input_values"][0].reshape(1, -1)
    age_out, g = ag_model(torch.from_numpy(x).float().to(DEVICE))
    age = float(age_out.item()) * 100
    gp = g[0].cpu().numpy()
    emo = emo_pipe({"raw": speech, "sampling_rate": SR}, top_k=1)[0]
    lid = lid_pipe({"raw": speech, "sampling_rate": SR}, top_k=1)[0]
    return age, gp, emo, lid


# ---------------- CPU part: audio processing ----------------
def analyze_voice(path):
    y, _ = librosa.load(path, sr=SR, mono=True, duration=30)
    if len(y) < SR // 2:
        raise gr.Error("Audio is too short. Please upload at least 1-2 seconds of speech.")
    y = y / (np.abs(y).max() + 1e-6)

    try:
        segs = vad_segments(y)
    except Exception:
        segs = [(a / SR, b / SR) for a, b in librosa.effects.split(y, top_db=30)]
    speech = np.concatenate([y[int(a * SR):int(b * SR)] for a, b in segs]) if segs else y
    speech_dur = sum(b - a for a, b in segs) or len(y) / SR
    if len(speech) < SR:
        speech = np.pad(speech, (0, SR - len(speech)))

    age, gp, emo, lid = gpu_infer(speech)
    gender = ["female", "male", "child"][int(gp.argmax())]

    f0, _, _ = librosa.pyin(y, fmin=65, fmax=500, sr=SR)
    times = librosa.times_like(f0, sr=SR)
    f0v = f0[~np.isnan(f0)]
    if len(f0v):
        pitch = dict(median=np.median(f0v), p10=np.percentile(f0v, 10), p90=np.percentile(f0v, 90))
    else:
        pitch = dict(median=np.nan, p10=np.nan, p90=np.nan)

    rms = uniform_filter1d(librosa.feature.rms(y=speech, frame_length=400, hop_length=160)[0], 5)
    peaks, _ = find_peaks(rms, height=0.1 * rms.max(), prominence=0.05 * rms.max(), distance=8)
    rate = len(peaks) / speech_dur

    return dict(age=age, bucket=age_bucket(age), gender=f"{gender} ({gp.max():.0%})",
                emotion=f"{EMO.get(emo['label'], emo['label'])} ({emo['score']:.0%})",
                language=f"{lid['label']} ({lid['score']:.0%})",
                pitch=pitch, rate=rate, speech_dur=speech_dur, f0=f0, times=times)


def compare(pa, pb):
    if pa is None or pb is None:
        raise gr.Error("Please provide both Voice A and Voice B.")
    A, B = analyze_voice(pa), analyze_voice(pb)
    rows = [
        ("Gender (estimated)", A["gender"], B["gender"]),
        ("Age (estimated)", f"{A['age']:.0f} yrs ({A['bucket']})", f"{B['age']:.0f} yrs ({B['bucket']})"),
        ("Emotion", A["emotion"], B["emotion"]),
        ("Language (ISO 639-3)", A["language"], B["language"]),
        ("Median pitch (Hz)", f"{A['pitch']['median']:.0f}", f"{B['pitch']['median']:.0f}"),
        ("Pitch range p10-p90 (Hz)", f"{A['pitch']['p10']:.0f}-{A['pitch']['p90']:.0f}",
                                     f"{B['pitch']['p10']:.0f}-{B['pitch']['p90']:.0f}"),
        ("Speaking rate (syllables/s)", f"{A['rate']:.1f}", f"{B['rate']:.1f}"),
        ("Speech duration (s)", f"{A['speech_dur']:.1f}", f"{B['speech_dur']:.1f}"),
    ]
    df = pd.DataFrame(rows, columns=["Attribute", "Voice A", "Voice B"])

    fig, ax = plt.subplots(figsize=(10, 3.5))
    ax.plot(A["times"], A["f0"], label="Voice A")
    ax.plot(B["times"], B["f0"], label="Voice B")
    ax.set_xlabel("time (s)")
    ax.set_ylabel("F0 (Hz)")
    ax.legend()
    fig.tight_layout()
    return df, fig


DESCRIPTION = (
    "Upload or record two voices (3-30 s of clear speech each) to compare estimated gender, age, "
    "emotion, language, pitch and speaking rate. Results are model estimates, not facts about a person. "
    "Only use recordings of people who have agreed to be analyzed."
)

demo = gr.Interface(
    fn=compare,
    inputs=[gr.Audio(type="filepath", label="Voice A"), gr.Audio(type="filepath", label="Voice B")],
    outputs=[gr.Dataframe(label="Attributes"), gr.Plot(label="Pitch contour")],
    title="Voice attribute comparison",
    description=DESCRIPTION,
    flagging_mode="never",
)

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