From 65288f6c3cfb725173b5dc32eed47b04dbd320aa Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 6 Sep 2026 13:13:25 -0400 Subject: [PATCH 1/3] feat(preview): retain the latest denoise frame server-side and replay it on reconnect Progress frames travelled over the socket as base64 data URLs and existed nowhere else. A client whose socket was dropped while its tab was hidden, or a page that reloaded mid-run, had nothing to show until the next denoising step, and nothing at all for a run that finished while it was away. Phase 2 of the preview restructuring (additive, no schema removals): - New `progress_previews` service on InvocationServices: an in-memory, lock-guarded map of the last emitted `invocation_progress` event per running queue item. Cleared when the session ends, whatever the outcome. No files, no TTL sweeps: a process restart empties it, which is the intended epoch. - `InvocationProgressEvent.revision`: monotonic per queue item and session on image-bearing frames. The client drops any frame at or below the revision it has shown, so a replay can never move the preview backwards. - `signal_progress` throttles image-bearing frames to one per 100 ms per item (final and indeterminate frames always pass), before the JPEG encode. - `subscribe_queue` replays the user's stored frames to the subscribing socket. - `GET /queue/{queue_id}/previews` returns the caller's frames as `ProgressPreviewDTO` (a dedicated model: exposing the event as a response model gave it a second OpenAPI schema variant that made every field optional in the generated client types). - webv2: the queue coordinator fetches the snapshot when the tab becomes visible and after re-adopting runs on reload, feeding each payload through the socket handler's revision gate. The gate reopens when an item goes back to pending/waiting. - A frame is dropped from the store as soon as its node completes or errors, so a replay never shows a finished denoise as running; the revision counter survives `clear` (bounded) so a workflow-call parent resuming with the same item and session keeps issuing higher revisions than any a client holds; the processor-error path clears the item too. - The SD/SDXL preview extension now reports steps completed rather than the index of the step just run, so its last frame reads 100% like the FLUX and SD3 callbacks and the throttle never drops it as a mid-run frame. Legacy `web` schema regenerated (additive only); its typecheck passes. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DoHBJ9QPWfq1qMTfzdzuke --- invokeai/app/api/dependencies.py | 2 + invokeai/app/api/routers/session_queue.py | 23 +++ invokeai/app/api/sockets.py | 20 +++ invokeai/app/services/events/events_base.py | 11 +- invokeai/app/services/events/events_common.py | 10 ++ invokeai/app/services/invocation_services.py | 13 ++ .../services/progress_previews/__init__.py | 0 .../progress_previews_base.py | 64 ++++++++ .../progress_previews_common.py | 36 +++++ .../progress_previews_default.py | 104 +++++++++++++ .../session_processor_default.py | 9 ++ .../app/services/shared/invocation_context.py | 26 +++- .../stable_diffusion/extensions/preview.py | 5 +- .../frontend/web/src/services/api/schema.ts | 109 +++++++++++++ .../webv2/src/features/queue/core/types.ts | 19 +++ .../queue/data/httpRealtimeQueueBackend.ts | 2 + .../src/features/queue/data/serverApi.ts | 6 +- .../queue/runtime/coordinator.test.ts | 137 +++++++++++++++- .../src/features/queue/runtime/coordinator.ts | 51 +++++- .../routers/test_session_queue_previews.py | 99 ++++++++++++ .../services/progress_previews/__init__.py | 0 .../test_progress_previews.py | 147 ++++++++++++++++++ .../progress_previews/test_signal_progress.py | 72 +++++++++ .../test_for_loop_processor.py | 2 + .../services/test_for_loop_session_runner.py | 2 + .../test_session_processor_callbacks.py | 2 + .../app/services/workflow_call_test_utils.py | 6 + .../test_progress_preview_replay_socketio.py | 73 +++++++++ tests/test_nodes.py | 11 +- 29 files changed, 1044 insertions(+), 17 deletions(-) create mode 100644 invokeai/app/services/progress_previews/__init__.py create mode 100644 invokeai/app/services/progress_previews/progress_previews_base.py create mode 100644 invokeai/app/services/progress_previews/progress_previews_common.py create mode 100644 invokeai/app/services/progress_previews/progress_previews_default.py create mode 100644 tests/app/routers/test_session_queue_previews.py create mode 100644 tests/app/services/progress_previews/__init__.py create mode 100644 tests/app/services/progress_previews/test_progress_previews.py create mode 100644 tests/app/services/progress_previews/test_signal_progress.py create mode 100644 tests/app/test_progress_preview_replay_socketio.py diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index aad84485ed7..f172c88b1bb 100644 --- a/invokeai/app/api/dependencies.py +++ b/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, diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index 39c606e38ea..e502b5134cb 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/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", diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index 33fa634f37d..685d9774899 100644 --- a/invokeai/app/api/sockets.py +++ b/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) diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index c23c21cbb86..f4e3565d710 100644 --- a/invokeai/app/services/events/events_base.py +++ b/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" diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index 353c2e026c9..f1837b23faf 100644 --- a/invokeai/app/services/events/events_common.py +++ b/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, ) diff --git a/invokeai/app/services/invocation_services.py b/invokeai/app/services/invocation_services.py index a32ed3f32b2..0e3d64898e5 100644 --- a/invokeai/app/services/invocation_services.py +++ b/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 diff --git a/invokeai/app/services/progress_previews/__init__.py b/invokeai/app/services/progress_previews/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/invokeai/app/services/progress_previews/progress_previews_base.py b/invokeai/app/services/progress_previews/progress_previews_base.py new file mode 100644 index 00000000000..493ddb2ed90 --- /dev/null +++ b/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 diff --git a/invokeai/app/services/progress_previews/progress_previews_common.py b/invokeai/app/services/progress_previews/progress_previews_common.py new file mode 100644 index 00000000000..039978cae1f --- /dev/null +++ b/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, + ) diff --git a/invokeai/app/services/progress_previews/progress_previews_default.py b/invokeai/app/services/progress_previews/progress_previews_default.py new file mode 100644 index 00000000000..e79b182b678 --- /dev/null +++ b/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] diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index d969836c1e9..721d35092c7 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -413,6 +413,9 @@ def _on_after_run_session(self, queue_item: SessionQueueItem) -> None: f"On after run session: queue item {queue_item.item_id}, session {queue_item.session_id}" ) + # The item's preview frame is disposable: whatever the outcome, nothing may replay it now. + self._services.progress_previews.clear(queue_item.item_id) + # If we are profiling, stop the profiler and dump the profile & stats if self._profiler is not None: profile_path = self._profiler.stop() @@ -473,6 +476,9 @@ def _on_after_run_node( f"On after run node: queue item {queue_item.item_id}, session {queue_item.session_id}, node {invocation.id} ({invocation.get_type()})" ) + # The node's denoise is over: its last frame must not be replayed as if still running. + self._services.progress_previews.clear_node(queue_item.item_id, invocation.id) + # Send complete event on successful runs self._services.events.emit_invocation_complete(invocation=invocation, queue_item=queue_item, output=output) @@ -495,6 +501,7 @@ def _on_node_error( - Emits an invocation error event. - Run any callbacks registered for this event. """ + self._services.progress_previews.clear_node(queue_item.item_id, invocation.id) self._services.logger.debug( f"On node error: queue item {queue_item.item_id}, session {queue_item.session_id}, node {invocation.id} ({invocation.get_type()})" @@ -1068,6 +1075,8 @@ def _on_non_fatal_processor_error( self._invoker.services.logger.error(error_traceback) if queue_item is not None: + # This path bypasses the runner's after-session hook; the item's frame must not outlive it. + self._invoker.services.progress_previews.clear(queue_item.item_id) try: queue_item = self._invoker.services.session_queue.set_queue_item_session( queue_item.item_id, queue_item.session diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 35ee5e11914..fca191924cf 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -886,13 +886,33 @@ def signal_progress( original size. """ - self._services.events.emit_invocation_progress( - queue_item=self._data.queue_item, + queue_item = self._data.queue_item + + if image is None: + self._services.events.emit_invocation_progress( + queue_item=queue_item, + invocation=self._data.invocation, + message=message, + percentage=percentage, + ) + return + + # Image-bearing frames are revisioned and throttled by the preview store, and the emitted + # event is retained there so a reconnecting client can be given the latest frame instead + # of waiting for the next step. Throttled frames are dropped before the JPEG encode. + previews = self._services.progress_previews + revision = previews.reserve_revision(queue_item.item_id, queue_item.session_id, percentage) + if revision is None: + return + event = self._services.events.emit_invocation_progress( + queue_item=queue_item, invocation=self._data.invocation, message=message, percentage=percentage, - image=ProgressImage.build(image, image_size) if image else None, + image=ProgressImage.build(image, image_size), + revision=revision, ) + previews.record(event) class WildcardsInterface(InvocationContextInterface): diff --git a/invokeai/backend/stable_diffusion/extensions/preview.py b/invokeai/backend/stable_diffusion/extensions/preview.py index 6256f475945..ba5287b73a5 100644 --- a/invokeai/backend/stable_diffusion/extensions/preview.py +++ b/invokeai/backend/stable_diffusion/extensions/preview.py @@ -53,7 +53,10 @@ def step_preview(self, ctx: DenoiseContext): self.callback( PipelineIntermediateState( - step=ctx.step_index, + # Steps completed, not the index of the step just run: the last frame then reports + # 100%, as the FLUX/SD3 callbacks already do, and the progress-image throttle never + # drops it as a mid-run frame. + step=ctx.step_index + 1, order=ctx.scheduler.order, total_steps=len(ctx.inputs.timesteps), timestep=int(ctx.timestep), # TODO: is there any code which uses it? diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index d48bfdc9e33..b518139f96b 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -2884,6 +2884,29 @@ export type paths = { patch?: never; trace?: never; }; + "/api/v1/queue/{queue_id}/previews": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Progress Previews + * @description 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. + */ + get: operations["get_progress_previews"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/api/v1/queue/{queue_id}/next": { parameters: { query?: never; @@ -20363,6 +20386,12 @@ export type components = { * @default null */ device: string | null; + /** + * Revision + * @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. + * @default null + */ + revision: number | null; }; /** * InvocationStartedEvent @@ -33766,6 +33795,54 @@ export type components = { */ dataURL: string; }; + /** + * ProgressPreviewDTO + * @description 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. + */ + ProgressPreviewDTO: { + /** + * Queue Id + * @description The ID of the queue + */ + queue_id: string; + /** + * Item Id + * @description The ID of the queue item + */ + item_id: number; + /** + * Session Id + * @description The ID of the session (aka graph execution state) + */ + session_id: string; + /** + * Invocation Source Id + * @description The ID of the prepared invocation's source node + */ + invocation_source_id: string; + /** + * Revision + * @description Monotonic per queue item and session; see InvocationProgressEvent + */ + revision: number | null; + /** + * Message + * @description A message to display + */ + message: string; + /** + * Percentage + * @description The percentage of the progress, or null if indeterminate + */ + percentage: number | null; + /** @description The latest denoising preview */ + image: components["schemas"]["ProgressImage"] | null; + }; /** * ProjectBoardItemDTO * @description One visible item on a project's board. @@ -51530,6 +51607,38 @@ export interface operations { }; }; }; + get_progress_previews: { + parameters: { + query?: never; + header?: never; + path: { + /** @description The queue id to perform this operation on */ + queue_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ProgressPreviewDTO"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_next_queue_item: { parameters: { query?: { diff --git a/invokeai/frontend/webv2/src/features/queue/core/types.ts b/invokeai/frontend/webv2/src/features/queue/core/types.ts index 8196697b200..929d9a7938e 100644 --- a/invokeai/frontend/webv2/src/features/queue/core/types.ts +++ b/invokeai/frontend/webv2/src/features/queue/core/types.ts @@ -262,9 +262,28 @@ export interface QueueFeatureCommands { * The queue feature's backend seam. It owns both command transport and realtime * events so runtimes never assemble HTTP calls and socket subscriptions. */ +/** + * The backend's preview snapshot (`ProgressPreviewDTO`): the fields of an + * `invocation_progress` socket event a preview consumer reads. The coordinator + * feeds it through the same handler as the socket event, where the revision + * gate drops anything the live stream already delivered. + */ +export interface QueueProgressPreviewPayload { + queue_id: string; + item_id: number; + session_id: string; + invocation_source_id: string; + revision: number | null; + message: string; + percentage: number | null; + image: { width: number; height: number; dataURL: string } | null; +} + export interface QueueBackendPort extends QueueFeatureCommands { acknowledgeEnqueue?(projectId: string, sourceQueueItemId: string): Promise; getEnqueueReceipt?(projectId: string, sourceQueueItemId: string): Promise; + /** The latest preview frame of each of the user's running items (`GET queue/{id}/previews`). */ + readProgressPreviews?(signal?: AbortSignal): Promise; cancelQueueItems(itemIds: number[]): Promise; cancelQueueItemsByBatchIds(batchIds: string[]): Promise; enqueueGenerate(request: QueueEnqueueGenerateRequest): Promise; diff --git a/invokeai/frontend/webv2/src/features/queue/data/httpRealtimeQueueBackend.ts b/invokeai/frontend/webv2/src/features/queue/data/httpRealtimeQueueBackend.ts index 42681df21bd..88fd9645fc4 100644 --- a/invokeai/frontend/webv2/src/features/queue/data/httpRealtimeQueueBackend.ts +++ b/invokeai/frontend/webv2/src/features/queue/data/httpRealtimeQueueBackend.ts @@ -13,6 +13,7 @@ import { clearScopedQueue, getCurrentQueueItem, getNextQueueItem, + getProgressPreviews, getQueueItem, getQueueItemIds, getQueueItemsByIds, @@ -65,6 +66,7 @@ export const queueBackend: QueueBackendPort = { }, readItemIds: async (order, scope, signal) => mapQueueItemIdsDTO(await getQueueItemIds(order, scope, signal)), readItemsById: async (itemIds, signal) => (await getQueueItemsByIds(itemIds, signal)).map(mapQueueItemDTO), + readProgressPreviews: getProgressPreviews, readNext: async (scope, signal) => { const item = await getNextQueueItem(scope, signal); diff --git a/invokeai/frontend/webv2/src/features/queue/data/serverApi.ts b/invokeai/frontend/webv2/src/features/queue/data/serverApi.ts index 46a29ff8fd1..60054297faf 100644 --- a/invokeai/frontend/webv2/src/features/queue/data/serverApi.ts +++ b/invokeai/frontend/webv2/src/features/queue/data/serverApi.ts @@ -1,4 +1,4 @@ -import type { QueueQueryScope } from '@features/queue/core/types'; +import type { QueueProgressPreviewPayload, QueueQueryScope } from '@features/queue/core/types'; import type { QueueAndProcessorStatusDTO, QueueItemIdsResultDTO, @@ -49,6 +49,10 @@ export const getQueueStatus = ( { signal } ); +/** The latest `invocation_progress` payload of each running item the user owns, with revisions. */ +export const getProgressPreviews = (signal?: AbortSignal): Promise => + apiFetchJson(buildQueueUrl('previews'), { signal }); + export const getCurrentQueueItem = ( scope: QueueQueryScope = {}, signal?: AbortSignal diff --git a/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.test.ts b/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.test.ts index 59e1a179f9c..1ed8df4e42c 100644 --- a/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.test.ts +++ b/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.test.ts @@ -3,6 +3,7 @@ import type { QueueEnqueueResult, QueueEnqueueGenerateRequest, QueueItemProgress, + QueueProgressPreviewPayload, QueueResultImage, } from '@features/queue/core/types'; import type { ActiveProgressTargetSink } from '@features/queue/data/activeProgressTargetStore'; @@ -121,6 +122,27 @@ const createStatusEvent = (overrides: Partial): Que ...overrides, }); +/** The REST snapshot carries the socket event's consumer-facing fields only. */ +const toPreviewPayload = (event: { + queue_id: string; + item_id: number; + session_id: string; + invocation_source_id: string; + revision: number; + message: string; + percentage: number | null; + image: { width: number; height: number; dataURL: string }; +}): QueueProgressPreviewPayload => ({ + image: event.image, + invocation_source_id: event.invocation_source_id, + item_id: event.item_id, + message: event.message, + percentage: event.percentage, + queue_id: event.queue_id, + revision: event.revision, + session_id: event.session_id, +}); + const createQueueBackendItem = (overrides: Partial): QueueBackendItem => ({ id: 1, status: 'in_progress', @@ -162,9 +184,12 @@ interface Harness { api: { [Key in Exclude< keyof QueueCoordinatorBackendPort, - 'emit' | 'on' | 'onConnectionChange' | 'getEnqueueReceipt' + 'emit' | 'on' | 'onConnectionChange' | 'getEnqueueReceipt' | 'readProgressPreviews' >]: ReturnType; - } & { getEnqueueReceipt?: ReturnType }; + } & { + getEnqueueReceipt?: ReturnType; + readProgressPreviews?: ReturnType Promise>>; + }; callbacks: { [Key in keyof QueueCoordinatorCallbacks]: ReturnType }; coordinator: QueueCoordinator; hub: ReturnType; @@ -208,6 +233,9 @@ const createHarness = (options: { galleryRefreshCoalesceMs?: number } = {}): Har Promise.resolve([createImage(`image-${itemId}.png`, sourceQueueItemId)]) ), listItems: vi.fn((): Promise => Promise.resolve([])), + readProgressPreviews: undefined as + | ReturnType Promise>> + | undefined, }; const callbacks = { onGalleryRefresh: vi.fn(), @@ -237,6 +265,9 @@ const createHarness = (options: { galleryRefreshCoalesceMs?: number } = {}): Har get getEnqueueReceipt() { return api.getEnqueueReceipt; }, + get readProgressPreviews() { + return api.readProgressPreviews; + }, emit: hub.emit, on: hub.on, onConnectionChange: hub.onConnectionChange, @@ -656,6 +687,108 @@ describe('queueCoordinator', () => { expect(harness.api.getItem).toHaveBeenCalledTimes(2); }); + it('applies the preview snapshot on visibility and lets the revision gate drop replays', async () => { + const listeners = new Map void>(); + const frame = (revision: number, dataURL: string) => ({ + ...createStatusEvent({ item_id: 1 }), + image: { dataURL, height: 32, width: 64 }, + invocation_source_id: 'denoise', + message: 'Denoising', + percentage: 0.5, + revision, + }); + + vi.stubGlobal('document', { + addEventListener: (type: string, listener: () => void) => listeners.set(type, listener), + removeEventListener: (type: string) => listeners.delete(type), + visibilityState: 'visible', + }); + + try { + harness.api.readProgressPreviews = vi.fn(() => + Promise.resolve([toPreviewPayload(frame(2, 'data:image/png;base64,snapshot'))]) + ); + harness.coordinator.connect(); + await harness.coordinator.submitGenerate('local-1', generateRequest); + harness.socket.fire('invocation_progress', frame(1, 'data:image/png;base64,live-1')); + + listeners.get('visibilitychange')?.(); + await Promise.resolve(); + await Promise.resolve(); + + // The snapshot is newer than the last live frame: applied. + expect(harness.progressImage.set.mock.calls.map(([image]) => (image as { dataUrl: string }).dataUrl)).toEqual([ + 'data:image/png;base64,live-1', + 'data:image/png;base64,snapshot', + ]); + + // The live stream has moved on; a second snapshot at the same revision is a replay. + harness.socket.fire('invocation_progress', frame(3, 'data:image/png;base64,live-3')); + listeners.get('visibilitychange')?.(); + await Promise.resolve(); + await Promise.resolve(); + + expect(harness.progressImage.set).toHaveBeenCalledTimes(3); + expect(harness.api.readProgressPreviews).toHaveBeenCalledTimes(2); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('fetches the preview snapshot after re-adopting runs on reload', async () => { + harness.api.readProgressPreviews = vi.fn(() => + Promise.resolve([ + toPreviewPayload({ + ...createStatusEvent({ item_id: 1 }), + image: { dataURL: 'data:image/png;base64,snapshot', height: 32, width: 64 }, + invocation_source_id: 'denoise', + message: 'Denoising', + percentage: 0.5, + revision: 4, + }), + ]) + ); + harness.api.getItem.mockResolvedValue( + createQueueBackendItem({ id: 1, origin: buildQueueItemOrigin('local-1', 'project-1'), status: 'in_progress' }) + ); + harness.coordinator.connect(); + + const outcomes = await harness.coordinator.reconcile([ + { backendItemIds: [1], id: 'local-1', projectId: 'project-1', status: 'running' }, + ]); + await Promise.resolve(); + await Promise.resolve(); + + expect(outcomes.get('local-1')?.kind).toBe('resumed'); + expect(harness.progressImage.set).toHaveBeenCalledWith( + { dataUrl: 'data:image/png;base64,snapshot', height: 32, width: 64 }, + { itemIndex: 1, queueItemId: 'local-1' } + ); + }); + + it('reopens the revision gate when an item goes back to waiting', async () => { + harness.coordinator.connect(); + await harness.coordinator.submitGenerate('local-1', generateRequest); + const frame = (revision: number, dataURL: string) => ({ + ...createStatusEvent({ item_id: 1 }), + image: { dataURL, height: 32, width: 64 }, + invocation_source_id: 'denoise', + message: 'Denoising', + percentage: 0.5, + revision, + }); + + harness.socket.fire('invocation_progress', frame(5, 'data:image/png;base64,first-leg')); + harness.socket.fire('queue_item_status_changed', createStatusEvent({ item_id: 1, status: 'waiting' })); + harness.socket.fire('queue_item_status_changed', createStatusEvent({ item_id: 1, status: 'in_progress' })); + harness.socket.fire('invocation_progress', frame(1, 'data:image/png;base64,second-leg')); + + expect(harness.progressImage.set.mock.calls.map(([image]) => (image as { dataUrl: string }).dataUrl)).toEqual([ + 'data:image/png;base64,first-leg', + 'data:image/png;base64,second-leg', + ]); + }); + it('drops a preview frame whose revision is not newer than the one already shown', async () => { harness.coordinator.connect(); await harness.coordinator.submitGenerate('local-1', generateRequest); diff --git a/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.ts b/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.ts index d6ea4303d4c..44360203f9e 100644 --- a/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.ts +++ b/invokeai/frontend/webv2/src/features/queue/runtime/coordinator.ts @@ -73,6 +73,7 @@ export type QueueCoordinatorBackendPort = Pick< | 'listItems' | 'on' | 'onConnectionChange' + | 'readProgressPreviews' >; export interface QueueCoordinatorCallbacks { @@ -466,6 +467,37 @@ export const createQueueCoordinator = ( } }; + /** + * Ask the backend for the latest preview frame of every running item and feed + * each through the socket handler, where the revision gate drops anything the + * live stream already delivered. Covers the frames lost while a hidden tab's + * socket was down, and the frames a reloaded page never saw. Best effort: the + * backend also replays them on `subscribe_queue`, and a failure here only means + * waiting for the next step. + */ + const refreshProgressPreviews = async (): Promise => { + if (!isActive() || waits.size === 0 || !backend.readProgressPreviews) { + return; + } + + let previews: Awaited>>; + + try { + previews = await backend.readProgressPreviews(); + } catch { + return; + } + + if (!isActive()) { + return; + } + + for (const preview of previews) { + // Structurally the socket payload; the port cannot name the event type. + handleProgress(preview as unknown as InvocationProgressEvent); + } + }; + const handleStatusChanged = (event: QueueItemStatusChangedEvent): void => { if (!isActive()) { return; @@ -481,6 +513,13 @@ export const createQueueCoordinator = ( } if (!isTerminalBackendStatus(event.status)) { + // Back to the queue (a workflow-call parent waiting on its child, a retry): + // whatever frames follow belong to a new leg, and after a backend restart + // their revisions start over. + if (event.status === 'pending' || event.status === 'waiting') { + latestFrameGates.delete(event.item_id); + } + return; } @@ -649,14 +688,16 @@ export const createQueueCoordinator = ( // is back. The outcome is on the backend already, so sweep on the // visibility edge itself rather than waiting for the reconnect edge. if (typeof document !== 'undefined') { + const visibilityDocument = document; const handleVisibilityChange = (): void => { - if (document.visibilityState === 'visible') { + if (visibilityDocument.visibilityState === 'visible') { void sweep(); + void refreshProgressPreviews(); } }; - document.addEventListener('visibilitychange', handleVisibilityChange); - detachers.push(() => document.removeEventListener('visibilitychange', handleVisibilityChange)); + visibilityDocument.addEventListener('visibilitychange', handleVisibilityChange); + detachers.push(() => visibilityDocument.removeEventListener('visibilitychange', handleVisibilityChange)); } sweepTimer = setInterval(() => { @@ -803,6 +844,10 @@ export const createQueueCoordinator = ( ); } + // A reloaded page has no frame for a run it just re-adopted; the socket only + // brings the next step's. + void refreshProgressPreviews(); + return outcomes; }; diff --git a/tests/app/routers/test_session_queue_previews.py b/tests/app/routers/test_session_queue_previews.py new file mode 100644 index 00000000000..d6da67443a8 --- /dev/null +++ b/tests/app/routers/test_session_queue_previews.py @@ -0,0 +1,99 @@ +"""`GET /queue/{queue_id}/previews` returns the caller's own latest preview frames only.""" + +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from invokeai.app.api.dependencies import ApiDependencies +from invokeai.app.api_app import app +from invokeai.app.services.events.events_common import InvocationProgressEvent +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.users.users_common import UserCreateRequest + + +class MockApiDependencies(ApiDependencies): + invoker: Invoker + + def __init__(self, invoker: Invoker) -> None: + self.invoker = invoker + + +@pytest.fixture +def setup_jwt_secret(): + from invokeai.app.services.auth.token_service import set_jwt_secret + + set_jwt_secret("test-secret-key-for-unit-tests-only-do-not-use-in-production") + + +@pytest.fixture +def client(): + return TestClient(app) + + +@pytest.fixture +def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker): + mock_invoker.services.configuration.multiuser = True + mock_deps = MockApiDependencies(mock_invoker) + monkeypatch.setattr("invokeai.app.api.routers.auth.ApiDependencies", mock_deps) + monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", mock_deps) + monkeypatch.setattr("invokeai.app.api.routers.session_queue.ApiDependencies", mock_deps) + yield + + +def _create_user(mock_invoker: Invoker, email: str, display_name: str, is_admin: bool = False) -> str: + user = mock_invoker.services.users.create( + UserCreateRequest(email=email, display_name=display_name, password="TestPass123", is_admin=is_admin) + ) + return user.user_id + + +def _login(client: TestClient, email: str) -> str: + response = client.post("/api/v1/auth/login", json={"email": email, "password": "TestPass123", "remember_me": False}) + assert response.status_code == 200, response.text + return response.json()["token"] + + +def _event(item_id: int, user_id: str) -> InvocationProgressEvent: + return InvocationProgressEvent( + queue_id="default", + item_id=item_id, + batch_id="batch-1", + user_id=user_id, + session_id=f"session-{item_id}", + invocation={"type": "add", "id": "node-1", "a": 1, "b": 2}, + invocation_source_id="node-1", + message="Denoising", + percentage=0.5, + image={"width": 64, "height": 64, "dataURL": "data:image/jpeg;base64,frame"}, + revision=2, + ) + + +def test_previews_are_scoped_to_the_caller_even_for_admins( + setup_jwt_secret: None, enable_multiuser: Any, mock_invoker: Invoker, client: TestClient +) -> None: + admin_id = _create_user(mock_invoker, "admin@test.com", "Admin", is_admin=True) + user_id = _create_user(mock_invoker, "user@test.com", "User") + previews = mock_invoker.services.progress_previews + previews.record(_event(1, admin_id)) + previews.record(_event(2, user_id)) + + admin_token = _login(client, "admin@test.com") + user_token = _login(client, "user@test.com") + + admin_response = client.get("/api/v1/queue/default/previews", headers={"Authorization": f"Bearer {admin_token}"}) + user_response = client.get("/api/v1/queue/default/previews", headers={"Authorization": f"Bearer {user_token}"}) + + assert admin_response.status_code == 200, admin_response.text + assert [entry["item_id"] for entry in admin_response.json()] == [1] + assert admin_response.json()[0]["revision"] == 2 + assert admin_response.json()[0]["image"]["dataURL"] == "data:image/jpeg;base64,frame" + assert user_response.status_code == 200 + assert [entry["item_id"] for entry in user_response.json()] == [2] + + +def test_previews_require_authentication_in_multiuser_mode( + setup_jwt_secret: None, enable_multiuser: Any, client: TestClient +) -> None: + assert client.get("/api/v1/queue/default/previews").status_code == 401 diff --git a/tests/app/services/progress_previews/__init__.py b/tests/app/services/progress_previews/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/app/services/progress_previews/test_progress_previews.py b/tests/app/services/progress_previews/test_progress_previews.py new file mode 100644 index 00000000000..89eb6de062c --- /dev/null +++ b/tests/app/services/progress_previews/test_progress_previews.py @@ -0,0 +1,147 @@ +from invokeai.app.services.events.events_common import InvocationProgressEvent +from invokeai.app.services.progress_previews.progress_previews_default import ( + PROGRESS_IMAGE_MIN_INTERVAL_S, + MemoryProgressPreviews, +) + +_COMMON_FIELDS = { + "queue_id": "default", + "batch_id": "batch-1", + "invocation": {"type": "add", "id": "node-1", "a": 1, "b": 2}, + "invocation_source_id": "node-1", + "message": "Denoising", + "image": {"width": 64, "height": 64, "dataURL": "data:image/jpeg;base64,frame"}, +} + + +def _event(item_id: int, user_id: str = "owner-1", session_id: str = "session-1", **overrides): + return InvocationProgressEvent( + **{**_COMMON_FIELDS, "item_id": item_id, "user_id": user_id, "session_id": session_id, **overrides} + ) + + +class _Clock: + def __init__(self) -> None: + self.now = 100.0 + + def __call__(self) -> float: + return self.now + + +def test_revisions_are_monotonic_per_item_and_restart_with_a_new_session() -> None: + clock = _Clock() + store = MemoryProgressPreviews(clock=clock) + + first = store.reserve_revision(1, "session-1", 0.1) + clock.now += 1 + second = store.reserve_revision(1, "session-1", 0.2) + clock.now += 1 + other_item = store.reserve_revision(2, "session-2", 0.1) + clock.now += 1 + new_session = store.reserve_revision(1, "session-1b", 0.1) + + assert (first, second) == (1, 2) + assert other_item == 1 + # A new session on the same item is a new epoch: the client resets its gate on session change. + assert new_session == 1 + + +def test_frames_inside_the_interval_are_dropped_except_final_and_indeterminate_ones() -> None: + clock = _Clock() + store = MemoryProgressPreviews(clock=clock) + + assert store.reserve_revision(1, "session-1", 0.1) == 1 + clock.now += PROGRESS_IMAGE_MIN_INTERVAL_S / 2 + assert store.reserve_revision(1, "session-1", 0.2) is None + # A final frame is never dropped: it is the one the preview settles on. + assert store.reserve_revision(1, "session-1", 1.0) == 2 + # Nor an indeterminate one. + assert store.reserve_revision(1, "session-1", None) == 3 + clock.now += PROGRESS_IMAGE_MIN_INTERVAL_S * 1.5 + assert store.reserve_revision(1, "session-1", 0.3) == 4 + + +def test_throttle_is_per_item() -> None: + clock = _Clock() + store = MemoryProgressPreviews(clock=clock) + + assert store.reserve_revision(1, "session-1", 0.1) == 1 + # Another item (another GPU) is not held back by the first one's frame. + assert store.reserve_revision(2, "session-2", 0.1) == 1 + + +def test_records_the_latest_event_and_lists_only_the_owners_running_items() -> None: + store = MemoryProgressPreviews() + + store.record(_event(1, revision=1)) + store.record(_event(1, revision=2)) + store.record(_event(2, user_id="owner-2", session_id="session-2", revision=1)) + store.record(_event(3, queue_id="other", session_id="session-3", revision=1)) + + assert store.get(1) is not None and store.get(1).revision == 2 + assert [event.item_id for event in store.list_for_user("owner-1")] == [1, 3] + assert [event.item_id for event in store.list_for_user("owner-1", "default")] == [1] + assert [event.item_id for event in store.list_for_user("owner-2")] == [2] + assert store.list_for_user("stranger") == [] + + +def test_clear_forgets_the_frame_but_keeps_the_revision_counter() -> None: + # A workflow-call parent suspends (session leg ends, frame cleared) and later resumes with + # the same item and session. Its next frames must outrank the ones a client already holds. + clock = _Clock() + store = MemoryProgressPreviews(clock=clock) + + assert store.reserve_revision(1, "session-1", 0.1) == 1 + store.record(_event(1, revision=1)) + + store.clear(1) + + assert store.get(1) is None + assert store.list_for_user("owner-1") == [] + # No throttle carries over, but the counter does. + assert store.reserve_revision(1, "session-1", 0.1) == 2 + + store.record(_event(1, revision=2)) + store.clear_all() + assert store.list_for_user("owner-1") == [] + assert store.reserve_revision(1, "session-1", 0.1) == 1 + + +def test_clear_node_drops_the_frame_only_if_that_node_produced_it() -> None: + # Once the denoise node completes, replaying its frame would show a finished denoise as + # still running. A later node's frame must not be dropped by the earlier node's completion. + store = MemoryProgressPreviews() + store.record(_event(1, revision=1)) + + store.clear_node(1, "other-node") + assert store.get(1) is not None + + store.clear_node(1, "node-1") + assert store.get(1) is None + assert store.list_for_user("owner-1") == [] + + +def test_frameless_counters_are_bounded() -> None: + from invokeai.app.services.progress_previews import progress_previews_default + + store = MemoryProgressPreviews() + limit = progress_previews_default.COUNTER_RETENTION_LIMIT + for item_id in range(limit + 10): + store.reserve_revision(item_id, f"session-{item_id}", 0.1) + store.clear(item_id) + + assert len(store._items) == limit + # The oldest counters went; the newest survive. + assert 0 not in store._items + assert limit + 9 in store._items + + +def test_reserved_revision_survives_a_record_with_the_same_session() -> None: + clock = _Clock() + store = MemoryProgressPreviews(clock=clock) + + revision = store.reserve_revision(1, "session-1", 0.1) + store.record(_event(1, revision=revision)) + clock.now += 1 + + assert store.reserve_revision(1, "session-1", 0.2) == 2 diff --git a/tests/app/services/progress_previews/test_signal_progress.py b/tests/app/services/progress_previews/test_signal_progress.py new file mode 100644 index 00000000000..1754f7f933e --- /dev/null +++ b/tests/app/services/progress_previews/test_signal_progress.py @@ -0,0 +1,72 @@ +"""`signal_progress` wires the preview store into the progress event path.""" + +from types import SimpleNamespace + +from PIL import Image + +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews +from invokeai.app.services.shared.invocation_context import UtilInterface + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +class _FakeEvents: + def __init__(self) -> None: + self.emitted: list[dict] = [] + + def emit_invocation_progress(self, **kwargs): + self.emitted.append(kwargs) + queue_item = kwargs["queue_item"] + return SimpleNamespace( + invocation=kwargs["invocation"], + item_id=queue_item.item_id, + session_id=queue_item.session_id, + revision=kwargs.get("revision"), + user_id=queue_item.user_id, + queue_id=queue_item.queue_id, + ) + + +def _util(clock: _Clock) -> tuple[UtilInterface, _FakeEvents, MemoryProgressPreviews]: + events = _FakeEvents() + previews = MemoryProgressPreviews(clock=clock) + services = SimpleNamespace(events=events, progress_previews=previews) + queue_item = SimpleNamespace(item_id=7, session_id="session-7", user_id="owner", queue_id="default") + data = SimpleNamespace(queue_item=queue_item, invocation=SimpleNamespace(id="node")) + return UtilInterface(services, data, is_canceled=lambda: False), events, previews # type: ignore[arg-type] + + +def test_image_frames_are_revisioned_throttled_and_retained() -> None: + clock = _Clock() + util, events, previews = _util(clock) + image = Image.new("RGB", (8, 8)) + + util.signal_progress("Denoising", 0.1, image, (64, 64)) + clock.now += 0.01 + util.signal_progress("Denoising", 0.2, image, (64, 64)) # inside the interval: dropped + clock.now += 1 + util.signal_progress("Denoising", 0.3, image, (64, 64)) + + assert [call["revision"] for call in events.emitted] == [1, 2] + assert [call["percentage"] for call in events.emitted] == [0.1, 0.3] + assert all(call["image"] is not None for call in events.emitted) + retained = previews.get(7) + assert retained is not None and retained.revision == 2 + + +def test_imageless_progress_is_neither_throttled_nor_revisioned() -> None: + clock = _Clock() + util, events, previews = _util(clock) + + util.signal_progress("Loading model") + util.signal_progress("Loading model", 0.5) + + assert len(events.emitted) == 2 + assert all("revision" not in call and call.get("image") is None for call in events.emitted) + assert previews.get(7) is None diff --git a/tests/app/services/session_processor/test_for_loop_processor.py b/tests/app/services/session_processor/test_for_loop_processor.py index 355a56a53be..9a389ebba74 100644 --- a/tests/app/services/session_processor/test_for_loop_processor.py +++ b/tests/app/services/session_processor/test_for_loop_processor.py @@ -10,6 +10,7 @@ from invokeai.app.invocations.fields import InputField, OutputField from invokeai.app.invocations.loops import ForInvocation, ForReturnInvocation from invokeai.app.services.events.events_common import QueueItemStatusChangedEvent +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews from invokeai.app.services.session_processor.session_processor_default import ( DefaultSessionProcessor, DefaultSessionRunner, @@ -251,6 +252,7 @@ def _build_processor( events=events, logger=_Logger(), performance_statistics=_Stats(), + progress_previews=MemoryProgressPreviews(), session_queue=queue, image_moves=None, ) diff --git a/tests/app/services/test_for_loop_session_runner.py b/tests/app/services/test_for_loop_session_runner.py index a3db4d493c0..85a0334fd42 100644 --- a/tests/app/services/test_for_loop_session_runner.py +++ b/tests/app/services/test_for_loop_session_runner.py @@ -8,6 +8,7 @@ from invokeai.app.invocations.fields import InputField, OutputField from invokeai.app.invocations.loops import ForInvocation, ForInvocationOutput, ForReturnInvocation, LoopState from invokeai.app.invocations.primitives import BooleanOutput +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews from invokeai.app.services.session_processor.session_processor_default import DefaultSessionRunner from invokeai.app.services.shared.graph import CollectInvocation, Graph, GraphExecutionState, IterateInvocation from invokeai.app.services.shared.invocation_context import InvocationContext @@ -167,6 +168,7 @@ def _build_runner( runner.start( services=SimpleNamespace( performance_statistics=_DummyStats(), + progress_previews=MemoryProgressPreviews(), events=events, logger=_DummyLogger(), configuration=_DummyConfig(), diff --git a/tests/app/services/test_session_processor_callbacks.py b/tests/app/services/test_session_processor_callbacks.py index bfca1b50623..2d7f9fd2255 100644 --- a/tests/app/services/test_session_processor_callbacks.py +++ b/tests/app/services/test_session_processor_callbacks.py @@ -5,6 +5,7 @@ import pytest +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews from invokeai.app.services.session_processor.session_processor_default import DefaultSessionRunner from invokeai.app.services.shared.graph import CollectInvocation, Graph, GraphExecutionState, IterateInvocation @@ -21,6 +22,7 @@ def test_after_run_node_callback_receives_control_node_inputs(monkeypatch: pytes services = SimpleNamespace( configuration=SimpleNamespace(node_cache_size=0), + progress_previews=MemoryProgressPreviews(), events=Mock(), logger=Mock(), performance_statistics=Mock(), diff --git a/tests/app/services/workflow_call_test_utils.py b/tests/app/services/workflow_call_test_utils.py index 34368cff033..e1c7e7b52aa 100644 --- a/tests/app/services/workflow_call_test_utils.py +++ b/tests/app/services/workflow_call_test_utils.py @@ -15,6 +15,7 @@ WorkflowReturnInvocation, WorkflowReturnOutput, ) +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews from invokeai.app.services.session_processor.session_processor_default import ( DefaultSessionProcessor, DefaultSessionRunner, @@ -815,6 +816,7 @@ def _build_runner(monkeypatch: pytest.MonkeyPatch) -> DefaultSessionRunner: (), { "performance_statistics": _DummyStats(), + "progress_previews": MemoryProgressPreviews(), "events": _DummyEvents(), "logger": _DummyLogger(), "configuration": _DummyConfig(), @@ -847,6 +849,7 @@ def _build_workflow_runner(monkeypatch: pytest.MonkeyPatch, session_queue=None): (), { "performance_statistics": _DummyStats(), + "progress_previews": MemoryProgressPreviews(), "events": events, "logger": _DummyLogger(), "configuration": _DummyConfig(), @@ -1214,6 +1217,7 @@ class DummyConfig: (), { "performance_statistics": DummyStats(), + "progress_previews": MemoryProgressPreviews(), "events": DummyEvents(), "logger": DummyLogger(), "configuration": DummyConfig(), @@ -1252,6 +1256,7 @@ def test_on_after_run_session_does_not_complete_incomplete_session(monkeypatch: (), { "performance_statistics": _DummyStats(), + "progress_previews": MemoryProgressPreviews(), "events": _DummyEvents(), "logger": _DummyLogger(), "configuration": _DummyConfig(), @@ -1369,6 +1374,7 @@ def test_run_persists_waiting_session_without_completing_queue_item(monkeypatch: (), { "performance_statistics": _DummyStats(), + "progress_previews": MemoryProgressPreviews(), "events": _DummyEvents(), "logger": _DummyLogger(), "configuration": _DummyConfig(), diff --git a/tests/app/test_progress_preview_replay_socketio.py b/tests/app/test_progress_preview_replay_socketio.py new file mode 100644 index 00000000000..c7f858d7c27 --- /dev/null +++ b/tests/app/test_progress_preview_replay_socketio.py @@ -0,0 +1,73 @@ +"""On `subscribe_queue`, a socket is sent the latest preview frame of each running item its user owns.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI + +from invokeai.app.api.dependencies import ApiDependencies +from invokeai.app.api.sockets import SocketIO +from invokeai.app.services.events.events_common import InvocationProgressEvent +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +def _event(item_id: int, user_id: str, queue_id: str = "default") -> InvocationProgressEvent: + return InvocationProgressEvent( + queue_id=queue_id, + item_id=item_id, + batch_id="batch-1", + user_id=user_id, + session_id=f"session-{item_id}", + invocation={"type": "add", "id": "node-1", "a": 1, "b": 2}, + invocation_source_id="node-1", + message="Denoising", + percentage=0.5, + image={"width": 64, "height": 64, "dataURL": "data:image/jpeg;base64,frame"}, + revision=3, + ) + + +@pytest.mark.anyio +async def test_subscribe_replays_only_the_users_previews_to_that_socket(monkeypatch) -> None: + previews = MemoryProgressPreviews() + previews.record(_event(1, "owner-1")) + previews.record(_event(2, "owner-2")) + previews.record(_event(3, "owner-1", queue_id="other")) + monkeypatch.setattr( + ApiDependencies, + "invoker", + SimpleNamespace(services=SimpleNamespace(progress_previews=previews)), + raising=False, + ) + + socketio = SocketIO(FastAPI()) + socketio._sio.enter_room = AsyncMock() + socketio._sio.emit = AsyncMock() + socketio._socket_users["sid-1"] = {"user_id": "owner-1", "is_admin": True, "authenticated": True} + + await socketio._handle_sub_queue("sid-1", {"queue_id": "default"}) + + socketio._sio.emit.assert_awaited_once_with( + event="invocation_progress", + data=previews.get(1).model_dump(mode="json"), + to="sid-1", + ) + + +@pytest.mark.anyio +async def test_subscribe_without_an_invoker_replays_nothing(monkeypatch) -> None: + monkeypatch.delattr(ApiDependencies, "invoker", raising=False) + socketio = SocketIO(FastAPI()) + socketio._sio.enter_room = AsyncMock() + socketio._sio.emit = AsyncMock() + socketio._socket_users["sid-1"] = {"user_id": "owner-1", "is_admin": False, "authenticated": True} + + await socketio._handle_sub_queue("sid-1", {"queue_id": "default"}) + + socketio._sio.emit.assert_not_awaited() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 47541fd4c12..43d309484fa 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -9,7 +9,7 @@ ) from invokeai.app.invocations.fields import InputField, OutputField from invokeai.app.invocations.image import ImageField -from invokeai.app.services.events.events_common import EventBase +from invokeai.app.services.events.events_common import EventBase, InvocationProgressEvent from invokeai.app.services.session_processor.session_processor_common import ProgressImage from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem from invokeai.app.services.shared.invocation_context import InvocationContext @@ -164,8 +164,13 @@ def emit_invocation_progress( message: str, percentage: float | None = None, image: "ProgressImage | None" = None, - ) -> None: - pass + revision: int | None = None, + ) -> InvocationProgressEvent: + if invocation is None: + # Model-load tests signal imageless progress outside any invocation; nothing retains + # the (unbuildable) event in that case. + return None # type: ignore[return-value] + return super().emit_invocation_progress(queue_item, invocation, message, percentage, image, revision) def wait_until(condition: Callable[[], bool], timeout: int = 10, interval: float = 0.1) -> None: From a4d940bffafae46b9056bf7e557c9b49a320b21e Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 6 Sep 2026 16:42:23 -0400 Subject: [PATCH 2/3] fix(ci): regenerate openapi.json and give socket test stubs the preview store The openapi-checks workflow diffs `invokeai/frontend/web/openapi.json` (separate from the typegen check's `schema.ts`); regenerated for the new previews route and the `revision` field. The `ProgressImage` component only reorders keys. `test_socket_privilege_revocation.py` binds `ApiDependencies` to hand-built services stubs; the subscribe replay now reads `progress_previews` from them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DoHBJ9QPWfq1qMTfzdzuke --- invokeai/frontend/web/openapi.json | 165 ++++++++++++++++-- tests/app/test_socket_privilege_revocation.py | 10 ++ 2 files changed, 164 insertions(+), 11 deletions(-) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 69ef493b4b1..3c91d6a743a 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -11146,6 +11146,58 @@ } } }, + "/api/v1/queue/{queue_id}/previews": { + "get": { + "tags": ["queue"], + "summary": "Get Progress Previews", + "description": "The latest denoising preview frame of each of the caller's running queue items: the same\npayloads as the `invocation_progress` socket events, with their revisions. A client whose socket\nwas dropped, or whose tab was hidden, reconciles from this instead of waiting for the next step.\nOwner-scoped: progress is personal UI, so even admins see only their own.", + "operationId": "get_progress_previews", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "queue_id", + "in": "path", + "required": true, + "schema": { + "type": "string", + "description": "The queue id to perform this operation on", + "title": "Queue Id" + }, + "description": "The queue id to perform this operation on" + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/ProgressPreviewDTO" + }, + "title": "Response 200 Get Progress Previews" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, "/api/v1/queue/{queue_id}/next": { "get": { "tags": ["queue"], @@ -51692,6 +51744,19 @@ "default": null, "description": "The device processing this session, e.g. 'cuda:1' (set only when running on a GPU)", "title": "Device" + }, + "revision": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "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.", + "title": "Revision" } }, "required": [ @@ -51708,7 +51773,8 @@ "message", "percentage", "image", - "device" + "device", + "revision" ], "title": "InvocationProgressEvent", "type": "object" @@ -79657,29 +79723,106 @@ "title": "PresetType" }, "ProgressImage": { - "description": "The progress image sent intermittently during processing", "properties": { "width": { - "description": "The effective width of the image in pixels", - "minimum": 1, + "type": "integer", + "minimum": 1.0, "title": "Width", - "type": "integer" + "description": "The effective width of the image in pixels" }, "height": { - "description": "The effective height of the image in pixels", - "minimum": 1, + "type": "integer", + "minimum": 1.0, "title": "Height", - "type": "integer" + "description": "The effective height of the image in pixels" }, "dataURL": { - "description": "The image data as a b64 data URL", + "type": "string", "title": "Dataurl", - "type": "string" + "description": "The image data as a b64 data URL" } }, + "type": "object", "required": ["width", "height", "dataURL"], "title": "ProgressImage", - "type": "object" + "description": "The progress image sent intermittently during processing" + }, + "ProgressPreviewDTO": { + "properties": { + "queue_id": { + "type": "string", + "title": "Queue Id", + "description": "The ID of the queue" + }, + "item_id": { + "type": "integer", + "title": "Item Id", + "description": "The ID of the queue item" + }, + "session_id": { + "type": "string", + "title": "Session Id", + "description": "The ID of the session (aka graph execution state)" + }, + "invocation_source_id": { + "type": "string", + "title": "Invocation Source Id", + "description": "The ID of the prepared invocation's source node" + }, + "revision": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Revision", + "description": "Monotonic per queue item and session; see InvocationProgressEvent" + }, + "message": { + "type": "string", + "title": "Message", + "description": "A message to display" + }, + "percentage": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Percentage", + "description": "The percentage of the progress, or null if indeterminate" + }, + "image": { + "anyOf": [ + { + "$ref": "#/components/schemas/ProgressImage" + }, + { + "type": "null" + } + ], + "description": "The latest denoising preview" + } + }, + "type": "object", + "required": [ + "queue_id", + "item_id", + "session_id", + "invocation_source_id", + "revision", + "message", + "percentage", + "image" + ], + "title": "ProgressPreviewDTO", + "description": "The latest preview frame of a running queue item, as served by `GET /queue/{queue_id}/previews`.\n\nThe fields of the `invocation_progress` socket event a preview consumer reads, without the\nserialized invocation. A separate model rather than the event itself: exposing the event as a\nresponse model would give it a second OpenAPI schema variant and change the generated client\ntypes for every socket consumer." }, "ProjectBoardItemDTO": { "properties": { diff --git a/tests/app/test_socket_privilege_revocation.py b/tests/app/test_socket_privilege_revocation.py index e09011b23f3..b094c0b5662 100644 --- a/tests/app/test_socket_privilege_revocation.py +++ b/tests/app/test_socket_privilege_revocation.py @@ -18,6 +18,7 @@ from invokeai.app.api.sockets import SOCKET_REVALIDATION_FAILURE_LIMIT, SocketIO from invokeai.app.services.events.events_common import UserAccessChangedEvent +from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews @pytest.fixture @@ -40,6 +41,7 @@ def _patch_multiuser_context( invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=lambda candidate_user_id: user if candidate_user_id == user_id else None), ) ) @@ -156,6 +158,7 @@ def _patch_record_agreeing_with(self, monkeypatch: pytest.MonkeyPatch, event: Us invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=lambda user_id: record), ) ) @@ -389,6 +392,7 @@ def get(user_id: str): invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=multiuser), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=get), events=SimpleNamespace(emit_user_access_changed=lambda **kwargs: emitted.append(kwargs)), session_processor=SimpleNamespace( @@ -622,6 +626,7 @@ def get(user_id: str): invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=get), events=SimpleNamespace(emit_user_access_changed=lambda **kwargs: None), ) @@ -942,6 +947,7 @@ def _patch_users(self, monkeypatch: pytest.MonkeyPatch, user) -> None: invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=lambda user_id: user), ) ) @@ -973,6 +979,7 @@ async def test_the_sweep_republishes_after_a_failed_demotion(self, monkeypatch: invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace( get=lambda user_id: SimpleNamespace(user_id="user-1", is_admin=False, is_active=True, token_epoch=0) ), @@ -1081,6 +1088,7 @@ def _patch_users(self, monkeypatch: pytest.MonkeyPatch, user) -> None: invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=lambda user_id: user), ) ) @@ -1136,6 +1144,7 @@ def explode(user_id: str): invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=explode), ) ) @@ -1193,6 +1202,7 @@ async def test_a_deferred_disconnect_is_applied_once_the_record_reads_again( invoker = SimpleNamespace( services=SimpleNamespace( configuration=SimpleNamespace(multiuser=True), + progress_previews=MemoryProgressPreviews(), users=SimpleNamespace(get=lambda user_id: deactivated), events=SimpleNamespace(emit_user_access_changed=lambda **kwargs: emitted.append(kwargs)), session_processor=SimpleNamespace(get_running_queue_item_owners=lambda: set()), From d88b40679fbed1f26ee4b981facb5cb977937862 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 6 Sep 2026 18:02:39 -0400 Subject: [PATCH 3/3] fix(preview): follow the slot's own frame, not the store-wide latest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With two slots live — a long video render next to a quick image batch — the single-frame preview read the store-wide "latest" frame and matched it to the followed placeholder. The latest belonged to whichever slot stepped last, and releasing that slot when its batch finished nulled it outright. The video slot's own frame was still stored per slot (the gallery cell kept showing it), but the preview rendered an empty card until the video's next step, minutes away. Reloading fixed it only because reconcile re-fed the frame. LivePreview now reads its followed slot's frame directly, and the store's latest falls back to the most recently updated remaining slot instead of null so the header toggle and the editor's Current Image node do not blank either. The store-wide matcher is gone. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DoHBJ9QPWfq1qMTfzdzuke (cherry picked from commit 0f36ead67ebefc82f4dc563626781f30d64d172a) --- .../queue/data/progressImageStore.test.ts | 31 +++++++++++++++++++ .../features/queue/data/progressImageStore.ts | 31 +++++++++++++++++-- .../PreviewNavigation.browser.test.tsx | 27 ++++++++++++++++ .../widgets/preview/PreviewWidgetView.test.ts | 30 +----------------- .../widgets/preview/PreviewWidgetView.tsx | 27 +++------------- 5 files changed, 92 insertions(+), 54 deletions(-) diff --git a/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.test.ts b/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.test.ts index 07ff7389ca0..724101b83a4 100644 --- a/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.test.ts +++ b/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.test.ts @@ -2,6 +2,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { consumeQueueItemSwapProgressImage, + getLatestProgressImage, getQueueItemBridgeProgressImage, getQueueItemSwapProgressImage, progressImageStore, @@ -133,3 +134,33 @@ describe('progressImageStore held frames', () => { expect(vi.getTimerCount()).toBe(0); }); }); + +describe('progressImageStore latest frame', () => { + beforeEach(() => { + progressImageStore.clear(); + }); + + it('falls back to the most recently updated live slot when the latest slot is released', () => { + progressImageStore.set(frame('video-1'), target('video')); + progressImageStore.set(frame('anima-1'), target('anima')); + progressImageStore.set(frame('video-2'), target('video')); + progressImageStore.set(frame('anima-2'), target('anima')); + + progressImageStore.clear(target('anima')); + + expect(getLatestProgressImage()).toEqual({ ...frame('video-2'), target: target('video') }); + + progressImageStore.clear(target('video')); + + expect(getLatestProgressImage()).toBeNull(); + }); + + it('keeps the latest frame when a slot that is not the latest is released', () => { + progressImageStore.set(frame('anima-1'), target('anima')); + progressImageStore.set(frame('video-1'), target('video')); + + progressImageStore.clear(target('anima')); + + expect(getLatestProgressImage()).toEqual({ ...frame('video-1'), target: target('video') }); + }); +}); diff --git a/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.ts b/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.ts index 5c7cdd48359..dd1d68915de 100644 --- a/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.ts +++ b/invokeai/frontend/webv2/src/features/queue/data/progressImageStore.ts @@ -54,6 +54,8 @@ interface SwapFrame { } const snapshotsByTarget = createKeyedTransientStore(); +/** Insertion order is recency: `set` re-inserts, so the last entry is the most recently updated slot. */ +const targetsByKey = new Map(); const bridgeFrames = createKeyedTransientStore(); const swapFrames = createKeyedTransientStore(); const swapExpiryTimers = new Map>(); @@ -116,6 +118,7 @@ export const progressImageStore = { if (!target) { latestSnapshotStore.patchSnapshot({ latestSnapshot: null }); snapshotsByTarget.clear(); + targetsByKey.clear(); bridgeFrames.clear(); for (const timer of swapExpiryTimers.values()) { @@ -132,9 +135,14 @@ export const progressImageStore = { const didClearLatest = isLatestTarget(target); snapshotsByTarget.delete(targetKey); + targetsByKey.delete(targetKey); if (didClearLatest) { - latestSnapshotStore.patchSnapshot({ latestSnapshot: null }); + // Another slot may still be live: a video rendering for minutes next to a + // quick image batch. Falling to null here left the single-frame preview + // blank until the video's next step, while the gallery cell — which reads + // its own slot — kept showing the frame. + latestSnapshotStore.patchSnapshot({ latestSnapshot: getMostRecentSnapshot() }); } }, /** Routing landed: these are the images the held swap frame may be painted over. */ @@ -169,11 +177,27 @@ export const progressImageStore = { latestSnapshotStore.patchSnapshot({ latestSnapshot: target ? { ...image, target } : image }); if (target) { - snapshotsByTarget.set(getTargetKey(target), image); + const targetKey = getTargetKey(target); + + snapshotsByTarget.set(targetKey, image); + targetsByKey.delete(targetKey); + targetsByKey.set(targetKey, target); } }, }; +function getMostRecentSnapshot(): LatestProgressImageSnapshot | null { + for (const [targetKey, target] of [...targetsByKey.entries()].reverse()) { + const image = snapshotsByTarget.get(targetKey); + + if (image) { + return { ...image, target }; + } + } + + return null; +} + registerAccountOwnedResource({ clear: () => progressImageStore.clear(), name: 'queue-progress-images', @@ -186,6 +210,9 @@ export const consumeQueueItemSwapProgressImage = (queueItemId: string): void => dropSwap(queueItemId); }; +export const getLatestProgressImage = (): LatestProgressImageSnapshot | null => + latestSnapshotStore.getSnapshot().latestSnapshot; + export const getQueueItemBridgeProgressImage = (queueItemId: string): ProgressImageSnapshot | null => bridgeFrames.get(queueItemId) ?? null; diff --git a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewNavigation.browser.test.tsx b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewNavigation.browser.test.tsx index ed30e128d68..4a011493520 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewNavigation.browser.test.tsx +++ b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewNavigation.browser.test.tsx @@ -122,6 +122,7 @@ const mocks = vi.hoisted(() => { recentImages, bridgeProgressImage: null as unknown, runningProgressTargets: undefined as unknown[] | undefined, + slotProgressImage: undefined as unknown, useActiveProgressTarget: vi.fn(() => null as unknown), useProgressImage: vi.fn(() => null as unknown), }; @@ -150,6 +151,16 @@ vi.mock('@features/queue/react', async (importOriginal) => ({ useFollowedProgressTargets: () => mockProgressTargets(), useProgressImage: () => mocks.useProgressImage(), useQueueItemBridgeProgressImage: () => mocks.bridgeProgressImage, + // The slot's own frame: derived from the "latest" mock by target unless a test overrides it. + useQueueItemProgressImage: (queueItemId: string, itemIndex: number) => { + if (mocks.slotProgressImage !== undefined) { + return mocks.slotProgressImage; + } + + const latest = mocks.useProgressImage() as { target?: { itemIndex: number; queueItemId: string } } | null; + + return latest?.target?.queueItemId === queueItemId && latest.target.itemIndex === itemIndex ? latest : null; + }, })); vi.mock('@features/gallery/queries', () => ({ @@ -474,6 +485,7 @@ beforeEach(() => { mocks.useProgressImage.mockReturnValue(null); mocks.bridgeProgressImage = null; mocks.runningProgressTargets = undefined; + mocks.slotProgressImage = undefined; }); afterEach(async () => { @@ -1444,6 +1456,21 @@ describe('preview keyboard navigation boundary', () => { expect(host?.textContent).toContain('64 × 64'); }); + it("shows the followed slot's own frame even when the store-wide latest frame is gone", async () => { + // A quick image batch finished next to a long video render while the tab was + // hidden: releasing the batch's slot cleared the latest frame. The video slot + // still has its frame and must not render an empty card until its next step. + mocks.project.queue.items = [queueItem]; + mocks.project.settings.showProgressImagesInViewer = true; + mocks.useActiveProgressTarget.mockReturnValue({ itemIndex: 1, queueItemId: 'queue-item-live' }); + mocks.useProgressImage.mockReturnValue(null); + mocks.slotProgressImage = { dataUrl: 'data:image/png;base64,video-step', height: 64, width: 64 }; + + await render(); + + expect(host?.querySelector('img[src="data:image/png;base64,video-step"]')).not.toBeNull(); + }); + it('follows a running slot over a settling one so a concurrent session is never hidden', async () => { // Multi-GPU: slot 1 completed and is settling, slot 2 is still streaming. // The single-frame preview must show slot 2 live, not slot 1's static frame. diff --git a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.test.ts b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.test.ts index 9abaf7ff5d1..4ab8032ed22 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.test.ts +++ b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.test.ts @@ -3,37 +3,9 @@ import type { GalleryItem } from '@features/gallery'; import { GALLERY_MAX_ROWS } from '@features/gallery/queries'; import { describe, expect, it } from 'vitest'; -import { getMatchingProgressImage, getVideoFrameCopyNotice } from './PreviewWidgetView'; +import { getVideoFrameCopyNotice } from './PreviewWidgetView'; import { mergePreviewBoardItems } from './usePreviewNavigation'; -describe('getMatchingProgressImage', () => { - const placeholder = { - backendItemId: null, - boardId: 'none', - height: 768, - id: 'queue-1:1', - itemIndex: 2, - queueItemId: 'queue-1', - width: 512, - }; - const progressImage = { - dataUrl: 'data:image/png;base64,abc', - height: 768, - target: { itemIndex: 2, queueItemId: 'queue-1' }, - width: 512, - }; - - it('returns progress only when it belongs to the current placeholder', () => { - expect(getMatchingProgressImage(progressImage, placeholder)).toBe(progressImage); - expect( - getMatchingProgressImage({ ...progressImage, target: { itemIndex: 1, queueItemId: 'queue-1' } }, placeholder) - ).toBeNull(); - expect( - getMatchingProgressImage({ ...progressImage, target: { itemIndex: 2, queueItemId: 'queue-2' } }, placeholder) - ).toBeNull(); - }); -}); - describe('mergePreviewBoardItems', () => { const item = (kind: GalleryItem['kind'], name: string, createdAt: string, starred = false): GalleryItem => { const base = { diff --git a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.tsx b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.tsx index b4324b95cf0..9cf4c62ca72 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.tsx +++ b/invokeai/frontend/webv2/src/workbench/widgets/preview/PreviewWidgetView.tsx @@ -41,11 +41,9 @@ import { useActiveProgressTargets, useFollowedProgressTargets, useItemProgress, - useProgressImage, useQueueItemBridgeProgressImage, useQueueItemProgressImage, useQueueItemSwapProgressImage, - type LatestProgressImageSnapshot, } from '@features/queue/react'; import { imageUrlToStreamingSource, @@ -156,22 +154,6 @@ const getBoardName = ( ): string => boardId === 'none' ? uncategorizedLabel : (boards.find((board) => board.id === boardId)?.name ?? unknownBoardLabel); -export const getMatchingProgressImage = ( - progressImage: LatestProgressImageSnapshot | null, - placeholder: GalleryQueuePlaceholder | null -): LatestProgressImageSnapshot | null => { - if ( - !progressImage?.target || - !placeholder || - progressImage.target.queueItemId !== placeholder.queueItemId || - progressImage.target.itemIndex !== placeholder.itemIndex - ) { - return null; - } - - return progressImage; -}; - const selectGenerateRecallValues = createGenerateFormValuesSelector(); /** @@ -212,7 +194,6 @@ export const PreviewWidgetView = ({ region, runtime }: WidgetViewProps) => { const { antialiasProgressImages, showProgressImagesInViewer } = useActiveProjectSelector( (project) => project.settings ); - const progressImage = useProgressImage(); const runningProgressTargets = useActiveProgressTargets(); const followedProgressTargets = useFollowedProgressTargets(); const { account, gallery, notifications, widgets } = useWorkbenchCommands(); @@ -265,7 +246,6 @@ export const PreviewWidgetView = ({ region, runtime }: WidgetViewProps) => { null, [followedProgressTargets, generationSequence.chronologicalSlots, liveGalleryPlaceholders] ); - const matchingProgressImage = getMatchingProgressImage(progressImage, activeGalleryPlaceholder); // Not while a similarity search is active: the grid hides pending items // there entirely, so following the generation would put Preview on a tile // the grid is not showing and, worse, hand the arrows the board listing @@ -654,7 +634,6 @@ export const PreviewWidgetView = ({ region, runtime }: WidgetViewProps) => { filmstripItems={isFilmstripVisible && density !== 'minimal' ? boardItems : null} isLoadingBoard={isLoadingBoard} placeholder={activeGalleryPlaceholder} - progressImage={matchingProgressImage} selectedIndex={navigationCursor} shouldAntialiasProgressImage={antialiasProgressImages} onNext={selectNextItem} @@ -950,7 +929,6 @@ const LivePreview = ({ filmstripItems, isLoadingBoard, placeholder, - progressImage, selectedIndex, shouldAntialiasProgressImage, onNext, @@ -962,13 +940,16 @@ const LivePreview = ({ filmstripItems: GalleryItem[] | null; isLoadingBoard: boolean; placeholder: GalleryQueuePlaceholder; - progressImage: LatestProgressImageSnapshot | null; selectedIndex: number; shouldAntialiasProgressImage: boolean; onNext: () => void; onPrevious: () => void; onSelectItem: (item: GalleryItem) => void; }) => { + // The followed slot's own frame, not the store-wide latest: with two slots + // live (a long video next to a quick image batch) the latest belongs to + // whichever stepped last, and releasing that slot must not blank this one. + const progressImage = useQueueItemProgressImage(placeholder.queueItemId, placeholder.itemIndex); // The previous slot's last frame stands in until this slot produces one of // its own (model load, text encoding) — otherwise a sequential batch drops // to an empty card between items.