feat(api): add chain pipeline stages to load images from S3 or URL (#184

, #185)

api/onnx_web/chain/__init__.py

@@ -9,7 +9,9 @@
 from .reduce_crop import reduce_crop
 from .reduce_thumbnail import reduce_thumbnail
 from .source_noise import source_noise
+from .source_s3 import source_s3
 from .source_txt2img import source_txt2img
+from .source_url import source_url
 from .upscale_bsrgan import upscale_bsrgan
 from .upscale_outpaint import upscale_outpaint
 from .upscale_resrgan import upscale_resrgan
@@ -27,7 +29,9 @@
     "reduce-crop": reduce_crop,
     "reduce-thumbnail": reduce_thumbnail,
     "source-noise": source_noise,
+    "source-s3": source_s3,
     "source-txt2img": source_txt2img,
+    "source-url": source_url,
     "upscale-bsrgan": upscale_bsrgan,
     "upscale-outpaint": upscale_outpaint,
     "upscale-resrgan": upscale_resrgan,

api/onnx_web/chain/persist_s3.py

@@ -37,7 +37,7 @@ def persist_s3(
 
     try:
         s3.upload_fileobj(data, bucket, output)
-        logger.info("saved image to %s/%s", bucket, output)
+        logger.info("saved image to s3://%s/%s", bucket, output)
     except Exception:
         logger.exception("error saving image to S3")
 

api/onnx_web/chain/source_s3.py

@@ -0,0 +1,42 @@
+from io import BytesIO
+from logging import getLogger
+from typing import Optional
+
+from boto3 import Session
+from PIL import Image
+
+from ..params import ImageParams, StageParams
+from ..server import ServerContext
+from ..worker import WorkerContext
+
+logger = getLogger(__name__)
+
+
+def source_s3(
+    _job: WorkerContext,
+    server: ServerContext,
+    _stage: StageParams,
+    _params: ImageParams,
+    source: Image.Image,
+    *,
+    source_key: str,
+    bucket: str,
+    endpoint_url: Optional[str] = None,
+    profile_name: Optional[str] = None,
+    stage_source: Optional[Image.Image] = None,
+    **kwargs,
+) -> Image.Image:
+    source = stage_source or source
+
+    session = Session(profile_name=profile_name)
+    s3 = session.client("s3", endpoint_url=endpoint_url)
+
+    try:
+        logger.info("loading image from s3://%s/%s", bucket, source_key)
+        data = BytesIO()
+        s3.download_fileobj(bucket, source_key, data)
+
+        data.seek(0)
+        return Image.open(data)
+    except Exception:
+        logger.exception("error loading image from S3")

api/onnx_web/chain/source_txt2img.py

@@ -33,7 +33,7 @@ def source_txt2img(
 
     if "stage_source" in kwargs:
         logger.warn(
-            "a source image was passed to a txt2img stage, but will be discarded"
+            "a source image was passed to a txt2img stage, and will be discarded"
         )
 
     latents = get_latents_from_seed(params.seed, size)

api/onnx_web/chain/source_url.py

@@ -0,0 +1,37 @@
+from io import BytesIO
+from logging import getLogger
+
+import requests
+from PIL import Image
+
+from ..params import ImageParams, StageParams
+from ..server import ServerContext
+from ..worker import WorkerContext
+
+logger = getLogger(__name__)
+
+
+def source_url(
+    _job: WorkerContext,
+    _server: ServerContext,
+    _stage: StageParams,
+    _params: ImageParams,
+    source: Image.Image,
+    *,
+    source_url: str,
+    stage_source: Image.Image,
+    **kwargs,
+) -> Image.Image:
+    source = stage_source or source
+    logger.info("loading image from URL source")
+
+    if source is not None:
+        logger.warn(
+            "a source image was passed to a source stage, and will be discarded"
+        )
+
+    response = requests.get(source_url)
+    output = Image.open(BytesIO(response.content))
+
+    logger.info("final output image size: %sx%s", output.width, output.height)
+    return output