apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

add missing arg to chain pipeline

This commit is contained in:
Sean Sube 2023-02-04 14:27:27 -06:00
parent 2f1f6b3a11
commit 8a81e8b810
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,10 +1,12 @@
from datetime import timedelta from datetime import timedelta
from logging import getLogger from logging import getLogger
from PIL import Image from PIL import Image
from os import path
from time import monotonic from time import monotonic
from typing import Any, List, Optional, Protocol, Tuple from typing import Any, List, Optional, Protocol, Tuple
from ..device_pool import (
JobContext,
)
from ..params import ( from ..params import (
ImageParams, ImageParams,
StageParams, StageParams,
@ -59,7 +61,7 @@ class ChainPipeline:
''' '''
self.stages.append(stage) self.stages.append(stage)
def __call__(self, ctx: ServerContext, params: ImageParams, source: Image.Image, **pipeline_kwargs) -> Image.Image: def __call__(self, job: JobContext, server: ServerContext, params: ImageParams, source: Image.Image, **pipeline_kwargs) -> Image.Image:
''' '''
TODO: handle List[Image] outputs TODO: handle List[Image] outputs
''' '''
@ -81,11 +83,11 @@ class ChainPipeline:
stage_params.tile_size) stage_params.tile_size)
def stage_tile(tile: Image.Image, _dims) -> Image.Image: def stage_tile(tile: Image.Image, _dims) -> Image.Image:
tile = stage_pipe(ctx, stage_params, params, tile, tile = stage_pipe(server, stage_params, params, tile,
**kwargs) **kwargs)
if is_debug(): if is_debug():
save_image(ctx, 'last-tile.png', tile) save_image(server, 'last-tile.png', tile)
return tile return tile
@ -93,14 +95,14 @@ class ChainPipeline:
image, stage_params.tile_size, stage_params.outscale, [stage_tile]) image, stage_params.tile_size, stage_params.outscale, [stage_tile])
else: else:
logger.info('image within tile size, running stage') logger.info('image within tile size, running stage')
image = stage_pipe(ctx, stage_params, params, image, image = stage_pipe(server, stage_params, params, image,
**kwargs) **kwargs)
logger.info('finished stage %s, result size: %sx%s', logger.info('finished stage %s, result size: %sx%s',
name, image.width, image.height) name, image.width, image.height)
if is_debug(): if is_debug():
save_image(ctx, 'last-stage.png', image) save_image(server, 'last-stage.png', image)
end = monotonic() end = monotonic()
duration = timedelta(seconds=(end - start)) duration = timedelta(seconds=(end - start))