Fix hang when images finish faster than they can be saved to document history

* usually only happened when comfy execution was cached
* _save_results is called from inside asyncio eventloop

Author: Acly

ai_diffusion/eventloop.py

@@ -48,11 +48,8 @@ def stop():
         pass
 
 
-def wait(future: asyncio.Future):
-    event = threading.Event()
-    future.add_done_callback(lambda _: event.set())
-    event.wait()
-    return future.result()
+def run_until_complete(future: asyncio.Future):
+    return _loop.run_until_complete(future)
 
 
 async def wait_until(condition: Callable[[], bool], iterations=10, no_error=False):

ai_diffusion/persistence.py

@@ -134,7 +134,7 @@ def __init__(self, model: Model):
     def __del__(self):
         try:
             if self._image_task is not None and not self._image_task.done():
-                eventloop.wait(self._image_task)
+                eventloop.run_until_complete(self._image_task)
         except Exception as e:
             log.warning(f"Persistence: failed to wait for image task completion: {e}")
 
@@ -229,12 +229,13 @@ def _save_results(self, job: Job):
         if job.kind in [JobKind.diffusion, JobKind.animation] and len(job.results) > 0:
             slot = self._slot_index
             self._slot_index += 1
+            self._image_task = eventloop.run(self._save_result_images(job, slot, self._image_task))
 
-            if self._image_task is not None and not self._image_task.done():
-                eventloop.wait(self._image_task)
-            self._image_task = eventloop.run(self._save_result_images(job, slot))
-
-    async def _save_result_images(self, job: Job, slot: int):
+    async def _save_result_images(
+        self, job: Job, slot: int, prev_task: asyncio.Future | None = None
+    ):
+        if prev_task is not None:
+            await prev_task
         if settings.multi_threading:
             loop = asyncio.get_running_loop()
             image_data, image_offsets = await loop.run_in_executor(