feat(api): switch to python logging

api/logging.yaml

@@ -0,0 +1,18 @@
+version: 1
+formatters:
+  simple:
+    format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+handlers:
+  console:
+    class: logging.StreamHandler
+    level: INFO
+    formatter: simple
+    stream: ext://sys.stdout
+loggers:
+  '':
+    level: INFO
+    handlers: [console]
+    propagate: True
+root:
+  level: INFO
+  handlers: [console]

api/onnx_web/chain/base.py

@@ -1,3 +1,5 @@
+from datetime import timedelta
+from logging import getLogger
 from PIL import Image
 from os import path
 from time import monotonic
@@ -15,6 +17,8 @@
     process_tiles,
 )
 
+logger = getLogger(__name__)
+
 
 class StageCallback(Protocol):
     def __call__(
@@ -57,21 +61,21 @@ def __call__(self, ctx: ServerContext, params: ImageParams, source: Image.Image,
         TODO: handle List[Image] outputs
         '''
         start = monotonic()
-        print('running pipeline on source image with dimensions %sx%s' %
-              source.size)
+        logger.info('running pipeline on source image with dimensions %sx%s',
+                    source.width, source.height)
         image = source
 
         for stage_pipe, stage_params, stage_kwargs in self.stages:
             name = stage_params.name or stage_pipe.__name__
             kwargs = stage_kwargs or {}
             kwargs = {**pipeline_kwargs, **kwargs}
 
-            print('running stage %s on result image with dimensions %sx%s, %s' %
-                  (name, image.width, image.height, kwargs))
+            logger.info('running stage %s on result image with dimensions %sx%s, %s',
+                        name, image.width, image.height, kwargs)
 
             if image.width > stage_params.tile_size or image.height > stage_params.tile_size:
-                print('source image larger than tile size of %s, tiling stage' % (
-                    stage_params.tile_size))
+                logger.info('source image larger than tile size of %s, tiling stage',
+                            stage_params.tile_size)
 
                 def stage_tile(tile: Image.Image, _dims) -> Image.Image:
                     tile = stage_pipe(ctx, stage_params, params, tile,
@@ -85,14 +89,15 @@ def stage_tile(tile: Image.Image, _dims) -> Image.Image:
                 image = process_tiles(
                     image, stage_params.tile_size, stage_params.outscale, [stage_tile])
             else:
-                print('source image within tile size, running stage')
+                logger.info('source image within tile size, running stage')
                 image = stage_pipe(ctx, stage_params, params, image,
                                    **kwargs)
 
-            print('finished stage %s, result size: %sx%s' %
-                  (name, image.width, image.height))
+            logger.info('finished stage %s, result size: %sx%s',
+                        name, image.width, image.height)
 
         end = monotonic()
-        duration = end - start
-        print('finished pipeline in %s seconds, result size: %sx%s' % (duration, image.width, image.height))
+        duration = timedelta(seconds=(end - start))
+        logger.info('finished pipeline in %s, result size: %sx%s',
+                    duration, image.width, image.height)
         return image

api/onnx_web/chain/blend_img2img.py

@@ -1,6 +1,7 @@
 from diffusers import (
     OnnxStableDiffusionImg2ImgPipeline,
 )
+from logging import getLogger
 from PIL import Image
 
 from ..diffusion import (
@@ -16,16 +17,18 @@
 
 import numpy as np
 
+logger = getLogger(__name__)
+
 
 def blend_img2img(
-    ctx: ServerContext,
-    stage: StageParams,
+    _ctx: ServerContext,
+    _stage: StageParams,
     params: ImageParams,
     source_image: Image.Image,
     *,
     strength: float,
 ) -> Image.Image:
-    print('generating image using img2img', params.prompt)
+    logger.info('generating image using img2img', params.prompt)
 
     pipe = load_pipeline(OnnxStableDiffusionImg2ImgPipeline,
                             params.model, params.provider, params.scheduler)
@@ -43,6 +46,6 @@ def blend_img2img(
     )
     output = result.images[0]
 
-    print('final output image size', output.size)
+    logger.info('final output image size', output.size)
     return output
 

api/onnx_web/chain/blend_inpaint.py

@@ -1,6 +1,7 @@
 from diffusers import (
     OnnxStableDiffusionInpaintPipeline,
 )
+from logging import getLogger
 from PIL import Image
 from typing import Callable, Tuple
 
@@ -31,6 +32,8 @@
 
 import numpy as np
 
+logger = getLogger(__name__)
+
 
 def blend_inpaint(
     ctx: ServerContext,
@@ -44,7 +47,7 @@ def blend_inpaint(
     mask_filter: Callable = mask_filter_none,
     noise_source: Callable = noise_source_histogram,
 ) -> Image.Image:
-    print('upscaling image by expanding borders', expand)
+    logger.info('upscaling image by expanding borders', expand)
 
     if mask_image is None:
         # if no mask was provided, keep the full source image
@@ -96,5 +99,5 @@ def outpaint(image: Image.Image, dims: Tuple[int, int, int]):
 
     output = process_tiles(source_image, SizeChart.auto, 1, [outpaint])
 
-    print('final output image size', output.size)
+    logger.info('final output image size', output.size)
     return output

api/onnx_web/chain/correct_gfpgan.py

@@ -1,4 +1,5 @@
 from gfpgan import GFPGANer
+from logging import getLogger
 from os import path
 from PIL import Image
 from realesrgan import RealESRGANer
@@ -16,6 +17,8 @@
     load_resrgan,
 )
 
+logger = getLogger(__name__)
+
 
 last_pipeline_instance = None
 last_pipeline_params = None
@@ -32,7 +35,7 @@ def load_gfpgan(ctx: ServerContext, upscale: UpscaleParams):
                           (upscale.correction_model))
 
     if last_pipeline_instance != None and face_path == last_pipeline_params:
-        print('reusing existing GFPGAN pipeline')
+        logger.info('reusing existing GFPGAN pipeline')
         return last_pipeline_instance
 
     # TODO: doesn't have a model param, not sure how to pass ONNX model
@@ -59,10 +62,10 @@ def correct_gfpgan(
     upsampler: Optional[RealESRGANer] = None,
 ) -> Image.Image:
     if upscale.correction_model is None:
-        print('no face model given, skipping')
+        logger.warn('no face model given, skipping')
         return image
 
-    print('correcting faces with GFPGAN model: %s' % upscale.correction_model)
+    logger.info('correcting faces with GFPGAN model: %s', upscale.correction_model)
     gfpgan = load_gfpgan(ctx, upscale)
 
     _, _, output = gfpgan.enhance(

api/onnx_web/chain/persist_disk.py

@@ -1,3 +1,4 @@
+from logging import getLogger
 from PIL import Image
 
 
@@ -10,6 +11,8 @@
     ServerContext,
 )
 
+logger = getLogger(__name__)
+
 
 def persist_disk(
     ctx: ServerContext,
@@ -21,5 +24,5 @@ def persist_disk(
 ) -> Image.Image:
     dest = base_join(ctx.output_path, output)
     source_image.save(dest)
-    print('saved image to %s' % (dest,))
+    logger.info('saved image to %s', dest)
     return source_image

api/onnx_web/chain/persist_s3.py

@@ -2,6 +2,7 @@
     Session,
 )
 from io import BytesIO
+from logging import getLogger
 from PIL import Image
 
 from ..params import (
@@ -12,6 +13,8 @@
     ServerContext,
 )
 
+logger = getLogger(__name__)
+
 
 def persist_s3(
     _ctx: ServerContext,
@@ -32,9 +35,9 @@ def persist_s3(
     data.seek(0)
 
     try:
-        response = s3.upload_fileobj(data, bucket, output)
-        print('saved image to %s' % (response))
+        s3.upload_fileobj(data, bucket, output)
+        logger.info('saved image to %s/%s', bucket, output)
     except Exception as err:
-        print('error saving image to S3: %s' % (err))
+        logger.error('error saving image to S3: %s', err)
 
     return source_image

api/onnx_web/chain/source_txt2img.py

@@ -1,6 +1,7 @@
 from diffusers import (
     OnnxStableDiffusionPipeline,
 )
+from logging import getLogger
 from PIL import Image
 
 from ..diffusion import (
@@ -18,6 +19,8 @@
 
 import numpy as np
 
+logger = getLogger(__name__)
+
 
 def source_txt2img(
     ctx: ServerContext,
@@ -27,10 +30,10 @@ def source_txt2img(
     *,
     size: Size,
 ) -> Image.Image:
-    print('generating image using txt2img', params.prompt)
+    logger.info('generating image using txt2img, %s steps', params.steps)
 
     if source_image is not None:
-        print('a source image was passed to a txt2img stage, but will be discarded')
+        logger.warn('a source image was passed to a txt2img stage, but will be discarded')
 
     pipe = load_pipeline(OnnxStableDiffusionPipeline,
                             params.model, params.provider, params.scheduler)
@@ -50,5 +53,5 @@ def source_txt2img(
     )
     output = result.images[0]
 
-    print('final output image size', output.size)
+    logger.info('final output image size: %sx%s', output.width, output.height)
     return output

api/onnx_web/chain/upscale_outpaint.py

@@ -1,6 +1,7 @@
 from diffusers import (
     OnnxStableDiffusionInpaintPipeline,
 )
+from logging import getLogger
 from PIL import Image
 from typing import Callable, Tuple
 
@@ -31,6 +32,8 @@
 
 import numpy as np
 
+logger = getLogger(__name__)
+
 
 def upscale_outpaint(
     ctx: ServerContext,
@@ -44,7 +47,7 @@ def upscale_outpaint(
     mask_filter: Callable = mask_filter_none,
     noise_source: Callable = noise_source_histogram,
 ) -> Image.Image:
-    print('upscaling image by expanding borders', expand)
+    logger.info('upscaling image by expanding borders: %s', expand)
 
     if mask_image is None:
         # if no mask was provided, keep the full source image
@@ -94,7 +97,7 @@ def outpaint(image: Image.Image, dims: Tuple[int, int, int]):
         )
         return result.images[0]
 
-    output = process_tiles(source_image, SizeChart.auto.value, 1, [outpaint])
+    output = process_tiles(source_image, SizeChart.auto, 1, [outpaint])
 
-    print('final output image size', output.size)
+    logger.info('final output image size: %sx%s', output.width, output.height)
     return output

api/onnx_web/chain/upscale_resrgan.py

@@ -1,4 +1,5 @@
 from basicsr.archs.rrdbnet_arch import RRDBNet
+from logging import getLogger
 from os import path
 from PIL import Image
 from realesrgan import RealESRGANer
@@ -17,6 +18,7 @@
 
 import numpy as np
 
+logger = getLogger(__name__)
 
 last_pipeline_instance = None
 last_pipeline_params = (None, None)
@@ -33,7 +35,7 @@ def load_resrgan(ctx: ServerContext, params: UpscaleParams, tile=0):
 
     cache_params = (model_path, params.format)
     if last_pipeline_instance != None and cache_params == last_pipeline_params:
-        print('reusing existing Real ESRGAN pipeline')
+        logger.info('reusing existing Real ESRGAN pipeline')
         return last_pipeline_instance
 
     # use ONNX acceleration, if available
@@ -76,13 +78,13 @@ def upscale_resrgan(
     *,
     upscale: UpscaleParams,
 ) -> Image.Image:
-    print('upscaling image with Real ESRGAN', upscale.scale)
+    logger.info('upscaling image with Real ESRGAN', upscale.scale)
 
     output = np.array(source_image)
     upsampler = load_resrgan(ctx, upscale, tile=stage.tile_size)
 
     output, _ = upsampler.enhance(output, outscale=upscale.outscale)
 
     output = Image.fromarray(output, 'RGB')
-    print('final output image size', output.size)
+    logger.info('final output image size', output.size)
     return output

api/onnx_web/chain/upscale_stable_diffusion.py

@@ -3,6 +3,7 @@
     DDPMScheduler,
     StableDiffusionUpscalePipeline,
 )
+from logging import getLogger
 from os import path
 from PIL import Image
 
@@ -20,6 +21,8 @@
 
 import torch
 
+logger = getLogger(__name__)
+
 
 last_pipeline_instance = None
 last_pipeline_params = (None, None)
@@ -33,7 +36,7 @@ def load_stable_diffusion(ctx: ServerContext, upscale: UpscaleParams):
     cache_params = (model_path, upscale.format)
 
     if last_pipeline_instance != None and cache_params == last_pipeline_params:
-        print('reusing existing Stable Diffusion upscale pipeline')
+        logger.info('reusing existing Stable Diffusion upscale pipeline')
         return last_pipeline_instance
 
     if upscale.format == 'onnx':
@@ -65,7 +68,7 @@ def upscale_stable_diffusion(
     *,
     upscale: UpscaleParams,
 ) -> Image.Image:
-    print('upscaling with Stable Diffusion')
+    logger.info('upscaling with Stable Diffusion')
 
     pipeline = load_stable_diffusion(ctx, upscale)
     generator = torch.manual_seed(params.seed)