Spaces:
Runtime error
Runtime error
| from sys import float_info | |
| import numpy as np | |
| from onnx import numpy_helper as onph | |
| from onnx.onnx_pb import AttributeProto, NodeProto, TensorProto | |
| from nodes.log import logger | |
| INT64_MIN, INT64_MAX = np.iinfo(np.int64).min, np.iinfo(np.int64).max | |
| FLOAT32_MAX = float_info.max | |
| APT = AttributeProto | |
| TPT = TensorProto | |
| def get_node_attr_ai(node: NodeProto, key: str) -> np.ndarray: | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return np.array( | |
| [max(min(i, INT64_MAX), INT64_MIN) for i in attr.ints], np.int64 | |
| ) | |
| return np.empty(0, np.int32) | |
| def set_node_attr_ai(node: NodeProto, key: str, value: np.ndarray) -> None: | |
| attr_group = AttributeProto(name=key, floats=value, type=APT.INTS) | |
| node.attribute.append(attr_group) | |
| def get_node_attr_af(node: NodeProto, key: str) -> np.ndarray: | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return np.array(list(attr.floats), np.float32) | |
| return np.empty(0, np.float32) | |
| def get_node_attr_i(node: NodeProto, key: str, default: int = 0) -> int: | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return max(min(attr.i, INT64_MAX), INT64_MIN) | |
| return default | |
| def get_node_attr_f(node: NodeProto, key: str, default: float = 0) -> float: | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return attr.f | |
| return default | |
| def get_node_attr_s(node: NodeProto, key: str, default: str = ""): | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return attr.s.decode("ascii") | |
| return default | |
| def get_node_attr_tensor(node: NodeProto, key: str) -> TensorProto: | |
| for attr in node.attribute: | |
| if attr.name == key: | |
| return attr.t | |
| return TensorProto() | |
| def get_node_attr_from_input_f(tp: TensorProto) -> float: | |
| shape_data = onph.to_array(tp) | |
| if tp.data_type in (TPT.FLOAT, TPT.FLOAT16, TPT.DOUBLE, TPT.INT32): | |
| f = shape_data.item(0) | |
| elif tp.data_type == TPT.INT64: | |
| f = max(min(shape_data.item(0), INT64_MAX), INT64_MIN) | |
| else: | |
| raise TypeError(f"Unknown data type {tp.data_type}") | |
| return f | |
| def get_node_attr_from_input_ai(tp: TensorProto) -> np.ndarray: | |
| if tp.data_type == TPT.INT32 or tp.data_type == TPT.INT64: | |
| shape_data = onph.to_array(tp) | |
| if shape_data.size == 1: | |
| shape_data = np.array([shape_data.item(0)], shape_data.dtype) | |
| return np.array( | |
| [ | |
| max(min(val, INT64_MAX), INT64_MIN) | |
| if tp.data_type == TPT.INT64 | |
| else val | |
| for val in shape_data | |
| ], | |
| shape_data.dtype, | |
| ) | |
| else: | |
| logger.error(f"chaiNNer: unknown data type {tp.data_type}") | |
| return np.empty(0, np.int32) | |
| def get_node_attr_from_input_af(tp: TensorProto) -> np.ndarray: | |
| if tp.data_type in (TPT.FLOAT, TPT.FLOAT16, TPT.DOUBLE): | |
| shape_data = onph.to_array(tp) | |
| return np.array(list(shape_data), shape_data.dtype) | |
| else: | |
| logger.error(f"chaiNNer: unknown data type {tp.data_type}") | |
| return np.empty(0, np.float32) | |
| def get_tensor_proto_data_size(tp: TensorProto, fpmode: int = TPT.FLOAT) -> int: | |
| if tp.raw_data: | |
| if fpmode == TPT.FLOAT16: | |
| return len(tp.raw_data) // 2 | |
| return len(tp.raw_data) // 4 | |
| elif tp.data_type == TPT.FLOAT or tp.data_type == TPT.FLOAT16: | |
| return len(tp.float_data) | |
| return 0 | |