fix: resolve 'all' models when asked (#321)

* fix: resolve 'all' models when asked

The dynamic I introduced with `AI_HORDE_MODEL_META_LARGE_MODELS` was confusing. Unfortunately, this is undoubtably going to create problems when workers start offering flux when they didn't intend to, but that's the cost of progress.

* style: fix

Author: tazlin

horde_sdk/ai_horde_worker/model_meta.py

@@ -57,7 +57,7 @@ def resolve_meta_instructions(
         for possible_instruction in possible_meta_instructions:
             # If the instruction is to load all models, return all model names
             if ImageModelLoadResolver.meta_instruction_regex_match(MetaInstruction.ALL_REGEX, possible_instruction):
-                return self.resolve_all_model_names()
+                return self.remove_large_models(set(self.resolve_all_model_names()))
 
             # If the instruction is to load the top N models, add the top N model names
             top_n_matches = ImageModelLoadResolver.meta_instruction_regex_match(
@@ -161,9 +161,16 @@ def remove_large_models(self, models: set[str]) -> set[str]:
             models = models - cascade_models - flux_models
         return models
 
-    def resolve_all_model_names(self) -> set[str]:
+    def resolve_all_model_names(
+        self,
+        ignore_large_models_env_var: bool = True,
+    ) -> set[str]:
         """Get the names of all models defined in the model reference.
 
+        Args:
+            ignore_large_models_env_var (bool): A boolean representing whether to ignore the environment variable for
+                large models, effectively returning all models regardless of `AI_HORDE_MODEL_META_LARGE_MODELS`.
+
         Returns:
             A set of strings representing the names of all models.
         """
@@ -173,7 +180,8 @@ def resolve_all_model_names(self) -> set[str]:
 
         all_models = set(sd_model_references.root.keys()) if sd_model_references is not None else set()
 
-        all_models = self.remove_large_models(all_models)
+        if not ignore_large_models_env_var:
+            all_models = self.remove_large_models(all_models)
 
         if not all_models:
             logger.error("No stable diffusion models found in model reference.")

tests/ai_horde_worker/test_model_meta_api_calls.py

@@ -31,13 +31,33 @@ def test_image_model_load_resolver_all(image_model_load_resolver: ImageModelLoad
 
     import os
 
-    os.environ["AI_HORDE_MODEL_META_LARGE_MODELS"] = "true"
-
-    all_model_names_with_large = image_model_load_resolver.resolve_all_model_names()
+    os.environ["AI_HORDE_MODEL_META_LARGE_MODELS"] = "1"
+    all_model_names_with_large = image_model_load_resolver.resolve_all_model_names(ignore_large_models_env_var=False)
 
     del os.environ["AI_HORDE_MODEL_META_LARGE_MODELS"]
 
-    assert len(all_model_names_with_large) > len(all_model_names)
+    assert len(all_model_names_with_large) == len(all_model_names)
+
+
+def test_image_model_load_resolver_all_without_large(image_model_load_resolver: ImageModelLoadResolver) -> None:
+    import os
+
+    all_model_names = image_model_load_resolver.resolve_all_model_names()
+
+    assert len(all_model_names) > 0
+
+    stored_value = os.environ.get("AI_HORDE_MODEL_META_LARGE_MODELS")
+
+    if "AI_HORDE_MODEL_META_LARGE_MODELS" in os.environ:
+        del os.environ["AI_HORDE_MODEL_META_LARGE_MODELS"]
+
+    all_model_names_without_large = image_model_load_resolver.resolve_all_model_names(
+        ignore_large_models_env_var=False,
+    )
+
+    if stored_value is not None:
+        os.environ["AI_HORDE_MODEL_META_LARGE_MODELS"] = stored_value
+    assert len(all_model_names) > len(all_model_names_without_large)
 
 
 def test_image_model_load_resolver_top_n(