feat(preview): retain the latest denoise frame server-side and replay it on reconnect

invokeai/app/api/dependencies.py

@@ -46,6 +46,7 @@
 from invokeai.app.services.names.names_default import SimpleNameService
 from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
 from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
+from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews
 from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite
 from invokeai.app.services.session_processor.session_processor_default import (
     DefaultSessionProcessor,
@@ -231,6 +232,7 @@ def initialize(
             events=events,
             image_files=image_files,
             image_moves=image_moves,
+            progress_previews=MemoryProgressPreviews(),
             image_records=image_records,
             images=images,
             invocation_cache=invocation_cache,

invokeai/app/api/routers/session_queue.py

@@ -10,6 +10,7 @@
 from invokeai.app.api.dependencies import ApiDependencies
 from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
 from invokeai.app.invocations.fields import ImageField, VideoField
+from invokeai.app.services.progress_previews.progress_previews_common import ProgressPreviewDTO
 from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus
 from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
@@ -635,6 +636,28 @@ def get_current_queue_item(
         raise HTTPException(status_code=500, detail=f"Unexpected error while getting current queue item: {e}")
 
 
+@session_queue_router.get(
+    "/{queue_id}/previews",
+    operation_id="get_progress_previews",
+    responses={
+        200: {"model": list[ProgressPreviewDTO]},
+    },
+)
+def get_progress_previews(
+    current_user: CurrentUserOrDefault,
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+) -> list[ProgressPreviewDTO]:
+    """The latest denoising preview frame of each of the caller's running queue items: the same
+    payloads as the `invocation_progress` socket events, with their revisions. A client whose socket
+    was dropped, or whose tab was hidden, reconciles from this instead of waiting for the next step.
+    Owner-scoped: progress is personal UI, so even admins see only their own."""
+    try:
+        previews = ApiDependencies.invoker.services.progress_previews.list_for_user(current_user.user_id, queue_id)
+        return [ProgressPreviewDTO.from_event(event) for event in previews]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error while getting progress previews: {e}")
+
+
 @session_queue_router.get(
     "/{queue_id}/next",
     operation_id="get_next_queue_item",

invokeai/app/api/sockets.py

@@ -831,6 +831,26 @@ async def _handle_sub_queue(self, sid: str, data: Any) -> None:
             f"Socket {sid} (user_id: {user_id}, is_admin: {is_admin}) subscribed to queue {queue_id} and user room {user_room}"
         )
 
+        await self._replay_progress_previews(sid, user_id, queue_id)
+
+    async def _replay_progress_previews(self, sid: str, user_id: str, queue_id: str) -> None:
+        """Send this socket the latest preview frame of each running queue item the user owns.
+
+        A reconnecting tab otherwise shows nothing until the next denoising step. The frames carry
+        their original revisions, so a socket that never lost them drops the duplicates. Sent to
+        this socket only: the user's other sockets already have them.
+        """
+        from invokeai.app.api.dependencies import ApiDependencies
+
+        invoker = getattr(ApiDependencies, "invoker", None)
+        if invoker is None:
+            return
+
+        for event in invoker.services.progress_previews.list_for_user(user_id, queue_id):
+            await self._sio.emit(
+                event=InvocationProgressEvent.__event_name__, data=event.model_dump(mode="json"), to=sid
+            )
+
     async def _handle_unsub_queue(self, sid: str, data: Any) -> None:
         await self._sio.leave_room(sid, QueueSubscriptionEvent(**data).queue_id)
 

invokeai/app/services/events/events_base.py

@@ -80,11 +80,14 @@ def emit_invocation_progress(
         message: str,
         percentage: float | None = None,
         image: "ProgressImage | None" = None,
-    ) -> None:
-        """Emitted at periodically during an invocation"""
-        self.dispatch(
-            InvocationProgressEvent.build(queue_item, invocation.get_event_invocation(), message, percentage, image)
+        revision: int | None = None,
+    ) -> InvocationProgressEvent:
+        """Emitted periodically during an invocation. Returns the event so the caller can retain it."""
+        event = InvocationProgressEvent.build(
+            queue_item, invocation.get_event_invocation(), message, percentage, image, revision
         )
+        self.dispatch(event)
+        return event
 
     def emit_invocation_complete(
         self, queue_item: "SessionQueueItem", invocation: "BaseInvocation", output: "BaseInvocationOutput"

invokeai/app/services/events/events_common.py

@@ -150,6 +150,14 @@ class InvocationProgressEvent(InvocationEventBase):
         default=None,
         description="The device processing this session, e.g. 'cuda:1' (set only when running on a GPU)",
     )
+    revision: int | None = Field(
+        default=None,
+        description=(
+            "Monotonic per queue item and session, set on image-bearing frames. A frame at or below a "
+            "revision already shown is stale: the same frame may reach a client twice (live stream and "
+            "reconnect replay) and must never move the preview backwards."
+        ),
+    )
 
     @classmethod
     def build(
@@ -159,6 +167,7 @@ def build(
         message: str,
         percentage: float | None = None,
         image: ProgressImage | None = None,
+        revision: int | None = None,
     ) -> "InvocationProgressEvent":
         # Report the GPU executing the session. Prefer the queue item's persisted device: the
         # thread-local session device is temporarily re-pinned to a borrowed idle GPU during
@@ -191,6 +200,7 @@ def build(
             image=image,
             message=message,
             device=device,
+            revision=revision,
         )
 
 

invokeai/app/services/invocation_services.py

@@ -41,6 +41,7 @@
     )
     from invokeai.app.services.model_relationships.model_relationships_base import ModelRelationshipsServiceABC
     from invokeai.app.services.names.names_base import NameServiceBase
+    from invokeai.app.services.progress_previews.progress_previews_base import ProgressPreviewsBase
     from invokeai.app.services.project_records.project_records_base import ProjectRecordsStorageBase
     from invokeai.app.services.session_processor.session_processor_base import SessionProcessorBase
     from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase
@@ -54,6 +55,12 @@
     from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
 
 
+def _default_progress_previews() -> "ProgressPreviewsBase":
+    from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews
+
+    return MemoryProgressPreviews()
+
+
 class InvocationServices:
     """Services that can be used by invocations"""
 
@@ -101,6 +108,7 @@ def __init__(
         image_index_records: "ImageIndexRecordsBase",
         image_index: "ImageIndexServiceBase",
         image_moves: "ImageMoveService | None" = None,
+        progress_previews: "ProgressPreviewsBase | None" = None,
     ):
         self.board_images = board_images
         self.board_image_records = board_image_records
@@ -122,6 +130,11 @@ def __init__(
         self.performance_statistics = performance_statistics
         self.session_queue = session_queue
         self.image_moves = image_moves
+        # Pure in-memory state with no dependencies, so callers (including the many test
+        # constructors) may leave it out.
+        self.progress_previews: "ProgressPreviewsBase" = (
+            progress_previews if progress_previews is not None else _default_progress_previews()
+        )
         self.session_processor = session_processor
         self.invocation_cache = invocation_cache
         self.names = names

invokeai/app/services/progress_previews/progress_previews_base.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+
+from invokeai.app.services.events.events_common import InvocationProgressEvent
+
+
+class ProgressPreviewsBase(ABC):
+    """The latest denoising preview per running queue item: best-effort and disposable.
+
+    Progress frames travel over the socket as base64 data URLs and otherwise exist nowhere. A
+    client that reconnects, or a tab that was hidden while its socket was dropped, needs to ask for
+    the last frame rather than wait for the next step. This store keeps the last *emitted* progress
+    event per queue item so it can be replayed on subscribe or fetched over HTTP.
+
+    Each frame carries a `revision` that is monotonic per queue item and session; the client drops
+    any frame at or below the revision it has already shown, so a replay can never move the preview
+    backwards. A process restart empties the store, which is the intended epoch semantics.
+
+    Image-bearing frames are also throttled here (`reserve_revision` returns None for a frame that
+    should be dropped) so that a fast GPU does not turn every denoising step into a JPEG encode and
+    a socket emit.
+    """
+
+    @abstractmethod
+    def reserve_revision(self, item_id: int, session_id: str, percentage: float | None) -> int | None:
+        """Allocate the revision for a new image-bearing frame, or None if the frame should be dropped.
+
+        Args:
+            item_id: The queue item the frame belongs to.
+            session_id: The session producing it; a new session on the same item starts a new epoch.
+            percentage: The frame's progress, if determinate. A final frame (1.0) is never dropped.
+        """
+        pass
+
+    @abstractmethod
+    def record(self, event: InvocationProgressEvent) -> None:
+        """Remember an emitted image-bearing progress event as the item's latest preview."""
+        pass
+
+    @abstractmethod
+    def get(self, item_id: int) -> InvocationProgressEvent | None:
+        """The latest preview for a queue item, if any."""
+        pass
+
+    @abstractmethod
+    def list_for_user(self, user_id: str, queue_id: str | None = None) -> list[InvocationProgressEvent]:
+        """The latest previews for every running queue item the user owns, oldest item first."""
+        pass
+
+    @abstractmethod
+    def clear_node(self, item_id: int, invocation_id: str) -> None:
+        """Forget the preview if this node produced it: the node is done, so a replay of its frame
+        would show a denoise that has already finished and mark the node as running again."""
+        pass
+
+    @abstractmethod
+    def clear(self, item_id: int) -> None:
+        """Forget a queue item's preview: its session leg ended, whatever the outcome. The item's
+        revision counter is kept so a later leg of the same session cannot reissue a revision a
+        client has already seen."""
+        pass
+
+    @abstractmethod
+    def clear_all(self) -> None:
+        pass

invokeai/app/services/progress_previews/progress_previews_common.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, Field
+
+from invokeai.app.services.events.events_common import InvocationProgressEvent
+from invokeai.app.services.session_processor.session_processor_common import ProgressImage
+
+
+class ProgressPreviewDTO(BaseModel):
+    """The latest preview frame of a running queue item, as served by `GET /queue/{queue_id}/previews`.
+
+    The fields of the `invocation_progress` socket event a preview consumer reads, without the
+    serialized invocation. A separate model rather than the event itself: exposing the event as a
+    response model would give it a second OpenAPI schema variant and change the generated client
+    types for every socket consumer.
+    """
+
+    queue_id: str = Field(description="The ID of the queue")
+    item_id: int = Field(description="The ID of the queue item")
+    session_id: str = Field(description="The ID of the session (aka graph execution state)")
+    invocation_source_id: str = Field(description="The ID of the prepared invocation's source node")
+    revision: int | None = Field(description="Monotonic per queue item and session; see InvocationProgressEvent")
+    message: str = Field(description="A message to display")
+    percentage: float | None = Field(description="The percentage of the progress, or null if indeterminate")
+    image: ProgressImage | None = Field(description="The latest denoising preview")
+
+    @classmethod
+    def from_event(cls, event: InvocationProgressEvent) -> "ProgressPreviewDTO":
+        return cls(
+            queue_id=event.queue_id,
+            item_id=event.item_id,
+            session_id=event.session_id,
+            invocation_source_id=event.invocation_source_id,
+            revision=event.revision,
+            message=event.message,
+            percentage=event.percentage,
+            image=event.image,
+        )

invokeai/app/services/progress_previews/progress_previews_default.py

@@ -0,0 +1,104 @@
+import time
+from dataclasses import dataclass
+from threading import Lock
+
+from invokeai.app.services.events.events_common import InvocationProgressEvent
+from invokeai.app.services.progress_previews.progress_previews_base import ProgressPreviewsBase
+
+# The floor between two image-bearing frames of one queue item. Ten frames a second is well past
+# what a viewer can follow; above it every step became a JPEG encode on the generation thread and a
+# socket emit for nothing. A final frame (percentage 1.0) and indeterminate frames always pass.
+PROGRESS_IMAGE_MIN_INTERVAL_S = 0.1
+
+# Revision counters outlive their frames (see `clear`). They are tiny, but bounded all the same:
+# past this many frameless entries the oldest are forgotten.
+COUNTER_RETENTION_LIMIT = 4096
+
+
+@dataclass
+class _ItemPreview:
+    session_id: str
+    revision: int = 0
+    last_emitted_at: float | None = None
+    event: InvocationProgressEvent | None = None
+    invocation_id: str | None = None
+
+
+class MemoryProgressPreviews(ProgressPreviewsBase):
+    """In-memory previews: one entry per queue item, shared by every worker thread."""
+
+    def __init__(self, min_interval_s: float = PROGRESS_IMAGE_MIN_INTERVAL_S, clock=time.monotonic) -> None:
+        self._min_interval_s = min_interval_s
+        self._clock = clock
+        self._lock = Lock()
+        self._items: dict[int, _ItemPreview] = {}
+
+    def reserve_revision(self, item_id: int, session_id: str, percentage: float | None) -> int | None:
+        now = self._clock()
+        with self._lock:
+            item = self._items.get(item_id)
+            if item is None or item.session_id != session_id:
+                item = _ItemPreview(session_id=session_id)
+                self._items[item_id] = item
+            is_final = percentage is None or percentage >= 1.0
+            if not is_final and item.last_emitted_at is not None and now - item.last_emitted_at < self._min_interval_s:
+                return None
+            item.revision += 1
+            item.last_emitted_at = now
+            return item.revision
+
+    def record(self, event: InvocationProgressEvent) -> None:
+        with self._lock:
+            item = self._items.get(event.item_id)
+            if item is None or item.session_id != event.session_id:
+                # Recorded without a reservation (or after a session change): adopt it as-is.
+                item = _ItemPreview(session_id=event.session_id, revision=event.revision or 0)
+                self._items[event.item_id] = item
+            item.event = event
+            item.invocation_id = event.invocation.id
+
+    def get(self, item_id: int) -> InvocationProgressEvent | None:
+        with self._lock:
+            item = self._items.get(item_id)
+            return item.event if item is not None else None
+
+    def list_for_user(self, user_id: str, queue_id: str | None = None) -> list[InvocationProgressEvent]:
+        with self._lock:
+            events = [item.event for item in self._items.values() if item.event is not None]
+        return [
+            event
+            for event in sorted(events, key=lambda event: event.item_id)
+            if event.user_id == user_id and (queue_id is None or event.queue_id == queue_id)
+        ]
+
+    def clear_node(self, item_id: int, invocation_id: str) -> None:
+        with self._lock:
+            item = self._items.get(item_id)
+            if item is not None and item.invocation_id == invocation_id:
+                item.event = None
+                item.invocation_id = None
+
+    def clear(self, item_id: int) -> None:
+        with self._lock:
+            item = self._items.get(item_id)
+            if item is None:
+                return
+            # The frame goes; the counter stays so a later leg of the same item and session (a
+            # workflow-call parent resuming after its child) keeps issuing higher revisions than
+            # any a client already holds. The throttle does not carry over.
+            item.event = None
+            item.invocation_id = None
+            item.last_emitted_at = None
+            self._evict_stale_counters()
+
+    def clear_all(self) -> None:
+        with self._lock:
+            self._items.clear()
+
+    def _evict_stale_counters(self) -> None:
+        """Caller holds the lock. Drop the oldest frameless entries past the retention cap."""
+        excess = len(self._items) - COUNTER_RETENTION_LIMIT
+        if excess <= 0:
+            return
+        for item_id in [item_id for item_id, item in self._items.items() if item.event is None][:excess]:
+            del self._items[item_id]