mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
29325a7214
Around the time we (I) implemented pydantic events, I noticed a short pause between progress images every 4 or 5 steps when generating with SDXL. It didn't happen with SD1.5, but I did notice that with SD1.5, we'd get 4 or 5 progress events simultaneously. I'd expect one event every ~25ms, matching my it/s with SD1.5. Mysterious! Digging in, I found an issue is related to our use of a synchronous queue for events. When the event queue is empty, we must call `asyncio.sleep` before checking again. We were sleeping for 100ms. Said another way, every time we clear the event queue, we have to wait 100ms before another event can be dispatched, even if it is put on the queue immediately after we start waiting. In practice, this means our events get buffered into batches, dispatched once every 100ms. This explains why I was getting batches of 4 or 5 SD1.5 progress events at once, but not the intermittent SDXL delay. But this 100ms wait has another effect when the events are put on the queue in intervals that don't perfectly line up with the 100ms wait. This is most noticeable when the time between events is >100ms, and can add up to 100ms delay before the event is dispatched. For example, say the queue is empty and we start a 100ms wait. Then, immediately after - like 0.01ms later - we push an event on to the queue. We still need to wait another 99.9ms before that event will be dispatched. That's the SDXL delay. The easy fix is to reduce the sleep to something like 0.01 seconds, but this feels kinda dirty. Can't we just wait on the queue and dispatch every event immediately? Not with the normal synchronous queue - but we can with `asyncio.Queue`. I switched the events queue to use `asyncio.Queue` (as seen in this commit), which lets us asynchronous wait on the queue in a loop. Unfortunately, I ran into another issue - events now felt like their timing was inconsistent, but in a different way than with the 100ms sleep. The time between pushing events on the queue and dispatching them was not consistently ~0ms as I'd expect - it was highly variable from ~0ms up to ~100ms. This is resolved by passing the asyncio loop directly into the events service and using its methods to create the task and interact with the queue. I don't fully understand why this resolved the issue, because either way we are interacting with the same event loop (as shown by `asyncio.get_running_loop()`). I suppose there's some scheduling magic happening.
39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
import asyncio
|
|
import threading
|
|
|
|
from fastapi_events.dispatcher import dispatch
|
|
|
|
from invokeai.app.services.events.events_base import EventServiceBase
|
|
from invokeai.app.services.events.events_common import EventBase
|
|
|
|
|
|
class FastAPIEventService(EventServiceBase):
|
|
def __init__(self, event_handler_id: int, loop: asyncio.AbstractEventLoop) -> None:
|
|
self.event_handler_id = event_handler_id
|
|
self._queue = asyncio.Queue[EventBase | None]()
|
|
self._stop_event = threading.Event()
|
|
self._loop = loop
|
|
self._loop.create_task(self._dispatch_from_queue(stop_event=self._stop_event))
|
|
|
|
super().__init__()
|
|
|
|
def stop(self, *args, **kwargs):
|
|
self._stop_event.set()
|
|
self._loop.call_soon_threadsafe(self._queue.put_nowait, None)
|
|
|
|
def dispatch(self, event: EventBase) -> None:
|
|
self._loop.call_soon_threadsafe(self._queue.put_nowait, event)
|
|
|
|
async def _dispatch_from_queue(self, stop_event: threading.Event):
|
|
"""Get events on from the queue and dispatch them, from the correct thread"""
|
|
while not stop_event.is_set():
|
|
try:
|
|
event = await self._queue.get()
|
|
if not event: # Probably stopping
|
|
continue
|
|
# Leave the payloads as live pydantic models
|
|
dispatch(event, middleware_id=self.event_handler_id, payload_schema_dump=False)
|
|
|
|
except asyncio.CancelledError as e:
|
|
raise e # Raise a proper error
|