Files
DiffSynth-Studio/diffsynth/utils/state_dict_converters/ace_step_tokenizer.py
2026-04-17 17:06:26 +08:00

28 lines
982 B
Python

"""
State dict converter for ACE-Step Tokenizer model.
The original checkpoint stores tokenizer and detokenizer weights at the top level:
- tokenizer.* (AceStepAudioTokenizer: audio_acoustic_proj, attention_pooler, quantizer)
- detokenizer.* (AudioTokenDetokenizer: embed_tokens, layers, proj_out)
These map directly to the AceStepTokenizer class which wraps both as
self.tokenizer and self.detokenizer submodules.
"""
def ace_step_tokenizer_converter(state_dict):
"""
Convert ACE-Step Tokenizer checkpoint keys to DiffSynth format.
The checkpoint keys `tokenizer.*` and `detokenizer.*` already match
the DiffSynth AceStepTokenizer module structure (self.tokenizer, self.detokenizer).
No key remapping needed — just extract the relevant keys.
"""
new_state_dict = {}
for key in state_dict:
if key.startswith("tokenizer.") or key.startswith("detokenizer."):
new_state_dict[key] = state_dict[key]
return new_state_dict