Merge pull request #72 from Haidra-Org/main

feat: Add support for gen_metadata key (#69)
Haidra-Org · Nov 1, 2023 · f99fb1a · f99fb1a
2 parents e0a16b0 + 9758cbe
commit f99fb1a
Show file tree

Hide file tree

Showing 14 changed files with 112 additions and 10 deletions.
diff --git a/.gitignore b/.gitignore
@@ -140,3 +140,4 @@ out.json
 .vscode/settings.json
 
 examples/requested_images/*.*
+_version.py
diff --git a/docs/request_field_names_and_descriptions.json b/docs/request_field_names_and_descriptions.json
@@ -345,6 +345,10 @@
         [
             "censored",
             null
+        ],
+        [
+            "gen_metadata",
+            null
         ]
     ],
     "StatsImageModelsRequest": [

diff --git a/horde_sdk/ai_horde_api/apimodels/base.py b/horde_sdk/ai_horde_api/apimodels/base.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import override
 
-from horde_sdk.ai_horde_api.consts import KNOWN_SAMPLERS, POST_PROCESSOR_ORDER_TYPE
+from horde_sdk.ai_horde_api.consts import KNOWN_SAMPLERS, METADATA_TYPE, METADATA_VALUE, POST_PROCESSOR_ORDER_TYPE
 from horde_sdk.ai_horde_api.endpoints import AI_HORDE_BASE_URL
 from horde_sdk.ai_horde_api.fields import JobID, WorkerID
 from horde_sdk.generic_api.apimodels import HordeRequest, HordeResponseBaseModel
@@ -198,3 +198,17 @@ class JobSubmitResponse(HordeResponseBaseModel):
     @classmethod
     def get_api_model_name(cls) -> str | None:
         return "GenerationSubmitted"
+
+
+class GenMetadataEntry(BaseModel):
+    """Represents a single generation metadata entry.
+
+    v2 API Model: `GenerationMetadataStable`
+    """
+
+    type_: METADATA_TYPE = Field(alias="type")
+    """The relevance of the metadata field."""
+    value: METADATA_VALUE = Field()
+    """The value of the metadata field."""
+    ref: str = Field(max_length=255)
+    """Optionally a reference for the metadata (e.g. a lora ID)"""
diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_status.py b/horde_sdk/ai_horde_api/apimodels/generate/_status.py
@@ -4,7 +4,7 @@
 from pydantic import BaseModel, Field, field_validator
 from typing_extensions import override
 
-from horde_sdk.ai_horde_api.apimodels.base import BaseAIHordeRequest, JobRequestMixin
+from horde_sdk.ai_horde_api.apimodels.base import BaseAIHordeRequest, GenMetadataEntry, JobRequestMixin
 from horde_sdk.ai_horde_api.apimodels.generate._progress import ResponseGenerationProgressMixin
 from horde_sdk.ai_horde_api.consts import GENERATION_STATE
 from horde_sdk.ai_horde_api.endpoints import AI_HORDE_API_ENDPOINT_SUBPATH
@@ -36,6 +36,8 @@ class ImageGeneration(BaseModel):
     """The seed which generated this image."""
     censored: bool
     """When true this image has been censored by the worker's safety filter."""
+    gen_metadata: list[GenMetadataEntry] | None = None
+    """Extra metadata about faulted or defaulted components of the generation"""
 
     @field_validator("id_", mode="before")
     def validate_id(cls, v: str | JobID) -> JobID | str:

diff --git a/horde_sdk/ai_horde_api/apimodels/generate/_submit.py b/horde_sdk/ai_horde_api/apimodels/generate/_submit.py
@@ -4,7 +4,12 @@
 from pydantic import model_validator
 from typing_extensions import override
 
-from horde_sdk.ai_horde_api.apimodels.base import BaseAIHordeRequest, JobRequestMixin, JobSubmitResponse
+from horde_sdk.ai_horde_api.apimodels.base import (
+    BaseAIHordeRequest,
+    GenMetadataEntry,
+    JobRequestMixin,
+    JobSubmitResponse,
+)
 from horde_sdk.ai_horde_api.consts import GENERATION_STATE
 from horde_sdk.ai_horde_api.endpoints import AI_HORDE_API_ENDPOINT_SUBPATH
 from horde_sdk.consts import HTTPMethod
@@ -25,6 +30,8 @@ class ImageGenerationJobSubmitRequest(BaseAIHordeRequest, JobRequestMixin, APIKe
     """The seed for this generation."""
     censored: bool = False
     """If True, this resulting image has been censored."""
+    gen_metadata: list[GenMetadataEntry] | None = None
+    """Extra metadata about faulted or defaulted components of the generation"""
 
     @model_validator(mode="after")
     def validate_generation(self) -> ImageGenerationJobSubmitRequest:

diff --git a/horde_sdk/ai_horde_api/consts.py b/horde_sdk/ai_horde_api/consts.py
@@ -180,3 +180,29 @@ class KNOWN_ALCHEMY_TYPES(StrEnum):
     CodeFormers = KNOWN_FACEFIXERS.GFPGAN
 
     strip_background = KNOWN_MISC_POST_PROCESSORS.strip_background
+
+
+class METADATA_TYPE(StrEnum):
+    """The generation metadata types that are known to the API.
+
+    (lora, ti, censorship, etc)
+    """
+
+    lora = auto()
+    ti = auto()
+    censorship = auto()
+    source_image = auto()
+    source_mask = auto()
+
+
+class METADATA_VALUE(StrEnum):
+    """The generation metadata values that are known to the API.
+
+    (download_failed, baseline_mismatch, etc)
+    """
+
+    download_failed = auto()
+    parse_failed = auto()
+    baseline_mismatch = auto()
+    csam = auto()
+    nsfw = auto()
diff --git a/tests/test_data/ai_horde_api/example_payloads/_v2_generate_submit_post.json b/tests/test_data/ai_horde_api/example_payloads/_v2_generate_submit_post.json
@@ -3,5 +3,12 @@
     "generation": "R2",
     "state": "ok",
     "seed": 0,
-    "censored": false
+    "censored": false,
+    "gen_metadata": [
+        {
+            "type": "lora",
+            "value": "download_failed",
+            "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+        }
+    ]
 }
diff --git a/tests/test_data/ai_horde_api/example_payloads/_v2_generate_text_async_post.json b/tests/test_data/ai_horde_api/example_payloads/_v2_generate_text_async_post.json
@@ -21,7 +21,10 @@
         "sampler_order": [
             0
         ],
-        "use_default_badwordsids": true
+        "use_default_badwordsids": true,
+        "stop_sequence": [
+            ""
+        ]
     },
     "softprompt": "a",
     "trusted_workers": false,

diff --git a/tests/test_data/ai_horde_api/example_payloads/_v2_generate_text_submit_post.json b/tests/test_data/ai_horde_api/example_payloads/_v2_generate_text_submit_post.json
@@ -1,5 +1,12 @@
 {
     "id": "00000000-0000-0000-0000-000000000000",
     "generation": "R2",
-    "state": "ok"
+    "state": "ok",
+    "gen_metadata": [
+        {
+            "type": "censorship",
+            "value": "csam",
+            "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+        }
+    ]
 }
diff --git a/tests/test_data/ai_horde_api/example_responses/_v2_generate_status_id_delete_200.json b/tests/test_data/ai_horde_api/example_responses/_v2_generate_status_id_delete_200.json
@@ -18,7 +18,14 @@
             "img": "",
             "seed": "",
             "id": "",
-            "censored": false
+            "censored": false,
+            "gen_metadata": [
+                {
+                    "type": "lora",
+                    "value": "download_failed",
+                    "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+                }
+            ]
         }
     ],
     "shared": false

diff --git a/tests/test_data/ai_horde_api/example_responses/_v2_generate_status_id_get_200.json b/tests/test_data/ai_horde_api/example_responses/_v2_generate_status_id_get_200.json
@@ -18,7 +18,14 @@
             "img": "",
             "seed": "",
             "id": "",
-            "censored": false
+            "censored": false,
+            "gen_metadata": [
+                {
+                    "type": "lora",
+                    "value": "download_failed",
+                    "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+                }
+            ]
         }
     ],
     "shared": false

diff --git a/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_pop_post_200.json b/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_pop_post_200.json
@@ -21,6 +21,9 @@
             0
         ],
         "use_default_badwordsids": true,
+        "stop_sequence": [
+            ""
+        ],
         "prompt": ""
     },
     "id": "",

diff --git a/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_status_id_delete_200.json b/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_status_id_delete_200.json
@@ -16,7 +16,14 @@
             "model": "",
             "state": "ok",
             "text": "",
-            "seed": 0
+            "seed": 0,
+            "gen_metadata": [
+                {
+                    "type": "censorship",
+                    "value": "csam",
+                    "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+                }
+            ]
         }
     ]
 }
diff --git a/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_status_id_get_200.json b/tests/test_data/ai_horde_api/example_responses/_v2_generate_text_status_id_get_200.json
@@ -16,7 +16,14 @@
             "model": "",
             "state": "ok",
             "text": "",
-            "seed": 0
+            "seed": 0,
+            "gen_metadata": [
+                {
+                    "type": "censorship",
+                    "value": "csam",
+                    "ref": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+                }
+            ]
         }
     ]
 }