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/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/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/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/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/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. 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/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()), 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: