From fd0cf2790b23668fa48eef9a28d28b9e618b1439 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 22 Nov 2023 23:41:00 +0200
Subject: [PATCH] attempt at memory management and remove unnecessary req

---
 comfyui_svd.py   | 6 ++++++
 requirements.txt | 1 -
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/comfyui_svd.py b/comfyui_svd.py
index 8c4e80d..7939d7b 100644
--- a/comfyui_svd.py
+++ b/comfyui_svd.py
@@ -3,6 +3,7 @@ from omegaconf import OmegaConf
 import math
 import torch
 import importlib
+import comfy.model_management
 
 def get_obj_from_str(string, reload=False, invalidate_cache=True):
     module, cls = string.rsplit(".", 1)
@@ -127,6 +128,11 @@ class SVDimg2vid:
 
     def generate(self, image, version, num_frames, num_steps, fps_id, motion_bucket_id, cond_aug, seed, decoding_t, lowvram_mode):
        
+        #since this is so memory intensive, try to get everything free
+        comfy.model_management.cleanup_models()
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
         device: str = "cuda"
 
         model_config = f"custom_nodes/ComfyUI-SVD/svd/configs/{version}.yaml"
diff --git a/requirements.txt b/requirements.txt
index 366333a..3254a24 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,5 @@
 einops
 omegaconf
-math
 pytorch-lightning==2.0.1
 open-clip-torch>=2.20.0
 sgm @ git+https://github.com/Stability-AI/generative-models
\ No newline at end of file