fix(api): unload old model before loading next one

ssube · Jan 22, 2023 · 9e26ee5 · 9e26ee5
1 parent abc1ae5
commit 9e26ee5
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/api/onnx_web/pipeline.py b/api/onnx_web/pipeline.py
@@ -10,6 +10,7 @@
 
 import gc
 import numpy as np
+import torch
 
 from .image import (
     expand_image,
@@ -54,6 +55,12 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu
         print('reusing existing pipeline')
         pipe = last_pipeline_instance
     else:
+        print('unloading previous pipeline')
+        last_pipeline_instance = None
+        last_pipeline_scheduler = None
+        gc.collect()
+        torch.cuda.empty_cache()
+
         print('loading different pipeline')
         pipe = pipeline.from_pretrained(
             model,