apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix output stage

This commit is contained in:
Sean Sube 2024-01-08 22:54:41 -06:00
parent 99bd6b0708
commit 0d1439734c
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 4 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
api/onnx_web/chain/persist_disk.py
View File

@ -3,7 +3,7 @@ from typing import List, Optional
from PIL import Image from PIL import Image
from ..output import save_image from ..output import save_image, save_result
from ..params import ImageParams, Size, SizeChart, StageParams from ..params import ImageParams, Size, SizeChart, StageParams
from ..server import ServerContext from ..server import ServerContext
from ..worker import WorkerContext from ..worker import WorkerContext
@ -18,28 +18,17 @@ class PersistDiskStage(BaseStage):
def run( def run(
self, self,
_worker: WorkerContext, worker: WorkerContext,
server: ServerContext, server: ServerContext,
_stage: StageParams, _stage: StageParams,
params: ImageParams, params: ImageParams,
sources: StageResult, sources: StageResult,
*, *,
output: List[str],
size: Optional[Size] = None,
stage_source: Optional[Image.Image] = None, stage_source: Optional[Image.Image] = None,
**kwargs, **kwargs,
) -> StageResult: ) -> StageResult:
logger.info("persisting %s images to disk: %s", len(sources), output) logger.info("persisting %s images to disk", len(sources))
for name, source, metadata in zip( save_result(server, sources, worker.job)
output, sources.as_images(), sources.metadata
):
dest = save_image(
server,
name,
source,
metadata=metadata,
)
logger.info("saved image to %s", dest)
return sources return sources