generate_podcast.py

#!/usr/bin/env python3
"""
Generate a NotebookLM-style podcast using Qwen3-TTS with voice cloning.

Uses the Base model with ICL (in-context learning) voice cloning for natural
prosody and speaker identity preservation.
"""

import time
import torch
import soundfile as sf
import numpy as np
from pathlib import Path

from qwen_tts import Qwen3TTSModel
from podcast_script import PODCAST_SCRIPT, REFERENCE_TEXT


# Generation parameters tuned for natural prosody
GENERATION_KWARGS = {
    "max_new_tokens": 2048,
    "do_sample": True,
    "top_k": 50,
    "top_p": 0.95,
    "temperature": 1.0,
    "repetition_penalty": 1.05,
    # Subtalker controls prosody/rhythm - higher values = more natural variation
    "subtalker_dosample": True,
    "subtalker_top_k": 80,
    "subtalker_top_p": 0.95,
    "subtalker_temperature": 1.1,
}

# Pause between speakers in seconds
PAUSE_BETWEEN_SPEAKERS = 0.4


def get_device_config():
    """Determine optimal device and dtype settings."""
    if torch.cuda.is_available():
        return "cuda:0", torch.bfloat16, "flash_attention_2"
    elif torch.backends.mps.is_available():
        return "mps", torch.float32, "eager"
    else:
        return "cpu", torch.float32, "eager"


def main():
    base_dir = Path(__file__).parent

    # Voice reference files (24kHz mono WAV recommended)
    nikhil_ref = base_dir / "nikhil_voice_clean.wav"
    anup_ref = base_dir / "anup_voice_clean.wav"

    output_dir = base_dir / "output"
    output_dir.mkdir(exist_ok=True)

    device, dtype, attn_impl = get_device_config()
    print(f"Using device: {device}")

    # Load Base model for voice cloning
    print("Loading Qwen3-TTS-12Hz-1.7B-Base model...")
    tts = Qwen3TTSModel.from_pretrained(
        "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
        device_map=device,
        dtype=dtype,
        attn_implementation=attn_impl,
    )
    print("Model loaded!")

    # Load reference audio
    print("Loading reference audio...")
    nikhil_audio, nikhil_sr = sf.read(str(nikhil_ref))
    anup_audio, anup_sr = sf.read(str(anup_ref))

    # Create voice clone prompts using ICL mode for richer voice capture
    print("Creating voice clone prompts (ICL mode)...")
    nikhil_prompt = tts.create_voice_clone_prompt(
        ref_audio=(nikhil_audio, nikhil_sr),
        ref_text=REFERENCE_TEXT,
        x_vector_only_mode=False,
    )
    anup_prompt = tts.create_voice_clone_prompt(
        ref_audio=(anup_audio, anup_sr),
        ref_text=REFERENCE_TEXT,
        x_vector_only_mode=False,
    )
    print("Voice prompts created!")

    # Generate segments
    all_audio = []
    sample_rate = None
    total = len(PODCAST_SCRIPT)

    print(f"\nGenerating {total} segments...")
    print(f"Params: temp={GENERATION_KWARGS['temperature']}, "
          f"subtalker_temp={GENERATION_KWARGS['subtalker_temperature']}")

    for i, segment in enumerate(PODCAST_SCRIPT):
        speaker = segment["speaker"]
        text = segment["text"]
        prompt = nikhil_prompt if speaker == "nikhil" else anup_prompt

        print(f"\n[{i+1}/{total}] {speaker}: {text[:50]}...")

        t0 = time.time()
        wavs, sr = tts.generate_voice_clone(
            text=text,
            language="English",
            voice_clone_prompt=prompt,
            **GENERATION_KWARGS,
        )
        t1 = time.time()

        audio = wavs[0]
        duration = len(audio) / sr
        print(f"    Generated {duration:.1f}s in {t1-t0:.2f}s")

        sample_rate = sr

        # Save individual segment
        segment_path = output_dir / f"segment_{i:02d}_{speaker}.wav"
        sf.write(str(segment_path), audio, sr)

        all_audio.append(audio)

        # Add pause between speakers
        pause = np.zeros(int(sr * PAUSE_BETWEEN_SPEAKERS))
        all_audio.append(pause)

    # Concatenate all segments
    print("\n" + "=" * 50)
    print("Concatenating...")
    final_audio = np.concatenate(all_audio)

    final_path = output_dir / "podcast_slm_2026.wav"
    sf.write(str(final_path), final_audio, sample_rate)

    duration_seconds = len(final_audio) / sample_rate
    print(f"\nPodcast generated: {final_path}")
    print(f"Duration: {duration_seconds/60:.1f} minutes")
    print("=" * 50)


if __name__ == "__main__":
    main()