fix(api): run GC after changing pipeline (#58)

api/onnx_web/pipeline.py

@@ -5,10 +5,10 @@
     OnnxStableDiffusionImg2ImgPipeline,
     OnnxStableDiffusionInpaintPipeline,
 )
-from os import environ
 from PIL import Image, ImageChops
 from typing import Any
 
+import gc
 import numpy as np
 
 from .image import (
@@ -19,6 +19,7 @@
     UpscaleParams,
 )
 from .utils import (
+    is_debug,
     safer_join,
     BaseParams,
     Border,
@@ -70,6 +71,9 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu
             model, subfolder='scheduler')
         last_pipeline_scheduler = scheduler
 
+    print('running garbage collection during pipeline change')
+    gc.collect()
+
     return pipe
 
 
@@ -167,7 +171,7 @@ def run_inpaint_pipeline(
         noise_source=noise_source,
         mask_filter=mask_filter)
 
-    if environ.get('DEBUG') is not None:
+    if is_debug():
         source_image.save(safer_join(ctx.output_path, 'last-source.png'))
         mask_image.save(safer_join(ctx.output_path, 'last-mask.png'))
         noise_image.save(safer_join(ctx.output_path, 'last-noise.png'))

api/onnx_web/serve.py

@@ -45,6 +45,7 @@
     UpscaleParams,
 )
 from .utils import (
+    is_debug,
     get_and_clamp_float,
     get_and_clamp_int,
     get_from_list,
@@ -57,6 +58,7 @@
     Size,
 )
 
+import gc
 import json
 import numpy as np
 
@@ -259,6 +261,9 @@ def load_params(context: ServerContext):
 CORS(app, origins=context.cors_origin)
 executor = Executor(app)
 
+if is_debug():
+    gc.set_debug(gc.DEBUG_STATS)
+
 
 # TODO: these two use context
 

api/onnx_web/utils.py

@@ -99,6 +99,10 @@ def tojson(self) -> Dict[str, int]:
         }
 
 
+def is_debug() -> bool:
+    return environ.get('DEBUG') is not None
+
+
 def get_and_clamp_float(args: Any, key: str, default_value: float, max_value: float, min_value=0.0) -> float:
     return min(max(float(args.get(key, default_value)), min_value), max_value)