#!/usr/bin/env python3
"""Render the bake-off script with VibeVoice-1.5B (single speaker, long-form).
Voice reference: a Kokoro-synthesized clip (Microsoft's preset wavs are no longer
hosted). One-shot generation — VibeVoice is built for long-form."""
import os, sys
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
from vibevoice.modular.modeling_vibevoice_inference import (
    VibeVoiceForConditionalGenerationInference,
)
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor

MODEL = "microsoft/VibeVoice-1.5B"
script = open(os.path.join(HERE, "script.txt")).read()
text = "Speaker 1: " + script.replace("\n\n", "\nSpeaker 1: ").replace("\n", " ")

print("[vibevoice] loading model...", file=sys.stderr)
processor = VibeVoiceProcessor.from_pretrained(MODEL)
model = VibeVoiceForConditionalGenerationInference.from_pretrained(
    MODEL, torch_dtype=torch.bfloat16, device_map="cuda",
    attn_implementation="sdpa")
model.eval()

inputs = processor(text=[text], voice_samples=[[os.path.join(HERE, "vv_voice_ref.wav")]],
                   padding=True, return_tensors="pt", return_attention_mask=True)
inputs = {k: (v.to("cuda") if hasattr(v, "to") else v) for k, v in inputs.items()}

print("[vibevoice] generating...", file=sys.stderr)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=None, cfg_scale=1.3,
                             tokenizer=processor.tokenizer,
                             generation_config={"do_sample": False}, verbose=False)

wav_path = os.path.join(HERE, "_vv.wav")
processor.save_audio(outputs.speech_outputs[0], output_path=wav_path)

import wave, lameenc
with wave.open(wav_path) as w:
    rate, nch, frames = w.getframerate(), w.getnchannels(), w.readframes(w.getnframes())
    dur = w.getnframes() // rate
enc = lameenc.Encoder()
enc.set_bit_rate(64); enc.set_in_sample_rate(rate); enc.set_channels(nch); enc.set_quality(5)
open(os.path.join(HERE, "voice_test_vibevoice.mp3"), "wb").write(bytes(enc.encode(frames) + enc.flush()))
os.remove(wav_path)
print(f"[vibevoice] DONE {dur}s")