From e0b788fa7c3fe7e4091aa813e26dbd07ff5fa3b8 Mon Sep 17 00:00:00 2001
From: yusing <yusing.wys@gmail.com>
Date: Tue, 18 Jun 2024 04:36:12 +0800
Subject: [PATCH] removed unload_before and unload_after option

---
 pyproject.toml      |  2 +-
 tensorrt_convert.py | 13 +------------
 2 files changed, 2 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 0d27a6d..bf651c8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_tensorrt"
 description = "TensorRT Node for ComfyUI\nThis node enables the best performance on NVIDIA RTX™ Graphics Cards  (GPUs) for Stable Diffusion by leveraging NVIDIA TensorRT."
-version = "0.1.1"
+version = "0.1.2"
 license = "LICENSE"
 dependencies = [
   "tensorrt>=10.0.1",
diff --git a/tensorrt_convert.py b/tensorrt_convert.py
index d4e4cde..8366b6c 100644
--- a/tensorrt_convert.py
+++ b/tensorrt_convert.py
@@ -153,15 +153,11 @@ def _convert(
         context_max,
         num_video_frames,
         is_static: bool,
-        unload_before: bool = True,
-        unload_after: bool = True,
     ):
         output_onnx = os.path.normpath(
             os.path.join(self.temp_dir, str(time.time()), "model.onnx")
         )
 
-        if unload_before:
-            comfy.model_management.unload_all_models()
         comfy.model_management.load_models_gpu([model], force_patch_weights=True)
         unet = model.model.diffusion_model
 
@@ -286,8 +282,7 @@ def forward(self, x, timesteps, context, y=None):
             dynamic_axes=dynamic_axes,
         )
 
-        if unload_after:
-            comfy.model_management.unload_all_models()
+        comfy.model_management.unload_all_models()
         comfy.model_management.soft_empty_cache()
 
         # TRT conversion starts here
@@ -606,8 +601,6 @@ def INPUT_TYPES(s):
                         "step": 1,
                     },
                 ),
-                "unload_before": ("BOOLEAN", {"default": True}),
-                "unload_after": ("BOOLEAN", {"default": True}),
             },
         }
 
@@ -620,8 +613,6 @@ def convert(
         width_opt,
         context_opt,
         num_video_frames,
-        unload_before: bool = True,
-        unload_after: bool = True,
     ):
         return super()._convert(
             model,
@@ -640,8 +631,6 @@ def convert(
             context_opt,
             num_video_frames,
             is_static=True,
-            unload_before=unload_before,
-            unload_after=unload_after,
         )