processing_koko_tts.py

import os
from transformers import AutoTokenizer

class KokoTTSProcessor:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, **kwargs)
        return cls(tokenizer)

    def save_pretrained(self, save_directory, **kwargs):
        os.makedirs(save_directory, exist_ok=True)
        self.tokenizer.save_pretrained(save_directory, **kwargs)

    def __call__(self, text=None, return_tensors="pt", **kwargs):
        if text is not None:
            return self.tokenizer(text, return_tensors=return_tensors, **kwargs)
        return {}