from pathlib import Path

import torch
import torchaudio
from torchaudio import transforms as T
from transformers import AutoModelForCausalLM, AutoTokenizer
from neucodec import NeuCodec

# model読み込み
tokenizer = AutoTokenizer.from_pretrained("tsukemono/neuTTS-JP-150m", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("tsukemono/neuTTS-JP-150m")
model.eval()
codec = NeuCodec.from_pretrained("neuphonic/neucodec")
codec = codec.eval()

# 参照音声のエンコード
waveform, sr = torchaudio.load("参照音源.mp3")
if waveform.shape[0] > 1:
    waveform = waveform.mean(dim=0, keepdim=True)
if sr != 16_000:
    waveform = T.Resample(sr, 16_000)(waveform)
waveform = waveform.unsqueeze(0)  # (B, 1, T_16k)
with torch.inference_mode():
    ref_codes = codec.encode_code(waveform).flatten().tolist()

# テキストをトークナイズしてプロンプト作成
text_ids = tokenizer(
    "ここに作成したいテキストを書いてください",
    add_special_tokens=False,
    return_attention_mask=False,
    return_token_type_ids=False,
)["input_ids"]

eos_id = int(tokenizer.eos_token_id)
input_ids = ref_codes + [eos_id] + text_ids + [eos_id]
input_ids = torch.tensor([input_ids], dtype=torch.long)

# 生成
with torch.inference_mode():
    generated = model.generate(
        input_ids=input_ids,
        repetition_penalty=1.1,
        max_new_tokens=1500,
    )

# 生成トークンから音声トークンだけ抽出
gen_ids = generated[0, input_ids.shape[1] :]
gen_ids = gen_ids[gen_ids < 65536]

# vocoderで音声作成 
with torch.inference_mode():
    audio_data = codec.decode_code(gen_ids.unsqueeze(0).unsqueeze(0)).cpu()
torchaudio.save("output.mp3", audio_data[0], 24_000, format="mp3")

サンプル音声

追記(2026/2/28)

学習元（llm-jp-3−150m）との兼ね合いで、ライセンスの変更を行いました。申し訳ありません。元のモデルとの変更点と致しましては下記になります

トークナイザーの一部トークン（外国語・記号など）の削除
トークナイザーの一部トークン（音声トークン）の追加
トークナイザーのエンコード部分に日本語のルビを振る処理の加筆
トークナイザーに合わせる形でのモデルのinput/output次元数の変更
モデルの重みの変更

Apache License, Version 2.0