From 930d4f6a698f6fc59347730e6746db406381e7cd Mon Sep 17 00:00:00 2001 From: tazlin Date: Sat, 9 Mar 2024 14:09:23 -0500 Subject: [PATCH] fix: check by value for KNOWN_UPSCALERS enum membership --- horde_sdk/ai_horde_api/apimodels/generate/_pop.py | 5 ++++- tests/ai_horde_api/test_ai_horde_api_models.py | 7 +++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py index fc25c27..af2d8f0 100644 --- 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 index 8ab985e..de134eb 100644 --- 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(