add dark mode with toggle and system preference

.vscode/launch.json

@@ -0,0 +1,21 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Python: Remote Attach",
+      "type": "python",
+      "request": "attach",
+      "connect": {
+        "host": "127.0.0.1",
+        "port": 5678
+      },
+      "pathMappings": [
+        {
+          "localRoot": "${workspaceFolder}",
+          "remoteRoot": "/opt/onnx-web"
+        }
+      ],
+      "justMyCode": true
+    }
+  ]
+}

api/entry.py

@@ -0,0 +1,29 @@
+import os
+
+def script_method(fn, _rcb=None):
+    return fn
+
+def script(obj, optimize=True, _frames_up=0, _rcb=None):
+    return obj
+
+import torch.jit
+torch.jit.script_method = script_method
+torch.jit.script = script
+
+import multiprocessing
+
+if __name__ == '__main__':
+    multiprocessing.freeze_support()
+    try:
+        from onnx_web.main import main
+        app, pool = main()
+        print("starting workers")
+        pool.start()
+        print("starting flask")
+        app.run("0.0.0.0", 5000, debug=False)
+        input("press the any key")
+        pool.join()
+    except Exception as e:
+        print(e)
+    finally:
+        os.system("pause")

api/onnx_web/__init__.py

@@ -16,6 +16,7 @@
     run_upscale_pipeline,
 )
 from .diffusers.stub_scheduler import StubScheduler
+from .diffusers.upscale import run_upscale_correction
 from .image import (
     expand_image,
     mask_filter_gaussian_multiply,
@@ -48,7 +49,6 @@
     apply_patch_facexlib,
     apply_patches,
 )
-from .upscale import run_upscale_correction
 from .utils import (
     base_join,
     get_and_clamp_float,

api/onnx_web/chain/base.py

@@ -16,6 +16,10 @@
 
 
 class StageCallback(Protocol):
+    """
+    Definition for a stage job function.
+    """
+
     def __call__(
         self,
         job: WorkerContext,
@@ -25,6 +29,9 @@ def __call__(
         source: Image.Image,
         **kwargs: Any
     ) -> Image.Image:
+        """
+        Run this stage against a source image.
+        """
         pass
 
 

api/onnx_web/chain/utils.py

@@ -1,4 +1,5 @@
 from logging import getLogger
+from math import ceil
 from typing import List, Protocol, Tuple
 
 from PIL import Image
@@ -9,7 +10,14 @@
 
 
 class TileCallback(Protocol):
+    """
+    Definition for a tile job function.
+    """
+
     def __call__(self, image: Image.Image, dims: Tuple[int, int, int]) -> Image.Image:
+        """
+        Run this stage against a single tile.
+        """
         pass
 
 
@@ -58,31 +66,13 @@ def process_tile_spiral(
     image = Image.new("RGB", (width * scale, height * scale))
     image.paste(source, (0, 0, width, height))
 
-    center_x = (width // 2) - (tile // 2)
-    center_y = (height // 2) - (tile // 2)
-
-    # TODO: should add/remove tiles when overlap != 0.5
-    tiles = [
-        (0, tile * -overlap),
-        (tile * overlap, tile * -overlap),
-        (tile * overlap, 0),
-        (tile * overlap, tile * overlap),
-        (0, tile * overlap),
-        (tile * -overlap, tile * overlap),
-        (tile * -overlap, 0),
-        (tile * -overlap, tile * -overlap),
-    ]
-
     # tile tuples is source, multiply by scale for dest
     counter = 0
+    tiles = generate_tile_spiral(width, height, tile, overlap=overlap)
     for left, top in tiles:
-        left = center_x + int(left)
-        top = center_y + int(top)
-
         counter += 1
         logger.debug("processing tile %s of %s, %sx%s", counter, len(tiles), left, top)
 
-        # TODO: only valid for scale == 1, resize source for others
         tile_image = image.crop((left, top, left + tile, top + tile))
 
         for filter in filters:
@@ -113,3 +103,85 @@ def process_tile_order(
     else:
         logger.warn("unknown tile order: %s", order)
         raise ValueError()
+
+
+def generate_tile_spiral(
+    width: int,
+    height: int,
+    tile: int,
+    overlap: float = 0.0,
+) -> List[Tuple[int, int]]:
+    spacing = 1.0 - overlap
+
+    # round dims up to nearest tiles
+    tile_width = ceil(width / tile)
+    tile_height = ceil(height / tile)
+
+    # start walking from the north-west corner, heading east
+    dir_height = 0
+    dir_width = 1
+
+    walk_height = tile_height
+    walk_width = tile_width
+
+    accum_height = 0
+    accum_width = 0
+
+    tile_top = 0
+    tile_left = 0
+
+    tile_coords = []
+    while walk_width > 0 and walk_height > 0:
+        # exhaust the current direction, then turn
+        while accum_width < walk_width and accum_height < walk_height:
+            # add a tile
+            logger.trace(
+                "adding tile at %s:%s, %s:%s, %s:%s",
+                tile_left,
+                tile_top,
+                accum_width,
+                accum_height,
+                walk_width,
+                walk_height,
+                spacing,
+            )
+            tile_coords.append((int(tile_left), int(tile_top)))
+
+            # move to the next
+            tile_top += dir_height * spacing * tile
+            tile_left += dir_width * spacing * tile
+
+            accum_height += abs(dir_height * spacing)
+            accum_width += abs(dir_width * spacing)
+
+        # reset for the next direction
+        accum_height = 0
+        accum_width = 0
+
+        # why tho
+        tile_top -= dir_height
+        tile_left -= dir_width
+
+        # turn right
+        if dir_width == 1 and dir_height == 0:
+            dir_width = 0
+            dir_height = 1
+        elif dir_width == 0 and dir_height == 1:
+            dir_width = -1
+            dir_height = 0
+        elif dir_width == -1 and dir_height == 0:
+            dir_width = 0
+            dir_height = -1
+        elif dir_width == 0 and dir_height == -1:
+            dir_width = 1
+            dir_height = 0
+
+        # step to the next tile as part of the turn
+        tile_top += dir_height
+        tile_left += dir_width
+
+        # shrink the last direction
+        walk_height -= abs(dir_height)
+        walk_width -= abs(dir_width)
+
+    return tile_coords

api/onnx_web/constants.py

@@ -0,0 +1,2 @@
+ONNX_MODEL = "model.onnx"
+ONNX_WEIGHTS = "weights.pb"

api/onnx_web/convert/__main__.py

@@ -12,6 +12,7 @@
 from transformers import CLIPTokenizer
 from yaml import safe_load
 
+from ..constants import ONNX_MODEL, ONNX_WEIGHTS
 from .correction_gfpgan import convert_correction_gfpgan
 from .diffusion.diffusers import convert_diffusion_diffusers
 from .diffusion.lora import blend_loras
@@ -297,7 +298,7 @@ def convert_models(ctx: ConversionContext, args, models: Models):
                                     path.join(
                                         dest,
                                         "text_encoder",
-                                        "model.onnx",
+                                        ONNX_MODEL,
                                     )
                                 )
 
@@ -341,13 +342,13 @@ def convert_models(ctx: ConversionContext, args, models: Models):
                                     path.join(
                                         dest,
                                         "text_encoder",
-                                        "model.onnx",
+                                        ONNX_MODEL,
                                     )
                                 )
 
                             if "unet" not in blend_models:
                                 blend_models["text_encoder"] = load_model(
-                                    path.join(dest, "unet", "model.onnx")
+                                    path.join(dest, "unet", ONNX_MODEL)
                                 )
 
                             # load models if not loaded yet
@@ -377,7 +378,7 @@ def convert_models(ctx: ConversionContext, args, models: Models):
 
                         for name in ["text_encoder", "unet"]:
                             if name in blend_models:
-                                dest_path = path.join(dest, name, "model.onnx")
+                                dest_path = path.join(dest, name, ONNX_MODEL)
                                 logger.debug(
                                     "saving blended %s model to %s", name, dest_path
                                 )
@@ -386,7 +387,7 @@ def convert_models(ctx: ConversionContext, args, models: Models):
                                     dest_path,
                                     save_as_external_data=True,
                                     all_tensors_to_one_file=True,
-                                    location="weights.pb",
+                                    location=ONNX_WEIGHTS,
                                 )
 
                 except Exception:
@@ -459,7 +460,7 @@ def main() -> int:
         "--half",
         action="store_true",
         default=False,
-        help="Export models for half precision, faster on some Nvidia cards.",
+        help="Export models for half precision, smaller and faster on most GPUs.",
     )
     parser.add_argument(
         "--opset",
@@ -477,16 +478,11 @@ def main() -> int:
     logger.info("CLI arguments: %s", args)
 
     ctx = ConversionContext.from_environ()
-    ctx.half = args.half
+    ctx.half = args.half or "onnx-fp16" in ctx.optimizations
     ctx.opset = args.opset
     ctx.token = args.token
     logger.info("converting models in %s using %s", ctx.model_path, ctx.training_device)
 
-    if ctx.half and ctx.training_device != "cuda":
-        raise ValueError(
-            "half precision model export is only supported on GPUs with CUDA"
-        )
-
     if not path.exists(ctx.model_path):
         logger.info("model path does not existing, creating: %s", ctx.model_path)
         makedirs(ctx.model_path)

api/onnx_web/convert/diffusion/diffusers.py

@@ -27,6 +27,7 @@
 from onnxruntime.transformers.float16 import convert_float_to_float16
 from torch.onnx import export
 
+from ...constants import ONNX_MODEL, ONNX_WEIGHTS
 from ...diffusers.load import optimize_pipeline
 from ...diffusers.pipeline_onnx_stable_diffusion_upscale import (
     OnnxStableDiffusionUpscalePipeline,
@@ -79,7 +80,7 @@ def onnx_export(
             f"{output_file}",
             save_as_external_data=external_data,
             all_tensors_to_one_file=True,
-            location="weights.pb",
+            location=ONNX_WEIGHTS,
         )
 
 
@@ -144,7 +145,7 @@ def convert_diffusion_diffusers(
             None,  # output attentions
             torch.tensor(True).to(device=ctx.training_device, dtype=torch.bool),
         ),
-        output_path=output_path / "text_encoder" / "model.onnx",
+        output_path=output_path / "text_encoder" / ONNX_MODEL,
         ordered_input_names=["input_ids"],
         output_names=["last_hidden_state", "pooler_output", "hidden_states"],
         dynamic_axes={
@@ -169,7 +170,7 @@ def convert_diffusion_diffusers(
 
     unet_in_channels = pipeline.unet.config.in_channels
     unet_sample_size = pipeline.unet.config.sample_size
-    unet_path = output_path / "unet" / "model.onnx"
+    unet_path = output_path / "unet" / ONNX_MODEL
     onnx_export(
         pipeline.unet,
         model_args=(
@@ -207,7 +208,7 @@ def convert_diffusion_diffusers(
         unet_model_path,
         save_as_external_data=True,
         all_tensors_to_one_file=True,
-        location="weights.pb",
+        location=ONNX_WEIGHTS,
         convert_attribute=False,
     )
     del pipeline.unet
@@ -233,7 +234,7 @@ def convert_diffusion_diffusers(
                 ).to(device=ctx.training_device, dtype=dtype),
                 False,
             ),
-            output_path=output_path / "vae" / "model.onnx",
+            output_path=output_path / "vae" / ONNX_MODEL,
             ordered_input_names=["latent_sample", "return_dict"],
             output_names=["sample"],
             dynamic_axes={
@@ -259,14 +260,14 @@ def convert_diffusion_diffusers(
                 ),
                 False,
             ),
-            output_path=output_path / "vae_encoder" / "model.onnx",
+            output_path=output_path / "vae_encoder" / ONNX_MODEL,
             ordered_input_names=["sample", "return_dict"],
             output_names=["latent_sample"],
             dynamic_axes={
                 "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
             },
             opset=ctx.opset,
-            half=ctx.half,
+            half=False,  # https://github.com/ssube/onnx-web/issues/290
         )
 
         # VAE DECODER
@@ -282,7 +283,7 @@ def convert_diffusion_diffusers(
                 ).to(device=ctx.training_device, dtype=dtype),
                 False,
             ),
-            output_path=output_path / "vae_decoder" / "model.onnx",
+            output_path=output_path / "vae_decoder" / ONNX_MODEL,
             ordered_input_names=["latent_sample", "return_dict"],
             output_names=["sample"],
             dynamic_axes={