apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): correctly handle image stacks in persist stages

This commit is contained in:
Sean Sube 2023-11-27 21:45:32 -06:00
parent 828c9511a6
commit 042181b9c5
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -31,7 +31,7 @@ class PersistDiskStage(BaseStage):
) -> StageResult:
logger.info("persisting %s images to disk: %s", len(sources), output)
for source, name in zip(sources, output):
for source, name in zip(sources.as_image(), output):
dest = save_image(server, name, source, params=params, size=size)
logger.info("saved image to %s", dest)

10
api/onnx_web/chain/persist_s3.py
View File

@ -1,6 +1,6 @@
from io import BytesIO
from logging import getLogger
from typing import Optional
from typing import List, Optional
from boto3 import Session
from PIL import Image
@ -23,7 +23,7 @@ class PersistS3Stage(BaseStage):
_params: ImageParams,
sources: StageResult,
*,
output: str,
output: List[str],
bucket: str,
endpoint_url: Optional[str] = None,
profile_name: Optional[str] = None,
@ -33,14 +33,14 @@ class PersistS3Stage(BaseStage):
session = Session(profile_name=profile_name)
s3 = session.client("s3", endpoint_url=endpoint_url)
for source in sources.as_image():
for source, name in zip(sources.as_image(), output):
data = BytesIO()
source.save(data, format=server.image_format)
data.seek(0)
try:
s3.upload_fileobj(data, bucket, output)
logger.info("saved image to s3://%s/%s", bucket, output)
s3.upload_fileobj(data, bucket, name)
logger.info("saved image to s3://%s/%s", bucket, name)
except Exception:
logger.exception("error saving image to S3")