| """Runnable inference for diffusion-gemma-asr-small. |
| |
| python inference.py audio.wav |
| """ |
| import sys, numpy as np, torch, soundfile as sf |
| from huggingface_hub import snapshot_download |
| from transformers import AutoTokenizer, WhisperFeatureExtractor |
| from model import AudioDiffusionGemma, AudioDiffusionConfig |
| from audio import audio_out_len, real_encoder_frames, WHISPER_ENC_FRAMES |
|
|
| MODEL_ID = "google/diffusiongemma-26B-A4B-it" |
| WHISPER_ID = "openai/whisper-small" |
| D_MODEL, VOCAB = 2816, 262144 |
| BOA, AUD, EOA = 256000, 258881, 258883 |
| EOS_IDS = {1, 106, 50} |
| SR, SUBSAMPLE = 16000, 8 |
| INSTRUCTION = "Transcribe this audio to text." |
| |
| CANVAS_LEN, TARGET_SEG, MAX_SEG = 192, 13.0, 18.0 |
|
|
|
|
| def load(ckpt="diffusion_asr_small.pt", device="cuda"): |
| model_dir = snapshot_download(MODEL_ID) |
| cfg = AudioDiffusionConfig(model_dir=model_dir, whisper_id=WHISPER_ID, whisper_dim=768, |
| d_model=D_MODEL, vocab_size=VOCAB, boa_token_id=BOA, |
| audio_token_id=AUD, eoa_token_id=EOA, subsample_factor=SUBSAMPLE) |
| model = AudioDiffusionGemma.from_pretrained(cfg, dtype=torch.bfloat16, device=device) |
| st = torch.load(ckpt, map_location=device) |
| if "lora" in st: |
| from peft import set_peft_model_state_dict |
| model.apply_lora(); set_peft_model_state_dict(model.base, st["lora"]) |
| model.projector.load_state_dict(st["projector"]); model.projector.to(device, torch.float32) |
| model.eval() |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| fe = WhisperFeatureExtractor.from_pretrained(WHISPER_ID) |
| return model, tok, fe |
|
|
|
|
| def _segment(wav): |
| target, search, fw = int(TARGET_SEG * SR), int(3.0 * SR), int(0.08 * SR) |
| segs, i = [], 0 |
| while i < len(wav): |
| if len(wav) - i <= int(MAX_SEG * SR): |
| segs.append(wav[i:]); break |
| lo, hi = i + target - search, min(len(wav), i + target + search) |
| region = wav[lo:hi]; nf = max(1, (len(region) - fw) // fw) |
| energy = np.array([float(np.dot(region[k*fw:k*fw+fw], region[k*fw:k*fw+fw])) for k in range(nf)]) |
| cut = lo + int(np.argmin(energy)) * fw + fw // 2 |
| segs.append(wav[i:cut]); i = cut |
| return segs |
|
|
|
|
| def _collapse(text): |
| out = [] |
| for w in text.split(): |
| if not out or out[-1] != w: |
| out.append(w) |
| return " ".join(out) |
|
|
|
|
| @torch.no_grad() |
| def transcribe(wav, model, tok, fe, max_steps=16, device="cuda"): |
| wav = np.asarray(wav, dtype=np.float32) |
| instr = tok(INSTRUCTION, add_special_tokens=False)["input_ids"] |
| prefix = [2] + instr + [BOA] |
| Ta = audio_out_len(WHISPER_ENC_FRAMES, SUBSAMPLE) |
| base = prefix + [AUD] * Ta + [EOA] |
| texts = [] |
| segs = _segment(wav) if len(wav) > MAX_SEG * SR else [wav] |
| for s in range(0, len(segs), 8): |
| sub = segs[s:s + 8] |
| mel = torch.stack([fe(c, sampling_rate=SR, return_tensors="pt").input_features[0] for c in sub]).to(device) |
| B = len(sub) |
| prompt_ids = torch.tensor([base] * B, dtype=torch.long, device=device) |
| prompt_mask = torch.zeros(B, len(prefix) + Ta + 1, dtype=torch.long, device=device) |
| for i, c in enumerate(sub): |
| n_real = min(Ta, audio_out_len(real_encoder_frames(len(c)), SUBSAMPLE)) |
| prompt_mask[i, :len(prefix)] = 1 |
| prompt_mask[i, len(prefix):len(prefix) + max(1, n_real)] = 1 |
| prompt_mask[i, len(prefix) + Ta] = 1 |
| canvas, _ = model.generate(prompt_ids, prompt_mask, mel, canvas_len=CANVAS_LEN, max_steps=max_steps) |
| for i in range(B): |
| ids = [] |
| for t in canvas[i].tolist(): |
| if t in EOS_IDS: |
| break |
| ids.append(t) |
| texts.append(_collapse(tok.decode(ids, skip_special_tokens=True))) |
| return " ".join(t for t in texts if t).strip() |
|
|
|
|
| if __name__ == "__main__": |
| path = sys.argv[1] if len(sys.argv) > 1 else "audio.wav" |
| wav, sr = sf.read(path) |
| if wav.ndim > 1: |
| wav = wav.mean(axis=1) |
| if sr != SR: |
| import librosa |
| wav = librosa.resample(wav.astype(np.float32), orig_sr=sr, target_sr=SR) |
| model, tok, fe = load() |
| print(transcribe(wav, model, tok, fe)) |
|
|
