| """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() |
|
|
