fix: support model_state for model load meta (#313)

* fix: support model_state for model load meta This fixes the behavior of workers attempting to load recently removed (and therefore previously known) models when using the `TOP` meta model load command * style: fix * style: fix
Haidra-Org · Dec 24, 2024 · e180e41 · e180e41
1 parent 9b7fe3d
commit e180e41
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/horde_sdk/ai_horde_api/ai_horde_clients.py b/horde_sdk/ai_horde_api/ai_horde_clients.py
@@ -61,7 +61,7 @@
     TextStatsModelsTotalResponse,
 )
 from horde_sdk.ai_horde_api.apimodels.base import BaseAIHordeRequest, JobRequestMixin
-from horde_sdk.ai_horde_api.consts import GENERATION_MAX_LIFE, PROGRESS_STATE
+from horde_sdk.ai_horde_api.consts import GENERATION_MAX_LIFE, MODEL_STATE, PROGRESS_STATE
 from horde_sdk.ai_horde_api.endpoints import AI_HORDE_BASE_URL
 from horde_sdk.ai_horde_api.exceptions import AIHordeImageValidationError, AIHordeRequestError
 from horde_sdk.ai_horde_api.fields import JobID, WorkerID
@@ -1121,14 +1121,20 @@ def image_stats_totals(
 
     def image_stats_models(
         self,
+        model_state: str | MODEL_STATE = MODEL_STATE.known,
     ) -> ImageStatsModelsResponse:
         """Get the stats for images by model.
 
         Returns:
             ImageStatsModelsResponse: The response from the API.
         """
+        model_state = MODEL_STATE(model_state)
+
         with AIHordeAPIClientSession() as horde_session:
-            response = horde_session.submit_request(ImageStatsModelsRequest(), ImageStatsModelsResponse)
+            response = horde_session.submit_request(
+                ImageStatsModelsRequest(model_state=model_state),
+                ImageStatsModelsResponse,
+            )
 
             if isinstance(response, RequestErrorResponse):
                 raise AIHordeRequestError(response)
@@ -1827,15 +1833,18 @@ async def image_stats_totals(
 
     async def image_stats_models(
         self,
+        model_state: str | MODEL_STATE = MODEL_STATE.known,
     ) -> ImageStatsModelsResponse:
         """Get the stats for images by model.
 
         Returns:
             ImageStatsModelsResponse: The response from the API.
         """
+        model_state = MODEL_STATE(model_state)
+
         if self._horde_client_session is not None:
             response = await self._horde_client_session.submit_request(
-                ImageStatsModelsRequest(),
+                ImageStatsModelsRequest(model_state=model_state),
                 ImageStatsModelsResponse,
             )
         else:

diff --git a/horde_sdk/ai_horde_worker/model_meta.py b/horde_sdk/ai_horde_worker/model_meta.py
@@ -8,6 +8,7 @@
 
 from horde_sdk.ai_horde_api.ai_horde_clients import AIHordeAPIManualClient
 from horde_sdk.ai_horde_api.apimodels import ImageStatsModelsRequest, ImageStatsModelsResponse, StatsModelsTimeframe
+from horde_sdk.ai_horde_api.consts import MODEL_STATE
 from horde_sdk.ai_horde_worker.bridge_data import MetaInstruction
 from horde_sdk.generic_api.apimodels import RequestErrorResponse
 
@@ -40,7 +41,10 @@ def resolve_meta_instructions(
             A set of strings representing the names of models to load.
         """
         # Get model stats from the API
-        stats_response = client.submit_request(ImageStatsModelsRequest(), ImageStatsModelsResponse)
+        stats_response = client.submit_request(
+            ImageStatsModelsRequest(model_state=MODEL_STATE.known),
+            ImageStatsModelsResponse,
+        )
         if isinstance(stats_response, RequestErrorResponse):
             raise Exception(f"Error getting stats for models: {stats_response.message}")