github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

use queue room/subscription instead of separate model loading room

Browse Source
This commit is contained in:
Mary Hipp
2024-11-06 16:30:02 -05:00
parent c63fe5e9bb
commit 28864f6d7f
2 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
invokeai/app/api/sockets.py
View File

@@ -31,6 +31,7 @@ from invokeai.app.services.events.events_common import (
ModelInstallErrorEvent,
ModelInstallStartedEvent,
ModelLoadCompleteEvent,
ModelLoadEventBase,
ModelLoadStartedEvent,
QueueClearedEvent,
QueueEventBase,
@@ -134,7 +135,7 @@ class SocketIO:
async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]):
await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"), room=event[1].queue_id)
async def _handle_model_load_event(self, event: FastAPIEvent[ModelEventBase]) -> None:
async def _handle_model_load_event(self, event: FastAPIEvent[ModelLoadEventBase]) -> None:
await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"), room=event[1].queue_id)
async def _handle_model_event(self, event: FastAPIEvent[ModelEventBase | DownloadEventBase]) -> None:

11
invokeai/app/services/events/events_common.py
View File

@@ -383,20 +383,19 @@ class DownloadErrorEvent(DownloadEventBase):
return cls(source=str(job.source), error_type=job.error_type, error=job.error)
class ModelEventBase(EventBase):
class ModelLoadEventBase(EventBase):
"""Base class for queue events"""
queue_id: str = Field(description="The ID of the queue")
@payload_schema.register
class ModelLoadStartedEvent(ModelEventBase):
class ModelLoadStartedEvent(ModelLoadEventBase):
"""Event model for model_load_started"""
__event_name__ = "model_load_started"
config: AnyModelConfig = Field(description="The model's config")
queue_id: str = Field(description="Queue ID to emit to")
submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
@classmethod
@@ -405,13 +404,12 @@ class ModelLoadStartedEvent(ModelEventBase):
@payload_schema.register
class ModelLoadCompleteEvent(ModelEventBase):
class ModelLoadCompleteEvent(ModelLoadEventBase):
"""Event model for model_load_complete"""
__event_name__ = "model_load_complete"
config: AnyModelConfig = Field(description="The model's config")
queue_id: str = Field(description="Queue ID to emit to")
submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
@classmethod
@@ -419,6 +417,9 @@ class ModelLoadCompleteEvent(ModelEventBase):
return cls(config=config, queue_id=queue_id, submodel_type=submodel_type)
class ModelEventBase(EventBase):
"""Base class for model events"""
@payload_schema.register
class ModelInstallDownloadStartedEvent(ModelEventBase):
"""Event model for model_install_download_started"""