Do a garbage collect after the interval even if nothing is running.

This commit is contained in:
comfyanonymous 2023-11-30 15:22:32 -05:00
parent 7f469203b7
commit 6b769bca01
2 changed files with 33 additions and 18 deletions

View File

@ -700,10 +700,12 @@ class PromptQueue:
self.server.queue_updated() self.server.queue_updated()
self.not_empty.notify() self.not_empty.notify()
def get(self): def get(self, timeout=None):
with self.not_empty: with self.not_empty:
while len(self.queue) == 0: while len(self.queue) == 0:
self.not_empty.wait() self.not_empty.wait(timeout=timeout)
if timeout is not None and len(self.queue) == 0:
return None
item = heapq.heappop(self.queue) item = heapq.heappop(self.queue)
i = self.task_counter i = self.task_counter
self.currently_running[i] = copy.deepcopy(item) self.currently_running[i] = copy.deepcopy(item)

45
main.py
View File

@ -89,23 +89,36 @@ def cuda_malloc_warning():
def prompt_worker(q, server): def prompt_worker(q, server):
e = execution.PromptExecutor(server) e = execution.PromptExecutor(server)
last_gc_collect = 0 last_gc_collect = 0
while True: need_gc = False
item, item_id = q.get() gc_collect_interval = 10.0
execution_start_time = time.perf_counter()
prompt_id = item[1]
e.execute(item[2], prompt_id, item[3], item[4])
q.task_done(item_id, e.outputs_ui)
if server.client_id is not None:
server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, server.client_id)
current_time = time.perf_counter() while True:
execution_time = current_time - execution_start_time timeout = None
print("Prompt executed in {:.2f} seconds".format(execution_time)) if need_gc:
if (current_time - last_gc_collect) > 10.0: timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
gc.collect()
comfy.model_management.soft_empty_cache() queue_item = q.get(timeout=timeout)
last_gc_collect = current_time if queue_item is not None:
print("gc collect") item, item_id = queue_item
execution_start_time = time.perf_counter()
prompt_id = item[1]
e.execute(item[2], prompt_id, item[3], item[4])
need_gc = True
q.task_done(item_id, e.outputs_ui)
if server.client_id is not None:
server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, server.client_id)
current_time = time.perf_counter()
execution_time = current_time - execution_start_time
print("Prompt executed in {:.2f} seconds".format(execution_time))
if need_gc:
current_time = time.perf_counter()
if (current_time - last_gc_collect) > gc_collect_interval:
gc.collect()
comfy.model_management.soft_empty_cache()
last_gc_collect = current_time
need_gc = False
async def run(server, address='', port=8188, verbose=True, call_on_start=None): async def run(server, address='', port=8188, verbose=True, call_on_start=None):
await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop()) await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop())