From 4a3bb9734218ede79f632b208a54da510d01ea00 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Thu, 19 Jan 2023 19:46:36 -0600
Subject: [PATCH] fix(api): run GC after changing pipeline (#58)

---
 api/onnx_web/pipeline.py | 8 ++++++--
 api/onnx_web/serve.py    | 5 +++++
 api/onnx_web/utils.py    | 4 ++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/api/onnx_web/pipeline.py b/api/onnx_web/pipeline.py
index 481758f8..7bf2c384 100644
--- a/api/onnx_web/pipeline.py
+++ b/api/onnx_web/pipeline.py
@@ -5,10 +5,10 @@ from diffusers import (
     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 @@ from .upscale import (
     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'))
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
index 5a630f3c..0c4b75a8 100644
--- a/api/onnx_web/serve.py
+++ b/api/onnx_web/serve.py
@@ -45,6 +45,7 @@ from .upscale import (
     UpscaleParams,
 )
 from .utils import (
+    is_debug,
     get_and_clamp_float,
     get_and_clamp_int,
     get_from_list,
@@ -57,6 +58,7 @@ from .utils import (
     Size,
 )
 
+import gc
 import json
 import numpy as np
 
@@ -259,6 +261,9 @@ app.config['EXECUTOR_PROPAGATE_EXCEPTIONS'] = True
 CORS(app, origins=context.cors_origin)
 executor = Executor(app)
 
+if is_debug():
+    gc.set_debug(gc.DEBUG_STATS)
+
 
 # TODO: these two use context
 
diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py
index 35088147..7903731d 100644
--- a/api/onnx_web/utils.py
+++ b/api/onnx_web/utils.py
@@ -99,6 +99,10 @@ class Size:
         }
 
 
+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)