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(