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
125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import pathlib
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
REFERENCE_ROOT = (
|
|
pathlib.Path(__file__).resolve().parents[3] / "_reference" / "GLM-TTS"
|
|
)
|
|
sys.path.insert(0, str(REFERENCE_ROOT))
|
|
|
|
from flow.flow import Flow
|
|
|
|
|
|
REFERENCE_TOKENS = [
|
|
29252, 4906, 833, 15564, 12971, 13808, 8457, 20137,
|
|
7313, 30766, 5123, 2098, 891, 2464, 16935, 26908,
|
|
20338, 286, 32542, 10092, 13792, 18489, 16200, 9895,
|
|
3221, 29590, 5542, 29513, 25025, 9777, 29029, 21176,
|
|
9781, 15068, 30891, 5324, 28699, 28591, 20197, 4599,
|
|
11932, 4759, 19157, 10647, 4938, 8754, 30267, 31575,
|
|
4731, 6326, 30991, 5040, 17687, 17687, 19635, 10780,
|
|
21045, 5387, 11503, 13228, 13228, 6164, 12240,
|
|
]
|
|
GENERATED_TOKENS = [
|
|
25568, 26211, 20216, 29473, 21992, 30971, 14699,
|
|
25619, 20563, 20602, 22277, 9540, 25669, 21373,
|
|
22055, 13228, 6387, 16809, 13228, 4503, 15102, 32001,
|
|
31164, 26617, 6852, 32401, 27451, 5997, 5590, 18014,
|
|
780, 26030, 17460, 15143, 18324, 5413, 23684, 28226,
|
|
16826, 31440, 19490, 15302, 3693, 17870, 515, 32307,
|
|
12991, 23343, 10209, 1118, 10442, 9769, 6294, 12860,
|
|
10055, 3886, 7692,
|
|
]
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("model_path", type=pathlib.Path)
|
|
parser.add_argument("--steps", type=int, default=10)
|
|
args = parser.parse_args()
|
|
|
|
model = Flow(
|
|
spkr_emb_adaLN=True,
|
|
speech_token_cfg=False,
|
|
remove_spkr_concat_condition=True,
|
|
mel_dim=80,
|
|
mel_framerate=50,
|
|
input_frame_rate=25,
|
|
)
|
|
checkpoint = torch.load(
|
|
args.model_path / "flow" / "flow.pt",
|
|
map_location="cpu",
|
|
weights_only=True,
|
|
)
|
|
model.load_state_dict(
|
|
checkpoint["model"] if "model" in checkpoint else checkpoint
|
|
)
|
|
model.eval()
|
|
|
|
prompt_frames = 126
|
|
prompt = torch.tensor(
|
|
[
|
|
math.sin((index + 1) * 0.011) * 0.5
|
|
for index in range(prompt_frames * 80)
|
|
],
|
|
dtype=torch.float32,
|
|
).reshape(1, prompt_frames, 80)
|
|
speaker = torch.tensor(
|
|
[[
|
|
math.sin((index + 1) * 0.03125)
|
|
for index in range(192)
|
|
]],
|
|
dtype=torch.float32,
|
|
)
|
|
full_frames = int(
|
|
(len(REFERENCE_TOKENS) + len(GENERATED_TOKENS)) / 25 * 50
|
|
)
|
|
noise = torch.tensor(
|
|
[
|
|
math.sin((index + 1) * 0.017)
|
|
for index in range(full_frames * 80)
|
|
],
|
|
dtype=torch.float32,
|
|
).reshape(1, full_frames, 80)
|
|
|
|
original_randn_like = torch.randn_like
|
|
|
|
def fixed_randn_like(value, *args, **kwargs):
|
|
del args, kwargs
|
|
if tuple(value.shape) == tuple(noise.shape):
|
|
return noise.clone()
|
|
return original_randn_like(value)
|
|
|
|
with torch.inference_mode(), patch(
|
|
"torch.randn_like", side_effect=fixed_randn_like
|
|
):
|
|
output, _ = model.inference_with_cache(
|
|
token=torch.tensor([GENERATED_TOKENS], dtype=torch.long),
|
|
prompt_token=torch.tensor(
|
|
[REFERENCE_TOKENS], dtype=torch.long
|
|
),
|
|
prompt_feat=prompt,
|
|
embedding=speaker,
|
|
n_timesteps=args.steps,
|
|
)
|
|
frame_major = output.permute(0, 2, 1).contiguous().view(-1)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"frames": int(output.shape[-1]),
|
|
"values": frame_major.tolist(),
|
|
},
|
|
separators=(",", ":"),
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|