From 1e958d29051e2d7a627b605de634486731d69feb Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 10:28:36 -0400 Subject: [PATCH 01/11] fix: default (or no-init) certain lists as `None` instead of `[]` --- horde_sdk/ai_horde_api/apimodels/base.py | 10 +++++----- horde_sdk/ai_horde_api/apimodels/generate/_pop.py | 6 ++---- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/horde_sdk/ai_horde_api/apimodels/base.py b/horde_sdk/ai_horde_api/apimodels/base.py index 0d410ea..12f5943 100644 --- a/horde_sdk/ai_horde_api/apimodels/base.py +++ b/horde_sdk/ai_horde_api/apimodels/base.py @@ -219,7 +219,7 @@ class ImageGenerateParamMixin(HordeAPIDataObject): karras: bool = True """Set to True if you want to use the Karras scheduling.""" tiling: bool = False - """Deprecated.""" + """Set to True if you want to use seamless tiling.""" hires_fix: bool = False """Set to True if you want to use the hires fix.""" hires_fix_denoising_strength: float | None = Field(default=None, ge=0, le=1) @@ -234,17 +234,17 @@ class ImageGenerateParamMixin(HordeAPIDataObject): """Set to True if you want the ControlNet map returned instead of a generated image.""" facefixer_strength: float | None = Field(default=None, ge=0, le=1) """The strength of the facefixer model.""" - loras: list[LorasPayloadEntry] = Field(default_factory=list) + loras: list[LorasPayloadEntry] | None = None """A list of lora parameters to use.""" - tis: list[TIPayloadEntry] = Field(default_factory=list) + tis: list[TIPayloadEntry] | None = None """A list of textual inversion (embedding) parameters to use.""" - extra_texts: list[ExtraTextEntry] = Field(default_factory=list) + extra_texts: list[ExtraTextEntry] | None = None """A list of extra texts and prompts to use in the comfyUI workflow.""" workflow: str | KNOWN_WORKFLOWS | None = None """The specific comfyUI workflow to use.""" transparent: bool | None = None """When true, will generate an image with a transparent background""" - special: dict[Any, Any] = Field(default_factory=dict) + special: dict[Any, Any] | None = None """Reserved for future use.""" use_nsfw_censor: bool = False """If the request is SFW, and the worker accidentally generates NSFW, it will send back a censored image.""" diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py index 03c5514..a004522 100644 --- a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py +++ b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py @@ -418,11 +418,9 @@ class PopInput(HordeAPIObject): max_length=1000, ) """The worker name, version and website.""" - models: list[str] | None = None + models: list[str] """The models this worker can generate.""" - name: str | None = Field( - None, - ) + name: str """The Name of the Worker.""" nsfw: bool | None = Field( False, From fef64b70ae416cf73e476dbf50e31452aeae2908 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 11:46:44 -0400 Subject: [PATCH 02/11] tests: `AI_HORDE_TESTING` and `HORDE_SDK_TESTING` I am putting these in place for future use --- .github/workflows/maintests.yml | 3 ++- .github/workflows/prtests.yml | 1 + tests/conftest.py | 18 +++++++++++++++++- tox.ini | 2 ++ 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/.github/workflows/maintests.yml b/.github/workflows/maintests.yml index fd458cb..51edb84 100644 --- a/.github/workflows/maintests.yml +++ b/.github/workflows/maintests.yml @@ -19,8 +19,9 @@ jobs: build: env: AIWORKER_CACHE_HOME: ${{ github.workspace }}/.cache - HORDE_MODEL_REFERENCE_MAKE_FOLDERS: 1 TESTS_ONGOING: 1 + HORDE_SDK_TESTING: 1 + HORDE_MODEL_REFERENCE_MAKE_FOLDERS: 1 runs-on: ubuntu-latest strategy: matrix: diff --git a/.github/workflows/prtests.yml b/.github/workflows/prtests.yml index 38990ce..144c384 100644 --- a/.github/workflows/prtests.yml +++ b/.github/workflows/prtests.yml @@ -23,6 +23,7 @@ jobs: env: AIWORKER_CACHE_HOME: ${{ github.workspace }}/.cache TESTS_ONGOING: 1 + HORDE_SDK_TESTING: 1 HORDE_MODEL_REFERENCE_MAKE_FOLDERS: 1 runs-on: ubuntu-latest strategy: diff --git a/tests/conftest.py b/tests/conftest.py index d29c349..dda7e2e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ import os import pathlib +from loguru import logger import pytest os.environ["TESTS_ONGOING"] = "1" @@ -15,6 +16,21 @@ def check_tests_ongoing_env_var() -> None: """Checks that the TESTS_ONGOING environment variable is set.""" assert os.getenv("TESTS_ONGOING", None) is not None, "TESTS_ONGOING environment variable not set" + AI_HORDE_TESTING = os.getenv("AI_HORDE_TESTING", None) + HORDE_SDK_TESTING = os.getenv("HORDE_SDK_TESTING", None) + if AI_HORDE_TESTING is None and HORDE_SDK_TESTING is None: + logger.warning( + "Neither AI_HORDE_TESTING nor HORDE_SDK_TESTING environment variables are set. " + "Is this a local development test run? If so, set AI_HORDE_TESTING=1 or HORDE_SDK_TESTING=1 to suppress " + "this warning", + ) + + if AI_HORDE_TESTING is not None: + logger.info("AI_HORDE_TESTING environment variable set.") + + if HORDE_SDK_TESTING is not None: + logger.info("HORDE_SDK_TESTING environment variable set.") + @pytest.fixture(scope="session") def ai_horde_api_key() -> str: @@ -30,7 +46,7 @@ def simple_image_gen_request(ai_horde_api_key: str) -> ImageGenerateAsyncRequest prompt="a cat in a hat", models=["Deliberate"], params=ImageGenerationInputPayload( - steps=1, + steps=5, n=1, ), ) diff --git a/tox.ini b/tox.ini index 0d0db14..cb3fe43 100644 --- a/tox.ini +++ b/tox.ini @@ -14,6 +14,8 @@ skip_empty = True description = base evironment passenv = AIWORKER_CACHE_HOME + HORDE_SDK_TESTING + AI_HORDE_TESTING TESTS_ONGOING [testenv:pre-commit] From 2ada3802e2442602cc485f0f8ca046f5588c0515 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 11:47:13 -0400 Subject: [PATCH 03/11] tests: `api_side_ci` mark and intial worker test --- pyproject.toml | 4 +- .../test_ai_worker_roundtrip_api_calls.py | 150 ++++++++++++++++++ tox.ini | 4 +- 3 files changed, 155 insertions(+), 3 deletions(-) create mode 100644 tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py diff --git a/pyproject.toml b/pyproject.toml index a9d3523..74a6393 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,7 +128,9 @@ exclude = [ concurrency = ["gevent"] [tool.pytest.ini_options] +# You can use `and`, `or`, `not` and parentheses to filter with the `-m` option markers = [ - # "slow: marks tests as slow (deselect with '-m \"not slow\"')", + "slow: marks tests as slow (deselect with '-m \"not slow\"')", "object_verify: marks tests that verify the API object structure and layout", + "api_side_ci: indicates that the test is intended to run during CI for the API", ] diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py new file mode 100644 index 0000000..6958b32 --- /dev/null +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -0,0 +1,150 @@ +import asyncio +import base64 +from typing import Any + +import aiohttp +import PIL.Image +import pytest +from loguru import logger +import yarl + +from horde_sdk.ai_horde_api.ai_horde_clients import ( + AIHordeAPIAsyncClientSession, + AIHordeAPIAsyncSimpleClient, + AIHordeAPISimpleClient, +) +from horde_sdk.ai_horde_api.apimodels import ( + KNOWN_ALCHEMY_TYPES, + AlchemyAsyncRequest, + AlchemyAsyncRequestFormItem, + AlchemyStatusResponse, + ImageGenerateAsyncRequest, + ImageGenerateAsyncResponse, + ImageGenerateCheckResponse, + ImageGenerateJobPopPayload, + ImageGenerateJobPopRequest, + ImageGenerateJobPopResponse, + ImageGenerateJobPopSkippedStatus, + ImageGenerateStatusResponse, + ImageGenerationInputPayload, + ImageGenerationJobSubmitRequest, + JobSubmitResponse, + LorasPayloadEntry, +) +from horde_sdk.ai_horde_api.apimodels.base import ExtraSourceImageEntry +from horde_sdk.ai_horde_api.consts import ( + GENERATION_STATE, + KNOWN_FACEFIXERS, + KNOWN_MISC_POST_PROCESSORS, + KNOWN_SOURCE_PROCESSING, + KNOWN_UPSCALERS, + POST_PROCESSOR_ORDER_TYPE, +) +from horde_sdk.ai_horde_api.fields import JobID +from horde_sdk.generic_api.apimodels import RequestErrorResponse + + +class TestImageWorkerRoundtrip: + async def fake_worker( + self, + aiohttp_session: aiohttp.ClientSession, + horde_client_session: AIHordeAPIAsyncClientSession, + image_gen_request: ImageGenerateAsyncRequest, + ) -> None: + assert image_gen_request.params is not None + + await asyncio.sleep(5) + + effective_resolution = image_gen_request.params.width * image_gen_request.params.height * 2 + + job_pop_request = ImageGenerateJobPopRequest( + name="fake CI worker", + bridge_agent="AI Horde Worker reGen:8.0.1-citests:https://github.com/Haidra-Org/horde-worker-reGen", + max_pixels=effective_resolution, + models=image_gen_request.models, + ) + + job_pop_response = await horde_client_session.submit_request( + job_pop_request, + job_pop_request.get_default_success_response_type(), + ) + + assert isinstance(job_pop_response, ImageGenerateJobPopResponse) + logger.info(f"{job_pop_response.log_safe_model_dump()}") + + assert job_pop_response.ids_present + + # We're going to send a blank image base64 encoded + fake_image = PIL.Image.new( + "RGB", + (image_gen_request.params.width, image_gen_request.params.height), + (255, 255, 255), + ) + + fake_image_bytes = fake_image.tobytes() + + r2_url = job_pop_response.r2_upload + + assert r2_url is not None + + async with aiohttp_session.put( + yarl.URL(r2_url, encoded=True), + data=fake_image_bytes, + skip_auto_headers=["content-type"], + timeout=aiohttp.ClientTimeout(total=10), + ) as response: + assert response.status == 200 + + assert job_pop_response.ids is not None + assert len(job_pop_response.ids) == 1 + + job_submit_request = ImageGenerationJobSubmitRequest( + id=job_pop_response.ids[0], + state=GENERATION_STATE.ok, + generation="R2", + seed="1312", + ) + + job_submit_response = await horde_client_session.submit_request( + job_submit_request, + job_submit_request.get_default_success_response_type(), + ) + + assert isinstance(job_submit_response, JobSubmitResponse) + + assert job_submit_response.reward is not None and job_submit_response.reward > 0 + + @pytest.mark.api_side_ci + @pytest.mark.asyncio + async def test_basic_image_roundtrip(self, simple_image_gen_request: ImageGenerateAsyncRequest) -> None: + aiohttp_session = aiohttp.ClientSession() + horde_client_session = AIHordeAPIAsyncClientSession(aiohttp_session) + + async with aiohttp_session, horde_client_session: + simple_client = AIHordeAPIAsyncSimpleClient(horde_client_session=horde_client_session) + + image_gen_task = asyncio.create_task(simple_client.image_generate_request(simple_image_gen_request)) + + fake_worker_task = asyncio.create_task( + self.fake_worker( + aiohttp_session, + horde_client_session, + simple_image_gen_request, + ), + ) + + await asyncio.gather(image_gen_task, fake_worker_task) + + image_gen_response, job_id = image_gen_task.result() + + assert isinstance(image_gen_response, ImageGenerateStatusResponse) + assert isinstance(job_id, JobID) + + assert len(image_gen_response.generations) == 1 + + generation = image_gen_response.generations[0] + assert generation.seed == "1312" + assert generation.img is not None + assert not generation.gen_metadata + + assert generation.censored is False diff --git a/tox.ini b/tox.ini index cb3fe43..ac5c805 100644 --- a/tox.ini +++ b/tox.ini @@ -35,7 +35,7 @@ deps = requests -r requirements.txt commands = - pytest tests {posargs} --cov + pytest tests {posargs} --cov -m "not api_side_ci" [testenv:tests-no-api-calls] @@ -50,4 +50,4 @@ deps = requests -r requirements.txt commands = - pytest tests {posargs} --ignore-glob=*api_calls.py --cov + pytest tests {posargs} --ignore-glob=*api_calls.py -m "not api_side_ci" From 5ebf090abb627b9f0365733f87bf1b3509178f92 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 11:59:42 -0400 Subject: [PATCH 04/11] tests: have fake worker check in first --- .../test_ai_worker_roundtrip_api_calls.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index 6958b32..5b0b1ab 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -45,6 +45,36 @@ class TestImageWorkerRoundtrip: + async def fake_worker_checkin( + self, + aiohttp_session: aiohttp.ClientSession, + horde_client_session: AIHordeAPIAsyncClientSession, + image_gen_request: ImageGenerateAsyncRequest, + ) -> None: + assert image_gen_request.params is not None + + effective_resolution = image_gen_request.params.width * image_gen_request.params.height + + job_pop_request = ImageGenerateJobPopRequest( + name="fake CI worker", + bridge_agent="AI Horde Worker reGen:8.0.1-citests:https://github.com/Haidra-Org/horde-worker-reGen", + max_pixels=effective_resolution, + models=image_gen_request.models, + ) + + job_pop_response = await horde_client_session.submit_request( + job_pop_request, + job_pop_request.get_default_success_response_type(), + ) + + assert isinstance(job_pop_response, ImageGenerateJobPopResponse) + logger.info(f"{job_pop_response.log_safe_model_dump()}") + + assert not job_pop_response.ids_present + assert job_pop_response.skipped is not None + + logger.info(f"Checked in as fake worker ({effective_resolution}): {job_pop_response.skipped}") + async def fake_worker( self, aiohttp_session: aiohttp.ClientSession, @@ -55,7 +85,7 @@ async def fake_worker( await asyncio.sleep(5) - effective_resolution = image_gen_request.params.width * image_gen_request.params.height * 2 + effective_resolution = image_gen_request.params.width * image_gen_request.params.height job_pop_request = ImageGenerateJobPopRequest( name="fake CI worker", @@ -123,6 +153,8 @@ async def test_basic_image_roundtrip(self, simple_image_gen_request: ImageGenera async with aiohttp_session, horde_client_session: simple_client = AIHordeAPIAsyncSimpleClient(horde_client_session=horde_client_session) + await self.fake_worker_checkin(aiohttp_session, horde_client_session, simple_image_gen_request) + image_gen_task = asyncio.create_task(simple_client.image_generate_request(simple_image_gen_request)) fake_worker_task = asyncio.create_task( From 64efb55858a098754009e6c42eeed357bdf55df6 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 12:10:40 -0400 Subject: [PATCH 05/11] tests: retry fake worker job pop --- .../test_ai_worker_roundtrip_api_calls.py | 88 +++++++++++-------- 1 file changed, 49 insertions(+), 39 deletions(-) diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index 5b0b1ab..1050c18 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -83,8 +83,6 @@ async def fake_worker( ) -> None: assert image_gen_request.params is not None - await asyncio.sleep(5) - effective_resolution = image_gen_request.params.width * image_gen_request.params.height job_pop_request = ImageGenerateJobPopRequest( @@ -94,55 +92,67 @@ async def fake_worker( models=image_gen_request.models, ) - job_pop_response = await horde_client_session.submit_request( - job_pop_request, - job_pop_request.get_default_success_response_type(), - ) + max_tries = 5 + try_wait = 5 + current_try = 0 - assert isinstance(job_pop_response, ImageGenerateJobPopResponse) - logger.info(f"{job_pop_response.log_safe_model_dump()}") + while True: + job_pop_response = await horde_client_session.submit_request( + job_pop_request, + job_pop_request.get_default_success_response_type(), + ) - assert job_pop_response.ids_present + assert isinstance(job_pop_response, ImageGenerateJobPopResponse) + logger.info(f"{job_pop_response.log_safe_model_dump()}") - # We're going to send a blank image base64 encoded - fake_image = PIL.Image.new( - "RGB", - (image_gen_request.params.width, image_gen_request.params.height), - (255, 255, 255), - ) + if not job_pop_response.ids_present: + if current_try >= max_tries: + raise RuntimeError("Max tries exceeded") - fake_image_bytes = fake_image.tobytes() + logger.info(f"Waiting {try_wait} seconds before retrying") + await asyncio.sleep(try_wait) + current_try += 1 + continue - r2_url = job_pop_response.r2_upload + # We're going to send a blank image base64 encoded + fake_image = PIL.Image.new( + "RGB", + (image_gen_request.params.width, image_gen_request.params.height), + (255, 255, 255), + ) - assert r2_url is not None + fake_image_bytes = fake_image.tobytes() - async with aiohttp_session.put( - yarl.URL(r2_url, encoded=True), - data=fake_image_bytes, - skip_auto_headers=["content-type"], - timeout=aiohttp.ClientTimeout(total=10), - ) as response: - assert response.status == 200 + r2_url = job_pop_response.r2_upload - assert job_pop_response.ids is not None - assert len(job_pop_response.ids) == 1 + assert r2_url is not None - job_submit_request = ImageGenerationJobSubmitRequest( - id=job_pop_response.ids[0], - state=GENERATION_STATE.ok, - generation="R2", - seed="1312", - ) + async with aiohttp_session.put( + yarl.URL(r2_url, encoded=True), + data=fake_image_bytes, + skip_auto_headers=["content-type"], + timeout=aiohttp.ClientTimeout(total=10), + ) as response: + assert response.status == 200 - job_submit_response = await horde_client_session.submit_request( - job_submit_request, - job_submit_request.get_default_success_response_type(), - ) + assert job_pop_response.ids is not None + assert len(job_pop_response.ids) == 1 + + job_submit_request = ImageGenerationJobSubmitRequest( + id=job_pop_response.ids[0], + state=GENERATION_STATE.ok, + generation="R2", + seed="1312", + ) + + job_submit_response = await horde_client_session.submit_request( + job_submit_request, + job_submit_request.get_default_success_response_type(), + ) - assert isinstance(job_submit_response, JobSubmitResponse) + assert isinstance(job_submit_response, JobSubmitResponse) - assert job_submit_response.reward is not None and job_submit_response.reward > 0 + assert job_submit_response.reward is not None and job_submit_response.reward > 0 @pytest.mark.api_side_ci @pytest.mark.asyncio From b11f3b2438230f0cb1ae709a21a03d0cae4f1bf6 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 12:20:10 -0400 Subject: [PATCH 06/11] fix: implement `ids_present` as intended --- horde_sdk/ai_horde_api/apimodels/generate/_pop.py | 11 ++++++++++- 1 file changed, 10 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 a004522..7de8309 100644 --- a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py +++ b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py @@ -253,8 +253,15 @@ def validate_id(cls, v: str | JobID) -> JobID | str: return v + _ids_present: bool = False + + @property + def ids_present(self) -> bool: + """Whether or not the IDs are present.""" + return self._ids_present + @model_validator(mode="after") - def ids_present(self) -> ImageGenerateJobPopResponse: + def validate_ids_present(self) -> ImageGenerateJobPopResponse: """Ensure that either id_ or ids is present.""" if self.model is None: if self.skipped.is_empty(): @@ -270,6 +277,8 @@ def ids_present(self) -> ImageGenerateJobPopResponse: logger.debug("Sorting IDs") self.ids.sort() + self._ids_present = True + return self @override From 76c0881bf269a85dc921dddc1af937b168d6d227 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 12:30:19 -0400 Subject: [PATCH 07/11] tests: remember to exit loop for fake worker --- tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index 1050c18..d2509cf 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -151,9 +151,10 @@ async def fake_worker( ) assert isinstance(job_submit_response, JobSubmitResponse) - assert job_submit_response.reward is not None and job_submit_response.reward > 0 + break + @pytest.mark.api_side_ci @pytest.mark.asyncio async def test_basic_image_roundtrip(self, simple_image_gen_request: ImageGenerateAsyncRequest) -> None: From fc05b9bb5f8d9a73bd42aa48bc672316b112f504 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 23 Jun 2024 14:49:50 -0400 Subject: [PATCH 08/11] fix: give fake worker higher resolution --- tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index d2509cf..8394161 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -53,7 +53,7 @@ async def fake_worker_checkin( ) -> None: assert image_gen_request.params is not None - effective_resolution = image_gen_request.params.width * image_gen_request.params.height + effective_resolution = (image_gen_request.params.width * image_gen_request.params.height) * 2 job_pop_request = ImageGenerateJobPopRequest( name="fake CI worker", @@ -83,7 +83,7 @@ async def fake_worker( ) -> None: assert image_gen_request.params is not None - effective_resolution = image_gen_request.params.width * image_gen_request.params.height + effective_resolution = (image_gen_request.params.width * image_gen_request.params.height) * 2 job_pop_request = ImageGenerateJobPopRequest( name="fake CI worker", From 5a0156ebfa588296195869ca26b663400ee98273 Mon Sep 17 00:00:00 2001 From: tazlin Date: Mon, 24 Jun 2024 10:15:33 -0400 Subject: [PATCH 09/11] feat: param to exclude addtl. fields w/ `log_safe_model_dump` --- horde_sdk/generic_api/apimodels.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/horde_sdk/generic_api/apimodels.py b/horde_sdk/generic_api/apimodels.py index 79b5243..377f03a 100644 --- a/horde_sdk/generic_api/apimodels.py +++ b/horde_sdk/generic_api/apimodels.py @@ -79,9 +79,14 @@ def get_extra_fields_to_exclude_from_log(self) -> set[str]: """Return an additional set of fields to exclude from the log_safe_model_dump method.""" return set() - def log_safe_model_dump(self) -> dict[Any, Any]: + def log_safe_model_dump(self, extra_exclude: set[str] | None = None) -> dict[Any, Any]: """Return a dict of the model's fields, with any sensitive fields redacted.""" - return self.model_dump(exclude=self.get_sensitive_fields() | self.get_extra_fields_to_exclude_from_log()) + if extra_exclude is None: + extra_exclude = set() + + return self.model_dump( + exclude=self.get_sensitive_fields() | self.get_extra_fields_to_exclude_from_log() | extra_exclude, + ) class HordeResponse(HordeAPIMessage): From e98c30ed6571c7e329c3db0232983c9e814c563c Mon Sep 17 00:00:00 2001 From: tazlin Date: Mon, 24 Jun 2024 10:17:06 -0400 Subject: [PATCH 10/11] tests: log `skipped` separate w/ fake worker --- tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index 8394161..f43de89 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -68,7 +68,7 @@ async def fake_worker_checkin( ) assert isinstance(job_pop_response, ImageGenerateJobPopResponse) - logger.info(f"{job_pop_response.log_safe_model_dump()}") + logger.info(f"{job_pop_response.log_safe_model_dump({'skipped'})}") assert not job_pop_response.ids_present assert job_pop_response.skipped is not None @@ -103,7 +103,8 @@ async def fake_worker( ) assert isinstance(job_pop_response, ImageGenerateJobPopResponse) - logger.info(f"{job_pop_response.log_safe_model_dump()}") + logger.info(f"{job_pop_response.log_safe_model_dump({'skipped'})}") + logger.info(f"Checked in as fake worker ({effective_resolution}): {job_pop_response.skipped}") if not job_pop_response.ids_present: if current_try >= max_tries: From 68454196f1e3a84e81a938653885804f3c14a5a2 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sat, 20 Jul 2024 19:03:10 -0400 Subject: [PATCH 11/11] fix: use max compat pydanic; style fix --- .pre-commit-config.yaml | 5 +++-- requirements.txt | 4 ++-- .../test_ai_worker_roundtrip_api_calls.py | 22 +------------------ tests/conftest.py | 2 +- 4 files changed, 7 insertions(+), 26 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 463a2a7..51ad9fc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,12 +21,13 @@ repos: pass_filenames: false additional_dependencies: [ pytest, - pydantic, + pydantic==2.7.4, types-Pillow, types-requests, types-pytz, types-setuptools, types-urllib3, types-aiofiles, - StrEnum + StrEnum, + horde_model_reference==0.8.1, ] diff --git a/requirements.txt b/requirements.txt index d6d29b6..1e9eff4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ -horde_model_reference~=0.7.0 +horde_model_reference~=0.8.1 -pydantic +pydantic==2.7.4 requests StrEnum loguru diff --git a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py index f43de89..2a6fb8d 100644 --- a/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py +++ b/tests/ai_horde_api/test_ai_worker_roundtrip_api_calls.py @@ -1,47 +1,27 @@ import asyncio -import base64 -from typing import Any import aiohttp import PIL.Image import pytest -from loguru import logger import yarl +from loguru import logger from horde_sdk.ai_horde_api.ai_horde_clients import ( AIHordeAPIAsyncClientSession, AIHordeAPIAsyncSimpleClient, - AIHordeAPISimpleClient, ) from horde_sdk.ai_horde_api.apimodels import ( - KNOWN_ALCHEMY_TYPES, - AlchemyAsyncRequest, - AlchemyAsyncRequestFormItem, - AlchemyStatusResponse, ImageGenerateAsyncRequest, - ImageGenerateAsyncResponse, - ImageGenerateCheckResponse, - ImageGenerateJobPopPayload, ImageGenerateJobPopRequest, ImageGenerateJobPopResponse, - ImageGenerateJobPopSkippedStatus, ImageGenerateStatusResponse, - ImageGenerationInputPayload, ImageGenerationJobSubmitRequest, JobSubmitResponse, - LorasPayloadEntry, ) -from horde_sdk.ai_horde_api.apimodels.base import ExtraSourceImageEntry from horde_sdk.ai_horde_api.consts import ( GENERATION_STATE, - KNOWN_FACEFIXERS, - KNOWN_MISC_POST_PROCESSORS, - KNOWN_SOURCE_PROCESSING, - KNOWN_UPSCALERS, - POST_PROCESSOR_ORDER_TYPE, ) from horde_sdk.ai_horde_api.fields import JobID -from horde_sdk.generic_api.apimodels import RequestErrorResponse class TestImageWorkerRoundtrip: diff --git a/tests/conftest.py b/tests/conftest.py index dda7e2e..bdec2f1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,8 +2,8 @@ import os import pathlib -from loguru import logger import pytest +from loguru import logger os.environ["TESTS_ONGOING"] = "1"