| |
| """Standalone inference for Pranavz/583.""" |
| import argparse |
| import re |
| import wave |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
| from transformers import AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer, MimiModel |
|
|
| AUDIO_RE = re.compile(r"^<(\d+)_(\d+)>$") |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--repo-id", default="Pranavz/583") |
| p.add_argument("--speaker", choices=["Ira", "Aisha", "Siya", "Zoya", "Silver"], default="Ira") |
| p.add_argument("--text", required=True) |
| p.add_argument("--output", default="output.wav") |
| p.add_argument("--temperature", type=float, default=0.8) |
| p.add_argument("--top-k", type=int, default=30) |
| p.add_argument("--max-new-tokens", type=int, default=2048) |
| p.add_argument("--device", default="cuda") |
| args = p.parse_args() |
|
|
| root = Path(snapshot_download(args.repo_id)) |
| dtype = torch.bfloat16 if args.device.startswith("cuda") else torch.float32 |
| tokenizer = AutoTokenizer.from_pretrained(root) |
| model = AutoModelForCausalLM.from_pretrained( |
| root, trust_remote_code=True, torch_dtype=dtype, attn_implementation="sdpa" |
| ).eval().to(args.device) |
| mimi = MimiModel.from_pretrained(root / "codec", torch_dtype=dtype).eval().to(args.device) |
| feature = AutoFeatureExtractor.from_pretrained(root / "codec") |
|
|
| vocab = tokenizer.get_vocab() |
| mapping = {} |
| for token, token_id in vocab.items(): |
| match = AUDIO_RE.match(token) |
| if match: |
| mapping[int(token_id)] = (int(match.group(1)), int(match.group(2))) |
| allowed = torch.tensor(sorted([*mapping, int(vocab["</audio>"])]), device=args.device) |
| prompt = f'<text>{args.speaker}: {args.text}<audio>' |
| inputs = tokenizer(prompt, return_tensors="pt").to(args.device) |
| output = model.generate_audio( |
| **inputs, |
| allowed_ids=allowed, |
| max_new_tokens=args.max_new_tokens, |
| min_new_tokens=8, |
| temperature=args.temperature, |
| top_k=args.top_k, |
| do_sample=True, |
| )[0].tolist() |
|
|
| start = len(inputs.input_ids[0]) |
| try: |
| end = output.index(int(vocab["</audio>"]), start) |
| except ValueError: |
| end = len(output) |
| values, frame, expected = [], [], 0 |
| for token_id in output[start:end]: |
| item = mapping.get(int(token_id)) |
| if item is None: |
| frame, expected = [], 0 |
| continue |
| code, q = item |
| if q == expected: |
| frame.append(code) |
| expected += 1 |
| if expected == 8: |
| values.extend(frame) |
| frame, expected = [], 0 |
| elif q == 0: |
| frame, expected = [code], 1 |
| else: |
| frame, expected = [], 0 |
| if not values: |
| raise RuntimeError("No complete Mimi-Q8 frames generated") |
| codes = torch.tensor(values, device=args.device).reshape(1, -1, 8).transpose(1, 2) |
| audio = mimi.decode(codes).audio_values[0, 0].float().cpu().clamp(-1, 1) |
| pcm = (audio.numpy() * 32767).astype("<i2") |
| with wave.open(args.output, "wb") as f: |
| f.setnchannels(1); f.setsampwidth(2); f.setframerate(int(feature.sampling_rate)); f.writeframes(pcm.tobytes()) |
| print(f"Saved {args.output} ({len(audio) / feature.sampling_rate:.2f}s)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
