fix: check by value for KNOWN_UPSCALERS enum membership

Haidra-Org · Mar 9, 2024 · cc6ebbf · cc6ebbf
1 parent 39653a7
commit cc6ebbf
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py
@@ -194,7 +194,10 @@ def has_upscaler(self) -> bool:
         if len(self.payload.post_processing) == 0:
             return False
 
-        return any(post_processing in KNOWN_UPSCALERS.__members__ for post_processing in self.payload.post_processing)
+        return any(
+            post_processing in KNOWN_UPSCALERS.__members__ or post_processing in KNOWN_UPSCALERS._value2member_map_
+            for post_processing in self.payload.post_processing
+        )
 
     @property
     def has_facefixer(self) -> bool:

diff --git a/tests/ai_horde_api/test_ai_horde_api_models.py b/tests/ai_horde_api/test_ai_horde_api_models.py
@@ -436,6 +436,11 @@ def test_ImageGenerateJobPopResponse() -> None:
         skipped=ImageGenerateJobPopSkippedStatus(),
     )
 
+    assert all(
+        post_processor in KNOWN_UPSCALERS._value2member_map_
+        for post_processor in test_response.payload.post_processing
+    )
+
     test_response = ImageGenerateJobPopResponse(
         id=None,
         ids=[JobID(root=UUID("00000000-0000-0000-0000-000000000000"))],
@@ -447,6 +452,8 @@ def test_ImageGenerateJobPopResponse() -> None:
         skipped=ImageGenerateJobPopSkippedStatus(),
     )
 
+    assert all(post_processor in KNOWN_UPSCALERS for post_processor in test_response.payload.post_processing)
+
 
 def test_ImageGenerateJobPopResponse_hashability() -> None:
     test_response_ids = ImageGenerateJobPopResponse(