removed unload_before and unload_after option

comfyanonymous · Jun 17, 2024 · e0b788f · e0b788f
1 parent 505b204
commit e0b788f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 13 deletions.
diff --git 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
@@ -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,
         )