fix: update model reference import and adjust model resolution logic

Author: tazlin

horde_sdk/ai_horde_worker/model_meta.py

@@ -3,7 +3,7 @@
 
 from horde_model_reference.meta_consts import MODEL_REFERENCE_CATEGORY, STABLE_DIFFUSION_BASELINE_CATEGORY
 from horde_model_reference.model_reference_manager import ModelReferenceManager
-from horde_model_reference.model_reference_records import StableDiffusion_ModelRecord
+from horde_model_reference.model_reference_records import ImageGenerationModelRecord
 from loguru import logger
 
 from horde_sdk.ai_horde_api.ai_horde_clients import AIHordeAPIManualClient
@@ -176,9 +176,9 @@ def resolve_all_model_names(
         """
         all_model_references = self._model_reference_manager.get_all_model_references()
 
-        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.stable_diffusion]
+        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.image_generation]
 
-        all_models = set(sd_model_references.root.keys()) if sd_model_references is not None else set()
+        all_models = set(sd_model_references.keys()) if sd_model_references is not None else set()
 
         if not ignore_large_models_env_var:
             all_models = self.remove_large_models(all_models)
@@ -198,21 +198,21 @@ def _resolve_sfw_nsfw_model_names(self, nsfw: bool) -> set[str]:
         """
         all_model_references = self._model_reference_manager.get_all_model_references()
 
-        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.stable_diffusion]
+        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.image_generation]
 
         found_models: set[str] = set()
 
         if sd_model_references is None:
             logger.error("No stable diffusion models found in model reference.")
             return found_models
 
-        for model in sd_model_references.root.values():
-            if not isinstance(model, StableDiffusion_ModelRecord):
-                logger.error(f"Model {model} is not a StableDiffusion_ModelRecord")
+        for model_name, model in sd_model_references.items():
+            if not isinstance(model, ImageGenerationModelRecord):
+                logger.error(f"Model {model_name} is not a ImageGenerationModelRecord")
                 continue
 
             if model.nsfw == nsfw:
-                found_models.add(model.name)
+                found_models.add(model_name)
 
         return found_models
 
@@ -240,21 +240,21 @@ def resolve_all_inpainting_models(self) -> set[str]:
         """
         all_model_references = self._model_reference_manager.get_all_model_references()
 
-        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.stable_diffusion]
+        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.image_generation]
 
         found_models: set[str] = set()
 
         if sd_model_references is None:
             logger.error("No stable diffusion models found in model reference.")
             return found_models
 
-        for model in sd_model_references.root.values():
-            if not isinstance(model, StableDiffusion_ModelRecord):
-                logger.error(f"Model {model} is not a StableDiffusion_ModelRecord")
+        for model_name, model in sd_model_references.items():
+            if not isinstance(model, ImageGenerationModelRecord):
+                logger.error(f"Model {model_name} is not a ImageGenerationModelRecord")
                 continue
 
             if model.inpainting:
-                found_models.add(model.name)
+                found_models.add(model_name)
 
         return found_models
 
@@ -269,21 +269,21 @@ def resolve_all_models_of_baseline(self, baseline: str) -> set[str]:
         """
         all_model_references = self._model_reference_manager.get_all_model_references()
 
-        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.stable_diffusion]
+        sd_model_references = all_model_references[MODEL_REFERENCE_CATEGORY.image_generation]
 
         found_models: set[str] = set()
 
         if sd_model_references is None:
             logger.error("No stable diffusion models found in model reference.")
             return found_models
 
-        for model in sd_model_references.root.values():
-            if not isinstance(model, StableDiffusion_ModelRecord):
-                logger.error(f"Model {model} is not a StableDiffusion_ModelRecord")
+        for model_name, model in sd_model_references.items():
+            if not isinstance(model, ImageGenerationModelRecord):
+                logger.error(f"Model {model_name} is not a ImageGenerationModelRecord")
                 continue
 
             if model.baseline == baseline:
-                found_models.add(model.name)
+                found_models.add(model_name)
 
         return found_models
 

requirements.txt

@@ -1,6 +1,6 @@
-horde_model_reference~=0.9.0
+horde_model_reference>=2.0.0
 
-pydantic==2.9.2
+pydantic>=2.9.2
 requests
 StrEnum
 loguru