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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde_model_reference
-Version: 4.0.2
+Version: 4.0.4
 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.4}/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.4'
+__version_tuple__ = version_tuple = (4, 0, 4)
 
-__commit_id__ = commit_id = 'gc8def96a6'
+__commit_id__ = commit_id = 'g5e3a18922'

{horde_model_reference-4.0.2 → horde_model_reference-4.0.4}/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,11 @@ 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:
+            if reference is not None:
+                safe_references[category] = reference
+            elif category not in categories_managed_elsewhere:
                 missing_references.append(category)
                 safe_references[category] = {}
-            elif reference is not None:
-                safe_references[category] = reference
 
         if missing_references:
             logger.error(f"Missing model references for categories: {missing_references}")
@@ -1122,13 +1122,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 +1198,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.4/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.4
 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.4}/tests/helpers.py

@@ -1,4 +1,5 @@
 from horde_model_reference import MODEL_REFERENCE_CATEGORY
+from horde_model_reference.meta_consts import categories_managed_elsewhere
 from horde_model_reference.model_reference_records import GenericModelRecord
 
 
@@ -13,10 +14,13 @@ def verify_model_references_structure(
     Raises:
         AssertionError: If validation fails.
     """
-    # Should have all categories
-    assert len(all_references) == len(MODEL_REFERENCE_CATEGORY)
+    # Should have all categories except those managed elsewhere
+    expected_count = len(MODEL_REFERENCE_CATEGORY) - len(categories_managed_elsewhere)
+    assert len(all_references) == expected_count
 
     for category in MODEL_REFERENCE_CATEGORY:
+        if category in categories_managed_elsewhere:
+            continue
         assert category in all_references, f"Category {category} is missing"
 
         references = all_references[category]

{horde_model_reference-4.0.2 → horde_model_reference-4.0.4}/tests/test_model_reference_manager.py

@@ -16,6 +16,7 @@ from horde_model_reference.meta_consts import (
     MODEL_PURPOSE,
     MODEL_REFERENCE_CATEGORY,
     ModelClassification,
+    categories_managed_elsewhere,
 )
 from horde_model_reference.model_reference_manager import ModelReferenceManager
 from horde_model_reference.model_reference_records import GenericModelRecord, ImageGenerationModelRecord
@@ -127,13 +128,13 @@ def test_manager(
     all_references = model_reference_manager.get_all_model_references(overwrite_existing=False)
 
     assert len(all_references) > 0
-    assert all(ref_cat in all_references for ref_cat in MODEL_REFERENCE_CATEGORY)
+    assert all(ref_cat in all_references for ref_cat in MODEL_REFERENCE_CATEGORY if ref_cat not in categories_managed_elsewhere if ref_cat not in categories_managed_elsewhere)
     assert all(ref is not None for ref in all_references.values())
 
     verify_model_references_structure(all_references)
 
     cached_references = model_reference_manager.get_all_model_references_or_none(overwrite_existing=False)
-    assert len(cached_references) == len(all_references)
+    assert len(cached_references) == len(all_references) + len(categories_managed_elsewhere)
 
 
 @pytest.mark.integration
@@ -168,12 +169,12 @@ async def test_manager_async(
     all_references = model_reference_manager.get_all_model_references(overwrite_existing=False)
 
     assert len(all_references) > 0
-    assert all(ref_cat in all_references for ref_cat in MODEL_REFERENCE_CATEGORY)
+    assert all(ref_cat in all_references for ref_cat in MODEL_REFERENCE_CATEGORY if ref_cat not in categories_managed_elsewhere)
 
     verify_model_references_structure(all_references)
 
     cached_references = model_reference_manager.get_all_model_references_or_none(overwrite_existing=False)
-    assert len(cached_references) == len(all_references)
+    assert len(cached_references) == len(all_references) + len(categories_managed_elsewhere)
 
 
 class TestSingleton:
@@ -513,7 +514,7 @@ class TestAsyncErgonomics:
         )
 
         first_result = await manager.get_all_model_references_async()
-        assert len(first_result) == len(MODEL_REFERENCE_CATEGORY)
+        assert len(first_result) == len(MODEL_REFERENCE_CATEGORY) - len(categories_managed_elsewhere)
         assert len(backend.async_calls) == 1
 
         second_result = await manager.get_all_model_references_async()

{horde_model_reference-4.0.2 → horde_model_reference-4.0.4}/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: