Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
89 lines
4.3 KiB
Python
89 lines
4.3 KiB
Python
# Dump prompt parity fixtures for MOSS-TTSD, in the schema the moss_tts_v15
|
|
# prompt_parity harness reads.
|
|
#
|
|
# Prompt-only, so the model weights are never loaded: the processor is built for
|
|
# its tokenizer and its message rendering. The codec IS loaded, because
|
|
# AutoProcessor constructs it -- a reference with audio needs real codes and
|
|
# there is no way to ask the reference for a prompt without one.
|
|
import json, os, sys
|
|
import torch
|
|
from transformers import AutoProcessor
|
|
|
|
MODEL = "/mnt/data/models/MOSS-TTSD-v1.0"
|
|
CODEC = "/mnt/data/models/MOSS-Audio-Tokenizer"
|
|
OUT = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True, codec_path=CODEC)
|
|
pm = sys.modules[type(proc).__module__]
|
|
n_vq = int(proc.model_config.n_vq)
|
|
sr = int(proc.model_config.sampling_rate)
|
|
|
|
def dump(name, msg):
|
|
ids = proc([msg], mode="generation")["input_ids"] # (1, rows, 1+n_vq)
|
|
rec = {
|
|
"name": name,
|
|
"text": msg.text,
|
|
"instruction": msg.instruction,
|
|
"language": msg.language,
|
|
# Recorded as the message actually carried them. Writing a value the
|
|
# message did not use makes the fixture describe a prompt nobody rendered.
|
|
"quality": msg.quality,
|
|
"sound_event": msg.sound_event,
|
|
"ambient_sound": msg.ambient_sound,
|
|
"content": msg._content,
|
|
"input_ids": ids[0].tolist(),
|
|
}
|
|
# Positional, one entry per speaker, null where that speaker brought no
|
|
# audio. The layout IS the test: a fixture that only listed the present
|
|
# speakers could not tell a "[S2]: None" line from a missing one.
|
|
#
|
|
# Stored [n_vq][frames], the codec layout the C++ ReferenceAudio takes, while
|
|
# the message wants [frames][channels] -- hence the transpose. Truncated to
|
|
# n_vq first: the tokenizer emits 32 RVQ layers and this model reads the
|
|
# first 16 (processing_moss_tts.py, "Always follow model RVQ channels"), so a
|
|
# fixture carrying all 32 would describe codes the model never saw.
|
|
if isinstance(msg.reference, list):
|
|
rec["reference_codes"] = [
|
|
None if r is None else r[:, :n_vq].transpose(0, 1).contiguous().tolist()
|
|
for r in msg.reference
|
|
]
|
|
json.dump(rec, open(f"{OUT}/ref_prompt_{name}.json", "w"))
|
|
print(f"wrote ref_prompt_{name}.json rows={ids.shape[1]} channels={ids.shape[2]} (1+{n_vq})")
|
|
|
|
DIALOGUE = "[S1] So what did you make of it? [S2] Honestly, I was not expecting that ending."
|
|
|
|
# 1. Two speakers, neither cloned: the "[S<n>]: None" branch, twice.
|
|
dump("dialogue_none", pm.UserMessage(text=DIALOGUE, reference=[None, None]))
|
|
|
|
# 2. No speakers named at all: the `reference is None` branch.
|
|
dump("dialogue_no_reference", pm.UserMessage(text=DIALOGUE))
|
|
|
|
# 3. One cloned, one not -- the mixed case, which is where a shape error in
|
|
# either branch shows up as a token shift rather than a wrong count.
|
|
torch.manual_seed(0)
|
|
wav = torch.zeros(1, sr) # one second; content is irrelevant to prompt shape
|
|
wav[0, ::137] = 0.2 # some structure so the codec does not emit a constant
|
|
codes = proc.encode_audios_from_wav([wav], sampling_rate=sr)
|
|
dump("dialogue_mixed", pm.UserMessage(text=DIALOGUE, reference=[codes[0], None]))
|
|
|
|
# 4. Continuation: the mode TTSD actually clones in. The user turn names the
|
|
# speakers, and an assistant turn carries the concatenated reference audio
|
|
# that the model continues from. The assistant span uses a different pair of
|
|
# slot tokens from a user span -- gen for the frames, delay for the trailing
|
|
# n_vq - 1 -- and the turn is left unterminated so generation continues it.
|
|
cont_msgs = [
|
|
pm.UserMessage(text=DIALOGUE, reference=[codes[0], None]),
|
|
pm.AssistantMessage(audio_codes_list=[codes[0]]),
|
|
]
|
|
ids = proc([cont_msgs], mode="continuation")["input_ids"]
|
|
json.dump({
|
|
"name": "dialogue_continuation",
|
|
"text": DIALOGUE,
|
|
"instruction": None, "language": None,
|
|
"quality": None, "sound_event": None, "ambient_sound": None,
|
|
"content": cont_msgs[0]._content,
|
|
"reference_codes": [codes[0][:, :n_vq].transpose(0, 1).contiguous().tolist(), None],
|
|
"assistant_codes": codes[0][:, :n_vq].transpose(0, 1).contiguous().tolist(),
|
|
"input_ids": ids[0].tolist(),
|
|
}, open(f"{OUT}/ref_prompt_dialogue_continuation.json", "w"))
|
|
print(f"wrote ref_prompt_dialogue_continuation.json rows={ids.shape[1]}")
|