1
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 15:04:50 +08:00

Move code to empty gpu cache to model_management.py

This commit is contained in:
comfyanonymous
2023-04-15 11:19:07 -04:00
parent f4c689ea89
commit deb2b93e79
2 changed files with 11 additions and 7 deletions

View File

@@ -10,7 +10,7 @@ import gc
import torch
import nodes
from model_management import xpu_available
import comfy.model_management
def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
valid_inputs = class_def.INPUT_TYPES()
@@ -204,12 +204,7 @@ class PromptExecutor:
self.server.send_sync("executing", { "node": None }, self.server.client_id)
gc.collect()
if torch.cuda.is_available():
if torch.version.cuda: #This seems to make things worse on ROCm so I only do it for cuda
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif xpu_available:
torch.xpu.empty_cache()
comfy.model_management.soft_empty_cache()
def validate_inputs(prompt, item):