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
99 lines
4.8 KiB
Python
99 lines
4.8 KiB
Python
"""Convert the fixed 512px SCRFD face detector used by LatentSync."""
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnx
|
|
from onnx import helper, numpy_helper
|
|
from safetensors.numpy import save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
graph = onnx.load(args.checkpoint).graph
|
|
constants = {item.name: numpy_helper.to_array(item) for item in graph.initializer}
|
|
shapes = {graph.input[0].name: [1, 3, 512, 512]}
|
|
nodes, weights = [], {}
|
|
for node in graph.node:
|
|
op, inputs, output = node.op_type, list(node.input), node.output[0]
|
|
attrs = {item.name: helper.get_attribute_value(item) for item in node.attribute}
|
|
if op == "Shape":
|
|
constants[output] = np.array(shapes[inputs[0]], dtype=np.int64)
|
|
continue
|
|
if op == "Gather":
|
|
constants[output] = np.take(constants[inputs[0]], constants[inputs[1]], axis=attrs["axis"])
|
|
continue
|
|
if op == "Unsqueeze":
|
|
constants[output] = np.expand_dims(constants[inputs[0]], tuple(attrs["axes"]))
|
|
continue
|
|
if op == "Slice":
|
|
if constants[inputs[3]].tolist() != [0]:
|
|
raise ValueError("Unexpected SCRFD shape slice")
|
|
constants[output] = constants[inputs[0]][int(constants[inputs[1]][0]):int(constants[inputs[2]][0])]
|
|
continue
|
|
if op == "Concat":
|
|
constants[output] = np.concatenate([np.atleast_1d(constants[key]) for key in inputs], axis=attrs["axis"])
|
|
continue
|
|
shape = list(shapes[inputs[0]])
|
|
item = {"op": op, "inputs": inputs[:1], "output": output}
|
|
prefix = f"face.scrfd.{output}"
|
|
if op == "Conv":
|
|
kernel = constants[inputs[1]].astype(np.float32)
|
|
if attrs.get("group", 1) != 1 or attrs["pads"][:2] != attrs["pads"][2:]:
|
|
raise ValueError("Unsupported SCRFD convolution")
|
|
weights[prefix + ".weight"] = np.ascontiguousarray(kernel)
|
|
weights[prefix + ".bias"] = np.ascontiguousarray(constants[inputs[2]], dtype=np.float32)
|
|
item.update({"kernel": list(kernel.shape), "stride": attrs["strides"],
|
|
"padding": attrs["pads"][:2], "dilation": attrs["dilations"]})
|
|
shape[1] = kernel.shape[0]
|
|
for i in range(2):
|
|
shape[i + 2] = (shape[i + 2] + 2 * attrs["pads"][i] - attrs["dilations"][i] * (kernel.shape[i + 2] - 1) - 1) // attrs["strides"][i] + 1
|
|
elif op in ("MaxPool", "AveragePool"):
|
|
if attrs["kernel_shape"] != [2, 2] or attrs["strides"] != [2, 2] or attrs["pads"] != [0, 0, 0, 0] or shape[2] % 2 or shape[3] % 2:
|
|
raise ValueError("Unsupported SCRFD pooling layout")
|
|
shape[2] //= 2
|
|
shape[3] //= 2
|
|
elif op == "Add":
|
|
if shape != shapes[inputs[1]]:
|
|
raise ValueError("SCRFD residual shape mismatch")
|
|
item["inputs"] = inputs
|
|
elif op == "Mul":
|
|
if constants[inputs[1]].size != 1:
|
|
raise ValueError("Expected SCRFD scalar scaling")
|
|
weights[prefix + ".scale"] = constants[inputs[1]].astype(np.float32).reshape(1, 1, 1, 1)
|
|
elif op == "Resize":
|
|
if attrs["mode"] != b"nearest" or attrs["coordinate_transformation_mode"] != b"asymmetric" or attrs["nearest_mode"] != b"floor":
|
|
raise ValueError("Unsupported SCRFD resize")
|
|
shape = constants[inputs[3]].tolist()
|
|
elif op == "Transpose":
|
|
item["axes"] = attrs["perm"]
|
|
shape = [shape[axis] for axis in attrs["perm"]]
|
|
elif op == "Reshape":
|
|
target = constants[inputs[1]].tolist()
|
|
elements = int(np.prod(shape))
|
|
if target.count(-1) == 1:
|
|
target[target.index(-1)] = elements // -int(np.prod(target))
|
|
if int(np.prod(target)) != elements:
|
|
raise ValueError("Invalid SCRFD reshape")
|
|
shape = target
|
|
elif op not in ("Relu", "Sigmoid"):
|
|
raise ValueError(f"Unsupported SCRFD operation: {op}")
|
|
shapes[output] = shape
|
|
item["shape"] = shape
|
|
nodes.append(item)
|
|
outputs = [value.name for value in graph.output]
|
|
config = {"input": graph.input[0].name, "input_shape": [1, 3, 512, 512],
|
|
"nodes": nodes, "outputs": outputs, "output_shapes": [shapes[key] for key in outputs]}
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
save_file(weights, args.output / "weights.safetensors")
|
|
(args.output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
print(f"Converted SCRFD: {len(nodes)} nodes, {len(weights)} tensors; heads {config['output_shapes']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|