File size: 8,420 Bytes
cd9b2d8 c1cc83c cd9b2d8 c1cc83c cd9b2d8 c1cc83c cd9b2d8 c1cc83c cd9b2d8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/usr/bin/env python3
"""Inference for LAION Humaneness Ears Base and Medium audio encoders."""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from safetensors.torch import load_file
from scipy.signal import resample_poly
from transformers import WhisperFeatureExtractor
from model import LayeredMultiTaskWhisper
def read_audio(path):
wave, rate = sf.read(path, dtype='float32', always_2d=True)
wave = wave.mean(axis=1)
duration = len(wave) / rate
if not .1 <= duration <= 30.:
raise ValueError(f'Audio must last 0.1–30 seconds; received {duration:.3f}s')
if not np.isfinite(wave).all():
raise ValueError('Audio contains nonfinite values')
if rate != 16000:
divisor = math.gcd(rate, 16000)
wave = resample_poly(wave, 16000 // divisor, rate // divisor).astype(np.float32)
return wave[:480000], duration
class Predictor:
"""Load once and reuse. Local safe weights; no teachers or ASR decoder."""
def __init__(self, repo_dir, model_size='medium', device='cpu', threads=4):
if model_size not in ('base', 'medium', 'small') or threads < 1:
raise ValueError('Use base/medium and at least one CPU thread')
# Accept the historical architecture label for existing callers.
model_size = 'medium' if model_size == 'small' else model_size
torch.set_num_threads(threads)
self.repo = Path(repo_dir)
self.model_size, self.device = model_size, device
self.norm = json.loads((self.repo / 'training_normalization.json').read_text())
self.display = json.loads((self.repo / 'display_normalization.json').read_text())
self.class_names = json.loads((self.repo / 'classes.json').read_text())['names']
self.names = self.norm['score_names']
if self.names != [r['name'] for r in self.display['targets']]:
raise ValueError('Display and training target orders differ')
directory = self.repo / model_size
self.model = LayeredMultiTaskWhisper(directory, len(self.class_names), initialize_pretrained=False)
self.model.load_state_dict(load_file(directory / 'model.safetensors', device='cpu'), strict=True)
self.model.to(device).eval()
self.extractor = WhisperFeatureExtractor.from_pretrained(directory, local_files_only=True)
def predict(self, audio, *, frame_threshold=.5, event_threshold=.5, max_events=32, include_frames=False):
if not 0 < frame_threshold < 1 or not 0 < event_threshold < 1 or max_events < 1:
raise ValueError('Thresholds must be between zero and one, max_events positive')
wave, duration = read_audio(audio)
features = self.extractor([wave], sampling_rate=16000, padding='longest', truncation=True,
max_length=480000, return_attention_mask=True, return_tensors='np')
mel = features['input_features'].astype(np.float32)
mask = features['attention_mask'].astype(np.int64)
if mel.shape[-1] % 2:
mel = np.pad(mel, ((0, 0), (0, 0), (0, 1)))
mask = np.pad(mask, ((0, 0), (0, 1)))
from contextlib import nullcontext
autocast = torch.autocast('cuda', dtype=torch.bfloat16) if self.device.startswith('cuda') else nullcontext()
with torch.inference_mode(), autocast:
out = self.model(torch.from_numpy(mel).to(self.device), torch.from_numpy(mask).to(self.device),
predict_events=True, event_threshold=event_threshold, max_pred_events=max_events)
standardized = out['scores'][0].float().cpu().numpy()
raw = standardized * np.asarray(self.norm['score_std']) + np.asarray(self.norm['score_mean'])
z = {name: ((float(value) - row['median']) / row['std']
if row['median'] is not None and row['std'] not in (None, 0) else None)
for name, value, row in zip(self.names, raw, self.display['targets'])}
count = min(int((mask.sum() + 1) // 2), out['frame'].shape[1])
frame = out['frame'][0, :count].float().sigmoid().cpu().numpy()
edges = np.diff(np.pad((frame >= frame_threshold).astype(np.int8), (1, 1)))
regions = [{'start_s': round(float(a) * .02, 6), 'end_s': round(min(float(b) * .02, duration), 6)}
for a, b in zip(np.flatnonzero(edges == 1), np.flatnonzero(edges == -1))
if float(a) * .02 < duration]
probabilities = out['event_class'][0].float().softmax(-1).cpu().numpy()
starts = out['predicted_event_starts'][0].cpu().numpy()
ends = out['predicted_event_ends'][0].cpu().numpy()
valid = out['predicted_event_valid'][0].cpu().numpy()
onset = out['onset'][0].float().sigmoid().cpu().numpy()
proposals = []
for j in np.flatnonzero(valid):
start, end = float(starts[j]) * .02, min(float(ends[j]) * .02, duration)
if start >= end:
continue
top = np.argsort(probabilities[j])[-3:][::-1]
proposals.append({'start_s': round(start, 6), 'end_s': round(end, 6),
'onset_probability': float(onset[int(starts[j])]),
'top3_classes': [{'label': self.class_names[int(k)],
'softmax_probability': float(probabilities[j, k])} for k in top]})
cps_z = float(out['cps'][0].float().cpu())
result = {
'model': self.model_size, 'checkpoint': 'best Gemini Flash 3.8 fine-tune; epoch 2',
'duration_s': duration,
'scores_raw': dict(zip(self.names, map(float, raw))),
'scores_training_z': dict(zip(self.names, map(float, standardized))),
'scores_z': z,
'characters_per_second': cps_z * self.norm['cps_std'] + self.norm['cps_mean'],
'characters_per_second_training_z': cps_z,
'estimated_burst_count': float(np.expm1(np.clip(raw[130], 0, 8))),
'orange_timbre_128d': out['timbre'][0].float().cpu().tolist(),
'orange_identity_250d': out['identity'][0].float().cpu().tolist(),
'binary_burst_regions': regions, 'burst_event_proposals': proposals,
'decoding': {'frame_threshold': frame_threshold, 'event_threshold': event_threshold, 'max_events': max_events},
'notes': 'Encoder only: no transcript or captions. Raw regression scores are not clipped to ordinal scales. '
'Display z-scores use the original train-only reference median/std; they are not probabilities. '
'Class probabilities are uncalibrated. Blend is only meaningful when vocal bursts are present.',
}
if include_frames:
result['burst_probability_per_20ms_frame'] = frame.tolist()
return result
def predict(audio, repo_dir, *, model_size='medium', device='cpu', threads=4, **options):
"""Convenience call; reuse Predictor for many clips to avoid reloading weights."""
return Predictor(repo_dir, model_size, device, threads).predict(audio, **options)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('audio', type=Path)
parser.add_argument('--repo-dir', type=Path, default=Path(__file__).resolve().parent)
parser.add_argument('--model', choices=('base', 'medium', 'small'), default='medium',
help='Release variant (small is a legacy alias for medium)')
parser.add_argument('--device', default='cpu')
parser.add_argument('--threads', type=int, default=4)
parser.add_argument('--frame-threshold', type=float, default=.5)
parser.add_argument('--event-threshold', type=float, default=.5)
parser.add_argument('--max-events', type=int, default=32)
parser.add_argument('--include-frame-probabilities', action='store_true')
args = parser.parse_args()
result = predict(args.audio, args.repo_dir, model_size=args.model, device=args.device, threads=args.threads,
frame_threshold=args.frame_threshold, event_threshold=args.event_threshold,
max_events=args.max_events, include_frames=args.include_frame_probabilities)
print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False))
if __name__ == '__main__':
main()
|