408 lines
12 KiB
Python
408 lines
12 KiB
Python
|
# Copyright 2025 XMOS LIMITED.
|
||
|
|
# This Software is subject to the terms of the XMOS Public Licence: Version 1.
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
import annotated_types
|
||
|
|
|
||
|
|
from audio_dsp.design.parse_json import Graph, make_pipeline, insert_forks, DspJson, pipeline_to_dspjson
|
||
|
|
|
||
|
|
from audio_dsp.models.stage import all_models
|
||
|
|
from audio_dsp.stages import all_stages
|
||
|
|
|
||
|
|
from typing import get_origin, get_args, Literal
|
||
|
|
from types import UnionType
|
||
|
|
|
||
|
|
|
||
|
|
def find_autoforks(graph):
|
||
|
|
for node in graph.nodes:
|
||
|
|
if 'AutoFork' in node.placement.name:
|
||
|
|
return True
|
||
|
|
|
||
|
|
assert False, "No AutoFork node found in the graph after insert_forks."
|
||
|
|
|
||
|
|
|
||
|
|
def test_no_shared_edge():
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "No shared edge",
|
||
|
|
"fs": 44100,
|
||
|
|
"inputs": [{
|
||
|
|
"name": "inputs",
|
||
|
|
"channels": 1
|
||
|
|
}],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "VolumeControl",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["inputs", 0]],
|
||
|
|
"name": "VolumeControl_1",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [{
|
||
|
|
"name": "outputs",
|
||
|
|
"input": [["VolumeControl_1", 0]]
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_no_shared_edge"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json()}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json()}")
|
||
|
|
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="pipeline_to_dspjson", producer_version="1.0", graph=new_graph)
|
||
|
|
new_json = pipeline_to_dspjson(a)
|
||
|
|
assert dsp_json.graph == new_json.graph, "Pipeline JSON does not match original"
|
||
|
|
|
||
|
|
for node in graph.nodes:
|
||
|
|
if 'AutoFork' in node.placement.name:
|
||
|
|
assert False, "AutoFork node found in the graph after insert_forks, but not needed."
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def test_shared_edge_from_graph_input():
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Shared edge from graph input",
|
||
|
|
"fs": 44100,
|
||
|
|
"inputs": [{
|
||
|
|
"name": "inputs",
|
||
|
|
"channels": 1
|
||
|
|
}],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "VolumeControl",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["inputs", 0]],
|
||
|
|
"name": "VolumeControl_A",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["inputs", 0]],
|
||
|
|
"name": "Mixer_B",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [{
|
||
|
|
"name": "outputs",
|
||
|
|
"input": [["VolumeControl_A", 0]]
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_shared_edge_from_graph_input"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json()}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json()}")
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
def test_shared_edge_from_producer_node():
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Shared edge from producer node",
|
||
|
|
"fs": 44100,
|
||
|
|
"inputs": [{
|
||
|
|
"name": "inputs",
|
||
|
|
"channels": 1
|
||
|
|
}],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "VolumeControl",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["inputs", 0]],
|
||
|
|
"name": "Producer",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["Producer", 0]],
|
||
|
|
"name": "Consumer_1",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["Producer", 0]],
|
||
|
|
"name": "Consumer_2",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [{
|
||
|
|
"name": "outputs",
|
||
|
|
"input": [["Consumer_1", 0]]
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_shared_edge_from_producer_node"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json()}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json()}")
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
def test_shared_edge_with_graph_output():
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Shared edge with graph output",
|
||
|
|
"fs": 44100,
|
||
|
|
"inputs": [{
|
||
|
|
"name": "inputs",
|
||
|
|
"channels": 1
|
||
|
|
}],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "VolumeControl",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["inputs", 0]],
|
||
|
|
"name": "Producer",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["Producer", 0]],
|
||
|
|
"name": "Consumer",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [{
|
||
|
|
"name": "outputs",
|
||
|
|
"input": [["Producer", 0]]
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_shared_edge_with_graph_output"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json()}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json()}")
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
def test_again():
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Stereo Mixer with Volume",
|
||
|
|
"fs": 48000,
|
||
|
|
"inputs": [{
|
||
|
|
"name": "stereo_in",
|
||
|
|
"channels": 2
|
||
|
|
}],
|
||
|
|
"nodes": [
|
||
|
|
{"op_type": "Mixer", "placement": {"input": [["stereo_in", 0], ["stereo_in", 1]], "name": "Mixer", "thread": 0}},
|
||
|
|
{"op_type": "VolumeControl", "placement": {"input": [["Mixer", 0]], "name": "Volume", "thread": 0}}
|
||
|
|
],
|
||
|
|
"outputs": [{
|
||
|
|
"name": "stereo_out",
|
||
|
|
"input": [["Volume", 0], ["Volume", 0]]
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_again"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json()}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json()}")
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
def test_multiple_inputs_outputs_non_shared():
|
||
|
|
"""
|
||
|
|
Test a graph with two inputs (one mono and one stereo) and two outputs (one mono and one stereo)
|
||
|
|
where no edge is shared between consumers.
|
||
|
|
"""
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Multiple Inputs and Outputs Non-Shared Test",
|
||
|
|
"fs": 48000,
|
||
|
|
"inputs": [
|
||
|
|
{"name": "mono_in", "channels": 1},
|
||
|
|
{"name": "stereo_in", "channels": 2}
|
||
|
|
],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "VolumeControl",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["mono_in", 0]],
|
||
|
|
"name": "VolumeControl_A",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["stereo_in", 0], ["stereo_in", 1]],
|
||
|
|
"name": "Mixer_B",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [
|
||
|
|
{"name": "mono_out", "input": [["VolumeControl_A", 0]]},
|
||
|
|
{"name": "stereo_out", "input": [["Mixer_B", 0], ["VolumeControl_A", 0]]}
|
||
|
|
]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_multiple_inputs_outputs_non_shared"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print("Test: Multiple Inputs and Outputs Non-Shared Test")
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json(indent=2)}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json(indent=2)}")
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
def test_multiple_inputs_outputs_shared():
|
||
|
|
"""
|
||
|
|
Test a graph with two inputs and two outputs where a node produces an output
|
||
|
|
that is shared by both graph outputs. A Fork should be inserted.
|
||
|
|
"""
|
||
|
|
json_str = """
|
||
|
|
{
|
||
|
|
"name": "Multiple Inputs and Outputs Shared Test",
|
||
|
|
"fs": 48000,
|
||
|
|
"inputs": [
|
||
|
|
{"name": "input1", "channels": 1},
|
||
|
|
{"name": "input2", "channels": 1}
|
||
|
|
],
|
||
|
|
"nodes": [
|
||
|
|
{
|
||
|
|
"op_type": "Mixer",
|
||
|
|
"config": {},
|
||
|
|
"placement": {
|
||
|
|
"input": [["input1", 0], ["input2", 0]],
|
||
|
|
"name": "Mixer",
|
||
|
|
"thread": 0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"outputs": [
|
||
|
|
{"name": "output1", "input": [["Mixer", 0]]},
|
||
|
|
{"name": "output2", "input": [["Mixer", 0]]}
|
||
|
|
]
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
graph = Graph.model_validate_json(json_str)
|
||
|
|
dsp_json = DspJson(ir_version=1, producer_name="test", producer_version="1.0", graph=graph)
|
||
|
|
a = make_pipeline(dsp_json)
|
||
|
|
|
||
|
|
a.draw(Path("test_multiple_inputs_outputs_shared"))
|
||
|
|
new_graph = insert_forks(graph)
|
||
|
|
print("Test: Multiple Inputs and Outputs Shared Test")
|
||
|
|
print(f"Before insert_forks: {graph.model_dump_json(indent=2)}")
|
||
|
|
print(f"After insert_forks: {new_graph.model_dump_json(indent=2)}")
|
||
|
|
|
||
|
|
find_autoforks(new_graph)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def test_all_stages_models():
|
||
|
|
all_m = all_models()
|
||
|
|
all_s = all_stages()
|
||
|
|
|
||
|
|
failed = False
|
||
|
|
for s in all_s:
|
||
|
|
if s.startswith("_") or s in ["DSPThreadStage", "PipelineStage"]:
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
assert s in all_m, f"Stage {s} not found in all models"
|
||
|
|
assert s == all_m[s].model_fields["op_type"].default, f"Stage {s} op_type mismatch"
|
||
|
|
except AssertionError as e:
|
||
|
|
print(e)
|
||
|
|
failed = True
|
||
|
|
continue
|
||
|
|
|
||
|
|
if "biquad" in s.lower() or "parametric" in s.lower():
|
||
|
|
continue
|
||
|
|
|
||
|
|
if "parameters" not in all_m[s].model_fields:
|
||
|
|
continue
|
||
|
|
|
||
|
|
|
||
|
|
try:
|
||
|
|
if type(all_s[s].set_parameters.__annotations__["parameters"]) is UnionType:
|
||
|
|
set_params_input = get_args(all_s[s].set_parameters.__annotations__["parameters"])
|
||
|
|
model_params_type = all_m[s].model_fields["parameters"].default_factory
|
||
|
|
assert model_params_type in set_params_input, f"Stage {s} set_parameters input type mismatch"
|
||
|
|
|
||
|
|
else:
|
||
|
|
set_params_input = all_s[s].set_parameters.__annotations__["parameters"].__name__
|
||
|
|
model_params_type = all_m[s].model_fields["parameters"].default_factory.__name__
|
||
|
|
assert set_params_input == model_params_type, f"Stage {s} set_parameters input type mismatch"
|
||
|
|
|
||
|
|
except AssertionError as e:
|
||
|
|
print(e)
|
||
|
|
failed = True
|
||
|
|
|
||
|
|
for field, value in type(all_m[s].model_fields["parameters"].default_factory()).model_fields.items():
|
||
|
|
try:
|
||
|
|
if get_origin(value.annotation) is list:
|
||
|
|
item_type = get_args(value.annotation)[0]
|
||
|
|
meta = get_args(item_type)[1].metadata
|
||
|
|
elif get_origin(value.annotation) is Literal:
|
||
|
|
continue # Skip Literal types
|
||
|
|
else:
|
||
|
|
meta = value.metadata
|
||
|
|
min = [g.ge for g in meta if isinstance(g, annotated_types.Ge)] or [
|
||
|
|
g.gt for g in meta if isinstance(g, annotated_types.Gt)]
|
||
|
|
max = [g.le for g in meta if isinstance(g, annotated_types.Le)] or [
|
||
|
|
g.lt for g in meta if isinstance(g, annotated_types.Lt)]
|
||
|
|
assert min, f"Minimum not defined for field {field} in {s}"
|
||
|
|
if field not in ["position", "delay"]:
|
||
|
|
assert max, f"Maximum not defined for field {field} in {s}"
|
||
|
|
assert min[0] < max[0], f"Range not correct for field {field} in {s}"
|
||
|
|
except AssertionError as e:
|
||
|
|
print(e)
|
||
|
|
failed = True
|
||
|
|
continue
|
||
|
|
|
||
|
|
assert not failed, "Some stages failed the test."
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
test_shared_edge_from_graph_input()
|