feat: Support for extra_source_images

Closes #165
Haidra-Org · Mar 21, 2024 · ea2075e · ea2075e
1 parent 702f0eb
commit ea2075e
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 1 deletion.
diff --git a/horde_sdk/ai_horde_api/apimodels/__init__.py b/horde_sdk/ai_horde_api/apimodels/__init__.py
@@ -36,7 +36,12 @@
     AlchemyStatusResponse,
     AlchemyUpscaleResult,
 )
-from horde_sdk.ai_horde_api.apimodels.base import GenMetadataEntry, LorasPayloadEntry, TIPayloadEntry
+from horde_sdk.ai_horde_api.apimodels.base import (
+    ExtraSourceImageEntry,
+    GenMetadataEntry,
+    LorasPayloadEntry,
+    TIPayloadEntry,
+)
 from horde_sdk.ai_horde_api.apimodels.generate._async import (
     ImageGenerateAsyncDryRunResponse,
     ImageGenerateAsyncRequest,
@@ -112,6 +117,7 @@
     "StatsModelsResponse",
     "StatsModelsTimeframe",
     "TIPayloadEntry",
+    "ExtraSourceImageEntry",
     "GenMetadataEntry",
     "UsageDetails",
     "UserAmountRecords",

diff --git a/horde_sdk/ai_horde_api/apimodels/base.py b/horde_sdk/ai_horde_api/apimodels/base.py
@@ -125,6 +125,18 @@ def strength_only_if_inject_ti(self) -> TIPayloadEntry:
         return self
 
 
+class ExtraSourceImageEntry(HordeAPIDataObject):
+    """Represents a single extra source image.
+
+    v2 API Model: `ExtraSourceImage`
+    """
+
+    image: str = Field(min_length=1)
+    """The URL of the image to download."""
+    strength: float = Field(default=1, ge=-5, le=5)
+    """The strength to apply to this image on various operations."""
+
+
 class ImageGenerateParamMixin(HordeAPIDataObject):
     """Mix-in class of some of the data included in a request to the `/v2/generate/async` endpoint.
 

diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_async.py b/horde_sdk/ai_horde_api/apimodels/generate/_async.py
@@ -4,6 +4,7 @@
 
 from horde_sdk.ai_horde_api.apimodels.base import (
     BaseAIHordeRequest,
+    ExtraSourceImageEntry,
     ImageGenerateParamMixin,
     JobResponseMixin,
 )
@@ -125,6 +126,8 @@ class ImageGenerateAsyncRequest(
     source_image: str | None = None
     source_processing: KNOWN_SOURCE_PROCESSING = KNOWN_SOURCE_PROCESSING.txt2img
     source_mask: str | None = None
+    extra_source_images: list[ExtraSourceImageEntry] | None = None
+    """Additional uploaded images which can be used for further operations."""
 
     @model_validator(mode="before")
     def validate_censor_nsfw(cls, values: dict) -> dict:

diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_pop.py b/horde_sdk/ai_horde_api/apimodels/generate/_pop.py
@@ -8,6 +8,7 @@
 
 from horde_sdk.ai_horde_api.apimodels.base import (
     BaseAIHordeRequest,
+    ExtraSourceImageEntry,
     ImageGenerateParamMixin,
 )
 from horde_sdk.ai_horde_api.apimodels.generate._submit import ImageGenerationJobSubmitRequest
@@ -109,6 +110,8 @@ class ImageGenerateJobPopResponse(HordeResponseBaseModel, ResponseRequiringFollo
     """If img_processing is set to 'inpainting' or 'outpainting', this parameter can be optionally provided as the
     mask of the areas to inpaint. If this arg is not passed, the inpainting/outpainting mask has to be embedded as
     alpha channel."""
+    extra_source_images: list[ExtraSourceImageEntry] | None = None
+    """Additional uploaded images which can be used for further operations."""
     r2_upload: str | None = None
     """(Obsolete) The r2 upload link to use to upload this image."""
     r2_uploads: list[str] | None = None