Skip to content

Commit

Permalink
Manually release memory footprint to avoid OOM that may occur after l…
Browse files Browse the repository at this point in the history
…ong-term use
  • Loading branch information
Your Name committed Mar 20, 2024
1 parent b2c2377 commit 858eed6
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions fooocusapi/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from fooocusapi.task_queue import QueueTask, TaskQueue, TaskOutputs

worker_queue: TaskQueue = None
last_model_name = None

def process_top():
import ldm_patched.modules.model_management
Expand Down Expand Up @@ -118,6 +119,17 @@ def yield_result(_, imgs, tasks, extension='png'):

try:
print(f"[Task Queue] Task queue start task, job_id={async_task.job_id}")
# clear memory
global last_model_name

if last_model_name is None:
last_model_name = async_task.req_param.base_model_name
if last_model_name != async_task.req_param.base_model_name:
model_management.cleanup_models() # key1
model_management.unload_all_models()
model_management.soft_empty_cache() # key2
last_model_name = async_task.req_param.base_model_name

worker_queue.start_task(async_task.job_id)

execution_start_time = time.perf_counter()
Expand Down

0 comments on commit 858eed6

Please sign in to comment.