horde-model-reference 4.0.2__tar.gz → 4.0.3__tar.gz

{horde_model_reference-4.0.2/src/horde_model_reference.egg-info → horde_model_reference-4.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde_model_reference
-Version: 4.0.2
+Version: 4.0.3
 Summary: A helper library providing a way to work with the lists of generation models, utility models, and any other related files required for the AI-Horde ecosystem.
 Author-email: tazlin <tazlin.on.github@gmail.com>, db0 <mail@dbzer0.com>
 License:                     GNU AFFERO GENERAL PUBLIC LICENSE

{horde_model_reference-4.0.2 → horde_model_reference-4.0.3}/src/horde_model_reference/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '4.0.2'
-__version_tuple__ = version_tuple = (4, 0, 2)
+__version__ = version = '4.0.3'
+__version_tuple__ = version_tuple = (4, 0, 3)
 
-__commit_id__ = commit_id = 'gc8def96a6'
+__commit_id__ = commit_id = 'g98fd7c136'

{horde_model_reference-4.0.2 → horde_model_reference-4.0.3}/src/horde_model_reference/model_reference_manager.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Generator, Iterable
+from collections.abc import Awaitable, Generator, Iterable, Mapping
 from pathlib import Path
 from threading import RLock
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, overload
@@ -962,11 +962,12 @@ class ModelReferenceManager:
         safe_references: dict[MODEL_REFERENCE_CATEGORY, dict[str, GenericModelRecord]] = {}
         missing_references: list[MODEL_REFERENCE_CATEGORY] = []
         for category, reference in all_references.items():
-            if category not in categories_managed_elsewhere:
-                missing_references.append(category)
-                safe_references[category] = {}
-            elif reference is not None:
+            if reference is not None:
                 safe_references[category] = reference
+            else:
+                safe_references[category] = {}
+                if category not in categories_managed_elsewhere:
+                    missing_references.append(category)
 
         if missing_references:
             logger.error(f"Missing model references for categories: {missing_references}")
@@ -1122,13 +1123,71 @@ class ModelReferenceManager:
             httpx_client=httpx_client,
         )
 
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.audio_generation],
+        overwrite_existing: bool = False,
+    ) -> dict[str, AudioGenerationModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.blip],
+        overwrite_existing: bool = False,
+    ) -> dict[str, BlipModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.clip],
+        overwrite_existing: bool = False,
+    ) -> dict[str, ClipModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.codeformer],
+        overwrite_existing: bool = False,
+    ) -> dict[str, CodeformerModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.controlnet],
+        overwrite_existing: bool = False,
+    ) -> dict[str, ControlNetModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.image_generation],
+        overwrite_existing: bool = False,
+    ) -> dict[str, ImageGenerationModelRecord]: ...
+
+    @overload
+    def get_model_reference(
+        self,
+        category: Literal[MODEL_REFERENCE_CATEGORY.text_generation],
+        overwrite_existing: bool = False,
+    ) -> dict[str, TextGenerationModelRecord]: ...
+
+    @overload
     def get_model_reference(
         self,
         category: MODEL_REFERENCE_CATEGORY,
         overwrite_existing: bool = False,
         *,
         source: SourceSelector = HORDE_SOURCE_ID,
-    ) -> dict[str, GenericModelRecord]:
+    ) -> dict[str, GenericModelRecord]: ...
+
+    def get_model_reference(
+        self,
+        category: MODEL_REFERENCE_CATEGORY = MODEL_REFERENCE_CATEGORY.image_generation,
+        overwrite_existing: bool = False,
+        *,
+        source: SourceSelector = HORDE_SOURCE_ID,
+    ) -> Mapping[str, GenericModelRecord]:
         """Return the model reference object for a specific category.
 
         Raises an exception if the model reference could not be found or parsed.
@@ -1140,7 +1199,7 @@ class ModelReferenceManager:
             source: Which source(s) to read from. See :meth:`get_model_reference_or_none`.
 
         Returns:
-            dict[str, GenericModelRecord]: The model reference object for the category.
+            Mapping[str, GenericModelRecord]: The model reference object for the category.
 
         """
         model_reference = self.get_model_reference_or_none(

{horde_model_reference-4.0.2 → horde_model_reference-4.0.3/src/horde_model_reference.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde_model_reference
-Version: 4.0.2
+Version: 4.0.3
 Summary: A helper library providing a way to work with the lists of generation models, utility models, and any other related files required for the AI-Horde ecosystem.
 Author-email: tazlin <tazlin.on.github@gmail.com>, db0 <mail@dbzer0.com>
 License:                     GNU AFFERO GENERAL PUBLIC LICENSE

{horde_model_reference-4.0.2 → horde_model_reference-4.0.3}/tests/test_scripts.py

@@ -1,5 +1,6 @@
 import os
 
+import horde_model_reference.meta_consts
 from horde_model_reference import ModelReferenceManager, ReplicateMode, horde_model_reference_paths
 from horde_model_reference.backends.github_backend import GitHubBackend
 from horde_model_reference.legacy.validate_sd import validate_legacy_stable_diffusion_db
@@ -9,7 +10,9 @@ from horde_model_reference.meta_consts import MODEL_REFERENCE_CATEGORY
 def test_download_all_model_references(model_reference_manager: ModelReferenceManager) -> None:
     """Test downloading all model references script."""
     download_models = model_reference_manager.get_all_model_references(overwrite_existing=True)
-    assert len(download_models) == len(MODEL_REFERENCE_CATEGORY.__members__)
+    assert len(download_models) == (
+        len(MODEL_REFERENCE_CATEGORY.__members__) - len(horde_model_reference.meta_consts.categories_managed_elsewhere)
+    )
 
 
 def test_validate_stable_diffusion_model_reference() -> None: