diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py index a01bb41..920bd32 100644 --- a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py +++ b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py @@ -70,6 +70,10 @@ class ImageGenerateJobPopSkippedStatus(HordeAPIDataObject): controlnet: int = Field(default=0, ge=0) """How many waiting requests were skipped because they requested a controlnet.""" + def is_empty(self) -> bool: + """Whether or not this object has any non-zero values.""" + return len(self.model_fields_set) == 0 + class ImageGenerateJobPopPayload(ImageGenerateParamMixin): prompt: str | None = None @@ -128,6 +132,9 @@ def validate_id(cls, v: str | JobID) -> JobID | str: @model_validator(mode="after") def ids_present(self) -> ImageGenerateJobPopResponse: """Ensure that either id_ or ids is present.""" + if self.skipped.is_empty() and self.model is None: + return self + if self.id_ is None and len(self.ids) == 0: raise ValueError("Neither id_ nor ids were present in the response.") 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 bfec2ff..55d0e7d 100644 --- a/tests/ai_horde_api/test_ai_horde_api_models.py +++ b/tests/ai_horde_api/test_ai_horde_api_models.py @@ -1,4 +1,5 @@ """Unit tests for AI-Horde API models.""" +import json from uuid import UUID import pytest @@ -604,3 +605,41 @@ def test_AlchemyPopResponse() -> None: container_multiple_forms = {test_alchemy_pop_response_multiple_forms} assert test_alchemy_pop_response_multiple_forms in container_multiple_forms assert test_alchemy_pop_response_multiple_forms_copy in container_multiple_forms + + +def test_ImageGenerateJobPopSkippedStatus() -> None: + testing_skipped_status = ImageGenerateJobPopSkippedStatus() + + assert testing_skipped_status.is_empty() + + +def test_problem_payload() -> None: + json_from_api = """ + { + "payload": { + "ddim_steps": 30, + "n_iter": 1, + "sampler_name": "k_euler_a", + "cfg_scale": 7.5, + "height": 512, + "width": 512, + "karras": false, + "tiling": false, + "hires_fix": false, + "image_is_control": false, + "return_control_map": false + }, + "id": null, + "ids": [], + "skipped": {}, + "model": null, + "source_image": null, + "source_processing": "img2img", + "source_mask": null, + "r2_upload": null, + "r2_uploads": null + }""" + + problem_payload = json.loads(json_from_api) + + ImageGenerateJobPopResponse.model_validate(problem_payload)