per device worker

api/.gitignore

@@ -1,6 +1,7 @@
 .coverage
 coverage.xml
 
+*.log
 *.swp
 *.pyc
 

api/launch-extras.bat

@@ -3,4 +3,4 @@ IF "%ONNX_WEB_EXTRA_MODELS%"=="" (set ONNX_WEB_EXTRA_MODELS=extras.json)
 python -m onnx_web.convert --sources --diffusion --upscaling --correction --extras=%ONNX_WEB_EXTRA_MODELS% --token=%HF_TOKEN%
 
 echo "Launching API server..."
-flask --app=onnx_web.serve run --host=0.0.0.0
+flask --app="onnx_web.serve:run" run --host=0.0.0.0

api/launch-extras.sh

@@ -25,4 +25,4 @@ python3 -m onnx_web.convert \
   --token=${HF_TOKEN:-}
 
 echo "Launching API server..."
-flask --app=onnx_web.serve run --host=0.0.0.0
+flask --app='onnx_web.main:run' run --host=0.0.0.0

api/launch.bat

@@ -2,4 +2,4 @@ echo "Downloading and converting models to ONNX format..."
 python -m onnx_web.convert --sources --diffusion --upscaling --correction --token=%HF_TOKEN%
 
 echo "Launching API server..."
-flask --app=onnx_web.serve run --host=0.0.0.0
+flask --app="onnx_web.serve:run" run --host=0.0.0.0

api/launch.sh

@@ -24,4 +24,4 @@ python3 -m onnx_web.convert \
   --token=${HF_TOKEN:-}
 
 echo "Launching API server..."
-flask --app=onnx_web.serve run --host=0.0.0.0
+flask --app='onnx_web.main:run' run --host=0.0.0.0

api/logging.yaml

@@ -5,14 +5,14 @@ formatters:
 handlers:
   console:
     class: logging.StreamHandler
-    level: INFO
+    level: DEBUG
     formatter: simple
     stream: ext://sys.stdout
 loggers:
   '':
-    level: INFO
+    level: DEBUG
     handlers: [console]
     propagate: True
 root:
-  level: INFO
+  level: DEBUG
   handlers: [console]

api/onnx_web/__init__.py

@@ -1,5 +1,10 @@
 from . import logging
-from .chain import correct_gfpgan, upscale_resrgan, upscale_stable_diffusion
+from .chain import (
+    correct_codeformer,
+    correct_gfpgan,
+    upscale_resrgan,
+    upscale_stable_diffusion,
+)
 from .diffusion.load import get_latents_from_seed, load_pipeline, optimize_pipeline
 from .diffusion.run import (
     run_blend_pipeline,
@@ -25,6 +30,7 @@
 from .onnx import OnnxNet, OnnxTensor
 from .params import (
     Border,
+    DeviceParams,
     ImageParams,
     Param,
     Point,
@@ -33,8 +39,6 @@
     UpscaleParams,
 )
 from .server import (
-    DeviceParams,
-    DevicePoolExecutor,
     ModelCache,
     ServerContext,
     apply_patch_basicsr,
@@ -51,3 +55,6 @@
     get_from_map,
     get_not_empty,
 )
+from .worker import (
+    DevicePoolExecutor,
+)

api/onnx_web/chain/__init__.py

@@ -13,3 +13,20 @@
 from .upscale_outpaint import upscale_outpaint
 from .upscale_resrgan import upscale_resrgan
 from .upscale_stable_diffusion import upscale_stable_diffusion
+
+CHAIN_STAGES = {
+    "blend-img2img": blend_img2img,
+    "blend-inpaint": blend_inpaint,
+    "blend-mask": blend_mask,
+    "correct-codeformer": correct_codeformer,
+    "correct-gfpgan": correct_gfpgan,
+    "persist-disk": persist_disk,
+    "persist-s3": persist_s3,
+    "reduce-crop": reduce_crop,
+    "reduce-thumbnail": reduce_thumbnail,
+    "source-noise": source_noise,
+    "source-txt2img": source_txt2img,
+    "upscale-outpaint": upscale_outpaint,
+    "upscale-resrgan": upscale_resrgan,
+    "upscale-stable-diffusion": upscale_stable_diffusion,
+}

api/onnx_web/chain/base.py

@@ -7,8 +7,9 @@
 
 from ..output import save_image
 from ..params import ImageParams, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
 from ..utils import is_debug
+from ..worker import ProgressCallback, WorkerContext
 from .utils import process_tile_order
 
 logger = getLogger(__name__)
@@ -17,7 +18,7 @@
 class StageCallback(Protocol):
     def __call__(
         self,
-        job: JobContext,
+        job: WorkerContext,
         ctx: ServerContext,
         stage: StageParams,
         params: ImageParams,
@@ -77,7 +78,7 @@ def append(self, stage: PipelineStage):
 
     def __call__(
         self,
-        job: JobContext,
+        job: WorkerContext,
         server: ServerContext,
         params: ImageParams,
         source: Image.Image,

api/onnx_web/chain/blend_img2img.py

@@ -7,13 +7,14 @@
 
 from ..diffusion.load import load_pipeline
 from ..params import ImageParams, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
+from ..worker import ProgressCallback, WorkerContext
 
 logger = getLogger(__name__)
 
 
 def blend_img2img(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     params: ImageParams,

api/onnx_web/chain/blend_inpaint.py

@@ -10,15 +10,16 @@
 from ..image import expand_image, mask_filter_none, noise_source_histogram
 from ..output import save_image
 from ..params import Border, ImageParams, Size, SizeChart, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
 from ..utils import is_debug
+from ..worker import ProgressCallback, WorkerContext
 from .utils import process_tile_order
 
 logger = getLogger(__name__)
 
 
 def blend_inpaint(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     stage: StageParams,
     params: ImageParams,

api/onnx_web/chain/blend_mask.py

@@ -3,18 +3,18 @@
 
 from PIL import Image
 
-from onnx_web.image import valid_image
-from onnx_web.output import save_image
-
+from ..image import valid_image
+from ..output import save_image
 from ..params import ImageParams, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
 from ..utils import is_debug
+from ..worker import ProgressCallback, WorkerContext
 
 logger = getLogger(__name__)
 
 
 def blend_mask(
-    _job: JobContext,
+    _job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/correct_codeformer.py

@@ -3,15 +3,16 @@
 from PIL import Image
 
 from ..params import ImageParams, StageParams, UpscaleParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 device = "cpu"
 
 
 def correct_codeformer(
-    job: JobContext,
+    job: WorkerContext,
     _server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/correct_gfpgan.py

@@ -5,8 +5,9 @@
 from PIL import Image
 
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
 from ..utils import run_gc
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
@@ -46,7 +47,7 @@ def load_gfpgan(
 
 
 def correct_gfpgan(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/persist_disk.py

@@ -4,13 +4,14 @@
 
 from ..output import save_image
 from ..params import ImageParams, StageParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 
 def persist_disk(
-    _job: JobContext,
+    _job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/persist_s3.py

@@ -5,13 +5,14 @@
 from PIL import Image
 
 from ..params import ImageParams, StageParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 
 def persist_s3(
-    _job: JobContext,
+    _job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/reduce_crop.py

@@ -3,13 +3,14 @@
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 
 def reduce_crop(
-    _job: JobContext,
+    _job: WorkerContext,
     _server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/reduce_thumbnail.py

@@ -3,13 +3,14 @@
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 
 def reduce_thumbnail(
-    _job: JobContext,
+    _job: WorkerContext,
     _server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/source_noise.py

@@ -4,13 +4,14 @@
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
 
 def source_noise(
-    _job: JobContext,
+    _job: WorkerContext,
     _server: ServerContext,
     _stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/source_txt2img.py

@@ -7,13 +7,14 @@
 
 from ..diffusion.load import get_latents_from_seed, load_pipeline
 from ..params import ImageParams, Size, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
+from ..worker import ProgressCallback, WorkerContext
 
 logger = getLogger(__name__)
 
 
 def source_txt2img(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     params: ImageParams,

api/onnx_web/chain/upscale_outpaint.py

@@ -10,15 +10,16 @@
 from ..image import expand_image, mask_filter_none, noise_source_histogram
 from ..output import save_image
 from ..params import Border, ImageParams, Size, SizeChart, StageParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
 from ..utils import is_debug
+from ..worker import ProgressCallback, WorkerContext
 from .utils import process_tile_grid, process_tile_order
 
 logger = getLogger(__name__)
 
 
 def upscale_outpaint(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     stage: StageParams,
     params: ImageParams,

api/onnx_web/chain/upscale_resrgan.py

@@ -6,8 +6,9 @@
 
 from ..onnx import OnnxNet
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import JobContext, ServerContext
+from ..server import ServerContext
 from ..utils import run_gc
+from ..worker import WorkerContext
 
 logger = getLogger(__name__)
 
@@ -96,7 +97,7 @@ def load_resrgan(
 
 
 def upscale_resrgan(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     stage: StageParams,
     _params: ImageParams,

api/onnx_web/chain/upscale_stable_diffusion.py

@@ -10,8 +10,9 @@
     OnnxStableDiffusionUpscalePipeline,
 )
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import JobContext, ProgressCallback, ServerContext
+from ..server import ServerContext
 from ..utils import run_gc
+from ..worker import ProgressCallback, WorkerContext
 
 logger = getLogger(__name__)
 
@@ -62,7 +63,7 @@ def load_stable_diffusion(
 
 
 def upscale_stable_diffusion(
-    job: JobContext,
+    job: WorkerContext,
     server: ServerContext,
     _stage: StageParams,
     params: ImageParams,