Merge pull request #112 from Haidra-Org/main

feat: Support n_iter batching (#109)

docs/request_field_names_and_descriptions.json

@@ -297,6 +297,10 @@
         [
             "allow_lora",
             null
+        ],
+        [
+            "amount",
+            null
         ]
     ],
     "ImageGenerateStatusRequest": [

docs/response_field_names_and_descriptions.json

@@ -302,6 +302,10 @@
             "id_",
             null
         ],
+        [
+            "ids",
+            null
+        ],
         [
             "payload",
             null
@@ -329,6 +333,10 @@
         [
             "r2_upload",
             null
+        ],
+        [
+            "r2_uploads",
+            null
         ]
     ],
     "JobSubmitResponse": [

horde_sdk/ai_horde_api/apimodels/generate/_pop.py

@@ -80,7 +80,9 @@ class ImageGenerateJobPopResponse(HordeResponseBaseModel, ResponseRequiringFollo
     """
 
     id_: JobID | None = Field(None, alias="id")
-    """The UUID for this image generation."""
+    """(Obsolete) The UUID for this image generation."""
+    ids: list[JobID]
+    """A list of UUIDs for image generation."""
 
     payload: ImageGenerateJobPopPayload
     """The parameters used to generate this image."""
@@ -97,7 +99,9 @@ class ImageGenerateJobPopResponse(HordeResponseBaseModel, ResponseRequiringFollo
     mask of the areas to inpaint. If this arg is not passed, the inpainting/outpainting mask has to be embedded as
     alpha channel."""
     r2_upload: str | None = None
-    """The r2 upload link to use to upload this image."""
+    """(Obsolete) The r2 upload link to use to upload this image."""
+    r2_uploads: list[str] | None = None
+    """The r2 upload links for each this image. Each index matches the ID in self.ids"""
 
     @field_validator("source_processing")
     def source_processing_must_be_known(cls, v: str | KNOWN_SOURCE_PROCESSING) -> str | KNOWN_SOURCE_PROCESSING:
@@ -177,6 +181,7 @@ class ImageGenerateJobPopRequest(BaseAIHordeRequest, APIKeyAllowedInRequestMixin
     allow_post_processing: bool = True
     allow_controlnet: bool = False
     allow_lora: bool = False
+    amount: int = 1
 
     @override
     @classmethod

horde_sdk/ai_horde_api/endpoints.py

@@ -32,6 +32,7 @@ class AI_HORDE_API_ENDPOINT_SUBPATH(GENERIC_API_ENDPOINT_SUBPATH):
     v2_generate_async = "/v2/generate/async"
 
     v2_generate_pop = "/v2/generate/pop"
+    v2_generate_pop_multi = "/v2/generate/pop_multi"
     v2_generate_submit = "/v2/generate/submit"
 
     # Note that `{id}`` (or any variable wrapped in curly braces) is dynamically replaced with the appropriate value

horde_sdk/ai_horde_worker/bridge_data.py

@@ -234,6 +234,13 @@ class ImageWorkerBridgeData(SharedHordeBridgeData):
     )
     """The number of threads to use for inference."""
 
+    max_batch: int = Field(
+        default=1,
+        ge=1,
+        le=20,  # 20 is the max per request on the horde
+    )
+    """The number of threads to use for inference."""
+
     require_upfront_kudos: bool = False
     """Whether to require upfront kudos for jobs. This effectively makes the worker reject jobs from
     the anonymous API key, or from keys with no kudos."""

tests/test_data/ai_horde_api/example_payloads/_v2_generate_async_post.json

@@ -1,7 +1,7 @@
 {
     "prompt": "a",
     "params": {
-        "sampler_name": "k_heun",
+        "sampler_name": "lcm",
         "cfg_scale": 7.5,
         "denoising_strength": 0.75,
         "seed": "The little seed that could",

tests/test_data/ai_horde_api/example_payloads/_v2_generate_pop_post.json

@@ -7,10 +7,10 @@
     "models": [
         "aaa"
     ],
-    "bridge_version": 1,
     "bridge_agent": "AI Horde Worker:24:https://github.com/db0/AI-Horde-Worker",
     "threads": 1,
     "require_upfront_kudos": false,
+    "amount": 1,
     "max_pixels": 262144,
     "blacklist": [
         ""

tests/test_data/ai_horde_api/example_payloads/_v2_generate_text_pop_post.json

@@ -7,10 +7,10 @@
     "models": [
         "aaa"
     ],
-    "bridge_version": 1,
     "bridge_agent": "AI Horde Worker:24:https://github.com/db0/AI-Horde-Worker",
     "threads": 1,
     "require_upfront_kudos": false,
+    "amount": 1,
     "max_length": 512,
     "max_context_length": 2048,
     "softprompts": [

tests/test_data/ai_horde_api/example_payloads/_v2_interrogate_pop_post.json

@@ -7,7 +7,6 @@
         "caption"
     ],
     "amount": 1,
-    "bridge_version": 1,
     "bridge_agent": "AI Horde Worker:24:https://github.com/db0/AI-Horde-Worker",
     "threads": 1,
     "max_tiles": 16

tests/test_data/ai_horde_api/example_responses/_v2_generate_pop_post_200.json

@@ -1,6 +1,6 @@
 {
     "payload": {
-        "sampler_name": "k_dpm_fast",
+        "sampler_name": "lcm",
         "cfg_scale": 7.5,
         "denoising_strength": 0.75,
         "seed": "The little seed that could",
@@ -45,6 +45,9 @@
         "use_nsfw_censor": false
     },
     "id": "",
+    "ids": [
+        "00000000-0000-0000-0000-000000000000"
+    ],
     "skipped": {
         "worker_id": 0.0,
         "performance": 0.0,
@@ -66,5 +69,8 @@
     "source_image": "",
     "source_processing": "img2img",
     "source_mask": "",
-    "r2_upload": ""
+    "r2_upload": "",
+    "r2_uploads": [
+        ""
+    ]
 }

tests/test_data/ai_horde_api/example_responses/_v2_generate_text_pop_post_200.json

@@ -28,6 +28,9 @@
         "prompt": ""
     },
     "id": "",
+    "ids": [
+        "00000000-0000-0000-0000-000000000000"
+    ],
     "skipped": {
         "worker_id": 0.0,
         "performance": 0.0,