demo_inference.py

import torch
import soundfile as sf
from pathlib import Path
import unicodedata
from transformers import (
    SpeechT5ForTextToSpeech,
    SpeechT5Processor,
    SpeechT5HifiGan,
    PreTrainedTokenizerFast,
)

MODEL_ID = "ahnhs2k/speecht5-korean"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def decompose_jamo(text):
    result = []
    for ch in text:
        name = unicodedata.name(ch, "")
        if "HANGUL SYLLABLE" in name:
            code = ord(ch) - 0xAC00
            result.append(chr(0x1100 + (code // 588)))
            result.append(chr(0x1161 + ((code % 588) // 28)))
            jong = code % 28
            if jong > 0:
                result.append(chr(0x11A7 + jong))
        else:
            result.append(ch)
    return result


def main():
    model = SpeechT5ForTextToSpeech.from_pretrained(MODEL_ID).to(DEVICE).eval()

    # Processor는 항상 원본에서 로드
    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")

    # Custom tokenizer 로드 후 processor에 덮어쓰기
    tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_ID)
    processor.tokenizer = tokenizer

    vocoder = SpeechT5HifiGan.from_pretrained(Path(__file__).resolve().parent / "vocoder").to(DEVICE).eval()

    # speaker embedding
    spk_path = Path(__file__).resolve().parent / "speaker_embedding.pth"
    spk_emb = torch.load(spk_path).to(DEVICE)

    text = "안녕하세요. 자모 토크나이저 기반 한국어 TTS 데모입니다."
    jamo_seq = decompose_jamo(text)

    enc = tokenizer(jamo_seq, is_split_into_words=True, add_special_tokens=True, return_tensors="pt")
    enc = {k: v.to(DEVICE) for k, v in enc.items()}

    with torch.no_grad():
        gen = model.generate_speech(enc["input_ids"], speaker_embeddings=spk_emb.unsqueeze(0), vocoder=vocoder)

    sf.write("demo_inference_output.wav", gen.cpu().numpy(), 16000)
    print("Saved demo_inference_output.wav")


if __name__ == "__main__":
    main()