apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

use pipeline output count

This commit is contained in:
Sean Sube 2023-09-13 08:43:31 -05:00
parent 046de9bf3a
commit 505e408dd6
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
api/onnx_web/chain/base.py
View File

@ -92,6 +92,13 @@ class ChainPipeline:
return steps
def outputs(self, params: ImageParams, sources: int):
outputs = sources
for callback, _params, _kwargs in self.stages:
outputs += callback.outputs(params, outputs)
return outputs
def __call__(
self,
worker: WorkerContext,

4
api/onnx_web/server/api.py
View File

@ -455,7 +455,7 @@ def chain(server: ServerContext, pool: DevicePoolExecutor):
logger.info("running chain pipeline with %s stages", len(pipeline.stages))
output = make_output_name(
server, "chain", base_params, base_size, count=len(pipeline.stages)
server, "chain", base_params, base_size, count=pipeline.outputs(base_params, 0))
)
job_name = output[0]
@ -471,7 +471,7 @@ def chain(server: ServerContext, pool: DevicePoolExecutor):
needs_device=device,
)
step_params = params.with_args(steps=pipeline.steps(base_params, base_size))
step_params = base_params.with_args(steps=pipeline.steps(base_params, base_size))
return jsonify(json_params(output, step_params, base_size))