File size: 3,063 Bytes
67faadf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Make a fixed-shape Piper encoder: 64 phoneme IDs in, 160 latent frames out.
Changes to data.encoder.onnx:
- input [1, N_IDS]; padding is masked by input_lengths
- total frame count (ReduceMax of predicted lengths) replaced by constant N_FRAMES;
speech longer than that is cut off, shorter is masked
- both RandomNormalLike nodes become inputs (noise_dp, noise_z) so windows can be
re-run with the same noise
- extra outputs: durations [1,1,N_IDS] (frames per ID) and frames [1] (real length)
"""
import sys
from pathlib import Path
import onnx
from onnx import TensorProto, helper
ROOT = Path(__file__).resolve().parent
N_IDS = int(sys.argv[1]) if len(sys.argv) > 1 else 64
N_FRAMES = int(sys.argv[2]) if len(sys.argv) > 2 else 160
def main() -> None:
model = onnx.load(ROOT / "data.encoder.onnx")
g = model.graph
g.ClearField("value_info")
nodes = []
for n in g.node:
if n.op_type == "Cast" and list(n.input) == ["/ReduceMax_output_0"]:
to = next(a.i for a in n.attribute if a.name == "to")
dtype = {TensorProto.INT64: TensorProto.INT64, TensorProto.INT32: TensorProto.INT32,
TensorProto.FLOAT: TensorProto.FLOAT}[to]
nodes.append(helper.make_node("Constant", [], list(n.output), name="FixedFrames",
value=helper.make_tensor("fixed_frames", dtype, [], [N_FRAMES])))
elif n.op_type == "RandomNormalLike" and n.name == "/dp/RandomNormalLike":
nodes.append(helper.make_node("Identity", ["noise_dp"], list(n.output), name="NoiseDp"))
elif n.op_type == "RandomNormalLike" and n.name == "/RandomNormalLike":
nodes.append(helper.make_node("Identity", ["noise_z"], list(n.output), name="NoiseZ"))
else:
nodes.append(n)
assert sum(x.name in ("FixedFrames", "NoiseDp", "NoiseZ") for x in nodes) == 3
nodes.append(helper.make_node("Identity", ["/Ceil_output_0"], ["durations"], name="DurationsOut"))
nodes.append(helper.make_node("Identity", ["/Clip_output_0"], ["frames"], name="FramesOut"))
g.ClearField("node")
g.node.extend(nodes)
g.input[0].CopyFrom(helper.make_tensor_value_info("input", TensorProto.INT64, [1, N_IDS]))
g.input[1].CopyFrom(helper.make_tensor_value_info("input_lengths", TensorProto.INT64, [1]))
g.input.append(helper.make_tensor_value_info("noise_dp", TensorProto.FLOAT, [1, 2, N_IDS]))
g.input.append(helper.make_tensor_value_info("noise_z", TensorProto.FLOAT, [1, 192, N_FRAMES]))
g.output[0].CopyFrom(helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 192, N_FRAMES]))
g.output.append(helper.make_tensor_value_info("durations", TensorProto.FLOAT, [1, 1, N_IDS]))
g.output.append(helper.make_tensor_value_info("frames", TensorProto.FLOAT, [1]))
model = onnx.shape_inference.infer_shapes(model)
onnx.checker.check_model(model)
out = ROOT / f"data.encoder.fixed{N_IDS}x{N_FRAMES}.onnx"
onnx.save(model, out)
print("saved", out.name)
if __name__ == "__main__":
main()
|