From 1793086d6ea5bed29225b549c3b9620f35e8fd0a Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Thu, 6 Jun 2024 07:18:12 -0700
Subject: [PATCH] add Image Ultra API

---
 __init__.py      |  3 ++-
 stability_api.py | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index d8275c2..525cb18 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,4 +1,4 @@
-from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure
+from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure
 
 NODE_CLASS_MAPPINGS = {
     "Stability Conservative Upscale": StabilityConservativeUpscale,
@@ -7,6 +7,7 @@ NODE_CLASS_MAPPINGS = {
     "Stability Erase": StabilityErase,
     "Stability Inpainting": StabilityInpainting,
     "Stability Image Core": StabilityCore,
+    "Stability Image Ultra": StabilityImageUltra,
     "Stability Search and Replace": StabilitySearchAndReplace,
     "Stability Outpainting": StabilityOutpainting,
     "Stability SD3": StabilitySD3,
diff --git a/stability_api.py b/stability_api.py
index d3ae22c..72f7498 100644
--- a/stability_api.py
+++ b/stability_api.py
@@ -161,6 +161,23 @@ class StabilityCore(StabilityBase):
     }
 
 
+class StabilityImageUltra(StabilityBase):
+    API_ENDPOINT = "stable-image/generate/ultra"
+    ACCEPT = "image/*"
+    INPUT_SPEC = {
+        "required": {
+            "prompt": ("STRING", {"multiline": True}),
+        },
+        "optional": {
+            "negative_prompt": ("STRING", {"multiline": True}),
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
+            "output_format": (["png", "jpeg"],),
+            "api_key_override": ("STRING", {"multiline": False}),
+        },
+    }
+
+
 class StabilityConservativeUpscale(StabilityBase):
     API_ENDPOINT = "stable-image/upscale/conservative"
     ACCEPT = "image/*"