From f7235cf82c3e13c836d266b95c9dcdb4653ec7d5 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Fri, 21 Apr 2023 20:59:14 +0530
Subject: [PATCH] Keep the task alive during step callbacks. Thanks Madrang

---
 ui/easydiffusion/task_manager.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ui/easydiffusion/task_manager.py b/ui/easydiffusion/task_manager.py
index 31fdaa6f..91adc04b 100644
--- a/ui/easydiffusion/task_manager.py
+++ b/ui/easydiffusion/task_manager.py
@@ -317,6 +317,9 @@ def thread_render(device):
             def step_callback():
                 global current_state_error
 
+                task_cache.keep(id(task), TASK_TTL)
+                session_cache.keep(task.task_data.session_id, TASK_TTL)
+
                 if (
                     isinstance(current_state_error, SystemExit)
                     or isinstance(current_state_error, StopAsyncIteration)