| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| from onnx import TensorProto, helper |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| FRAMES = 64 |
|
|
|
|
| def main() -> None: |
| model = onnx.load(ROOT / "data.decoder.onnx") |
| model.graph.ClearField("value_info") |
| model.graph.input[0].CopyFrom( |
| helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 192, FRAMES]) |
| ) |
| model.graph.output[0].CopyFrom( |
| helper.make_tensor_value_info(model.graph.output[0].name, TensorProto.FLOAT, [1, 1, FRAMES * 256]) |
| ) |
|
|
| scale = 1 |
| nodes = [] |
| mask_scales = set() |
| for node in model.graph.node: |
| if node.op_type == "ConvTranspose": |
| stride = next(attr.ints[0] for attr in node.attribute if attr.name == "strides") |
| scale *= stride |
| if node.op_type in {"Conv", "ConvTranspose"}: |
| original = node.output[0] |
| unmasked = f"{original}_unmasked" |
| node.output[0] = unmasked |
| nodes.append(node) |
| mask_name = f"valid_mask_{scale}" |
| nodes.append(helper.make_node("Mul", [unmasked, mask_name], [original], name=f"{node.name}/ValidMask")) |
| mask_scales.add(scale) |
| else: |
| nodes.append(node) |
|
|
| model.graph.ClearField("node") |
| model.graph.node.extend(nodes) |
| for current_scale in sorted(mask_scales): |
| model.graph.input.append( |
| helper.make_tensor_value_info( |
| f"valid_mask_{current_scale}", TensorProto.FLOAT, [1, 1, FRAMES * current_scale] |
| ) |
| ) |
|
|
| model = onnx.shape_inference.infer_shapes(model) |
| onnx.checker.check_model(model) |
| onnx.save(model, ROOT / "data.decoder.fixed64.onnx") |
| print("mask scales", sorted(mask_scales)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
