From 9f95b2c27f8294104973931a0f8c50c51346d71a Mon Sep 17 00:00:00 2001 From: jakechai <61836864+jakechai@users.noreply.github.com> Date: Fri, 4 Oct 2024 09:36:19 +0800 Subject: [PATCH 1/9] fix Search And Replace remove "Search Prompt" --- stability_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stability_api.py b/stability_api.py index 72f7498..a2785c3 100644 --- a/stability_api.py +++ b/stability_api.py @@ -287,7 +287,7 @@ class StabilitySearchAndReplace(StabilityBase): INPUT_SPEC = { "required": { "image": ("IMAGE",), - "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"), + "search_prompt": ("STRING", {"multiline": True}), "prompt": ("STRING", {"multiline": True}), }, "optional": { From 8141eeb3f7e466fcfcd411f019d7faf3fa8fd1bd Mon Sep 17 00:00:00 2001 From: jakechai <61836864+jakechai@users.noreply.github.com> Date: Fri, 4 Oct 2024 10:10:33 +0800 Subject: [PATCH 2/9] add SearchAndRecolor ControlStyle FastUpscale --- stability_api.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/stability_api.py b/stability_api.py index a2785c3..ce727d1 100644 --- a/stability_api.py +++ b/stability_api.py @@ -372,3 +372,55 @@ class StabilityControlStructure(StabilityBase): "api_key_override": ("STRING", {"multiline": False}), }, } + + +class StabilitySearchAndRecolor(StabilityBase): + API_ENDPOINT = "stable-image/edit/search-and-recolor" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "select_prompt": ("STRING", {"multiline": True}), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True}), + "grow_mask": ("INT", {"default": 3, "min": 0, "max": 20, "step": 1}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "api_key_override": ("STRING", {"multiline": False}), + "output_format": (["png", "webp", "jpeg"],), + }, + } + + +class StabilityControlStyle(StabilityBase): + API_ENDPOINT = "stable-image/control/structure" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + }, + } + + +class StabilityFastUpscale(StabilityBase): + API_ENDPOINT = "stable-image/upscale/fast" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + }, + "optional": { + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + } + } From aa4b2005641714a9b0bb764f7aa34d55e65e2387 Mon Sep 17 00:00:00 2001 From: jakechai <61836864+jakechai@users.noreply.github.com> Date: Fri, 4 Oct 2024 10:11:50 +0800 Subject: [PATCH 3/9] add SearchAndRecolor ControlStyle FastUpscale --- __init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/__init__.py b/__init__.py index 525cb18..b7b07e0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure +from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure, StabilitySearchAndRecolor, StabilityControlStyle, StabilityFastUpscale NODE_CLASS_MAPPINGS = { "Stability Conservative Upscale": StabilityConservativeUpscale, @@ -13,4 +13,7 @@ NODE_CLASS_MAPPINGS = { "Stability SD3": StabilitySD3, "Stability Control Skech": StabilityControlSketch, "Stability Control Structure": StabilityControlStructure, + "Stability Search And Recolor": StabilitySearchAndRecolor, + "Stability Control Style": StabilityControlStyle, + "Stability Fast Upscale": StabilityFastUpscale, } From d3c67acbef906ccbd119633d2f8ebf0cfd3bf645 Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Wed, 23 Oct 2024 09:29:58 +0800 Subject: [PATCH 4/9] Add SD3.5. --- stability_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stability_api.py b/stability_api.py index ce727d1..a9a4c04 100644 --- a/stability_api.py +++ b/stability_api.py @@ -304,7 +304,7 @@ class StabilitySD3(StabilityBase): ACCEPT = "image/*" INPUT_SPEC = { "required": { - "model": (["sd3", "sd3-turbo"],), + "model": (["sd3.5-large", "sd3.5-large-turbo", "sd3-large", "sd3-large-turbo", "sd3-medium"],), "prompt": ("STRING", {"multiline": True}), }, "optional": { From d41163b6aa30716f080e93417b408befeea3667e Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Tue, 26 Nov 2024 14:15:46 +0800 Subject: [PATCH 5/9] Update __init__.py --- __init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index b7b07e0..5a3921a 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure, StabilitySearchAndRecolor, StabilityControlStyle, StabilityFastUpscale +from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure, StabilitySearchAndRecolor, StabilityControlStyle, StabilityFastUpscale, StabilityReplaceRelight NODE_CLASS_MAPPINGS = { "Stability Conservative Upscale": StabilityConservativeUpscale, @@ -11,9 +11,10 @@ NODE_CLASS_MAPPINGS = { "Stability Search and Replace": StabilitySearchAndReplace, "Stability Outpainting": StabilityOutpainting, "Stability SD3": StabilitySD3, - "Stability Control Skech": StabilityControlSketch, + "Stability Control Sketch": StabilityControlSketch, "Stability Control Structure": StabilityControlStructure, "Stability Search And Recolor": StabilitySearchAndRecolor, "Stability Control Style": StabilityControlStyle, "Stability Fast Upscale": StabilityFastUpscale, + "Stability Replace Background and Relight": StabilityReplaceRelight, } From d569e1119f8e77ccde989d0965eb1b17d56fc95d Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Tue, 26 Nov 2024 14:16:29 +0800 Subject: [PATCH 6/9] Update stability_api.py --- stability_api.py | 81 ++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 75 insertions(+), 6 deletions(-) diff --git a/stability_api.py b/stability_api.py index a9a4c04..917533b 100644 --- a/stability_api.py +++ b/stability_api.py @@ -59,7 +59,28 @@ class StabilityBase: files = self._get_files(buffered, **kwargs) else: kwargs.pop("strength", None) - + + image_subject = kwargs.get('subject_image', None) + if image_subject is not None: + buffered_sub = BytesIO() + image_subject = ToPILImage()(image_subject.squeeze(0).permute(2,0,1)) + image_subject.save(buffered_sub, format="PNG") + files = self._get_files_sub(buffered_sub, **kwargs) + + image_bg_ref = kwargs.get('background_reference', None) + if image_bg_ref is not None: + buffered_bg_ref = BytesIO() + image_bg_ref = ToPILImage()(image_bg_ref.squeeze(0).permute(2,0,1)) + image_bg_ref.save(buffered_bg_ref, format="PNG") + files["background_reference"] = buffered_bg_ref.getvalue() + + image_lt_ref = kwargs.get('light_reference', None) + if image_lt_ref is not None: + buffered_lt_ref = BytesIO() + image_lt_ref = ToPILImage()(image_lt_ref.squeeze(0).permute(2,0,1)) + image_lt_ref.save(buffered_lt_ref, format="PNG") + files["light_reference"] = buffered_lt_ref.getvalue() + style = kwargs.get('style', False) if style is False: kwargs.pop('style_preset', None) @@ -76,7 +97,25 @@ class StabilityBase: headers["Accept"] = self.ACCEPT data = self._get_data(**kwargs) - + + if kwargs.get('subject_image', None) is not None: + if 'subject_image' in data: + del data['subject_image'] + + if 'background_reference' in data: + del data['background_reference'] + + if 'light_reference' in data: + del data['light_reference'] + + if kwargs.get('light_reference', None) is not None: + if 'light_source_direction' in data: + del data['light_source_direction'] + else: + if 'light_source_direction' in data and data['light_source_direction'] == 'none': + if 'light_source_strength' in data: + del data['light_source_strength'] + req = PreparedRequest() req.prepare_method('POST') req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None) @@ -92,7 +131,7 @@ class StabilityBase: while True: response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers) if response.status_code == 200: - if self.ACCEPT == "image/*": + if self.ACCEPT == "image/*" or self.ACCEPT == "*/*": return self._return_image(response) if self.ACCEPT == "video/*": return self._return_video(response) @@ -137,7 +176,12 @@ class StabilityBase: return { "image": buffered.getvalue() } - + + def _get_files_sub(self, buffered, **kwargs): + return { + "subject_image": buffered.getvalue() + } + def _get_data(self, **kwargs): return {k: v for k, v in kwargs.items() if k != "image"} @@ -198,8 +242,8 @@ class StabilityConservativeUpscale(StabilityBase): class StabilityCreativeUpscale(StabilityBase): API_ENDPOINT = "stable-image/upscale/creative" - POLL_ENDPOINT = "stable-image/upscale/creative/result/" - ACCEPT = "image/*" + POLL_ENDPOINT = "results/" + ACCEPT = "*/*" INPUT_SPEC = { "required": { "image": ("IMAGE",), @@ -424,3 +468,28 @@ class StabilityFastUpscale(StabilityBase): "api_key_override": ("STRING", {"multiline": False}), } } + +class StabilityReplaceRelight(StabilityBase): + API_ENDPOINT = "stable-image/edit/replace-background-and-relight" + POLL_ENDPOINT = "results/" + ACCEPT = "*/*" + INPUT_SPEC = { + "required": { + "subject_image": ("IMAGE",), + }, + "optional": { + "background_reference": ("IMAGE",), + "background_prompt": ("STRING", {"multiline": True}), + "foreground_prompt": ("STRING", {"multiline": True}), + "negative_prompt": ("STRING", {"multiline": True}), + "preserve_original_subject": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}), + "original_background_depth": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "keep_original_background": ("BOOLEAN", {"default": False}), + "light_source_direction": (["none", "above", "below", "left", "right"],), + "light_reference": ("IMAGE",), + "light_source_strength": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + } + } From 0230c644b521ec99b126e68e14da36fccda4d62d Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Tue, 26 Nov 2024 18:16:10 +0800 Subject: [PATCH 7/9] Add missing parameters. --- stability_api.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/stability_api.py b/stability_api.py index 917533b..39a9037 100644 --- a/stability_api.py +++ b/stability_api.py @@ -213,8 +213,10 @@ class StabilityImageUltra(StabilityBase): "prompt": ("STRING", {"multiline": True}), }, "optional": { + "image": ("IMAGE",), "negative_prompt": ("STRING", {"multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), "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}), @@ -252,7 +254,7 @@ class StabilityCreativeUpscale(StabilityBase): "optional": { "negative_prompt": ("STRING", {"multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}), + "creativity": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 0.35, "step": 0.01}), "output_format": (["png", "webp", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), } @@ -280,6 +282,7 @@ class StabilityInpainting(StabilityBase): }, "optional": { "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "grow_mask": ("INT", {"default": 5, "min": 0, "max": 20}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "output_format": (["png", "webp", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), @@ -307,6 +310,7 @@ class StabilityErase(StabilityBase): "mask": ("MASK",) }, "optional": { + "grow_mask": ("INT", {"default": 5, "min": 0, "max": 20}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "output_format": (["png", "webp", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), @@ -336,6 +340,7 @@ class StabilitySearchAndReplace(StabilityBase): }, "optional": { "negative_prompt": ("STRING", {"multiline": True}), + "grow_mask": ("INT", {"default": 3, "min": 0, "max": 20}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "api_key_override": ("STRING", {"multiline": False}), "output_format": (["png", "webp", "jpeg"],), @@ -348,12 +353,13 @@ class StabilitySD3(StabilityBase): ACCEPT = "image/*" INPUT_SPEC = { "required": { - "model": (["sd3.5-large", "sd3.5-large-turbo", "sd3-large", "sd3-large-turbo", "sd3-medium"],), + "model": (["sd3.5-large", "sd3.5-large-turbo", "sd3.5-medium", "sd3-large", "sd3-large-turbo", "sd3-medium"],), "prompt": ("STRING", {"multiline": True}), }, "optional": { "image": ("IMAGE",), "negative_prompt": ("STRING", {"multiline": True}), + "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 10.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), @@ -376,6 +382,7 @@ class StabilityOutpainting(StabilityBase): }, "optional": { "prompt": ("STRING", {"multiline": True}), + "creativity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "output_format": (["png", "webp", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), From 2988e7a8a6701d92c4232af6ba6057ea268f50f2 Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Sun, 22 Dec 2024 13:33:09 +0800 Subject: [PATCH 8/9] Add Replace Background and Relight. Add aspect_ratio resolution. --- stability_api.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/stability_api.py b/stability_api.py index 39a9037..a87a5d8 100644 --- a/stability_api.py +++ b/stability_api.py @@ -52,8 +52,9 @@ class StabilityBase: image = kwargs.get('image', None) if image is not None: - kwargs["mode"] = "image-to-image" - kwargs.pop("aspect_ratio", None) + if self.API_ENDPOINT != "stable-image/control/style": + kwargs["mode"] = "image-to-image" + kwargs.pop("aspect_ratio", None) image = ToPILImage()(image.squeeze(0).permute(2,0,1)) image.save(buffered, format="PNG") files = self._get_files(buffered, **kwargs) @@ -98,6 +99,9 @@ class StabilityBase: data = self._get_data(**kwargs) + if kwargs.get('aspect_ratio', None) is not None: + data['aspect_ratio'] = data['aspect_ratio'].split("(", 1)[0] + if kwargs.get('subject_image', None) is not None: if 'subject_image' in data: del data['subject_image'] @@ -197,7 +201,7 @@ class StabilityCore(StabilityBase): "negative_prompt": ("STRING", {"multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "output_format": (["png", "webp", "jpeg"],), - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "aspect_ratio": (["1:1(1536, 1536)", "5:4(1632, 1344)", "3:2(1824, 1248)", "16:9(2016, 1152)", "21:9(2304, 960)", "4:5(1344, 1632)", "2:3(1248, 1824)", "9:16(1152, 2016)", "9:21(960, 2304)"],), "style": ("BOOLEAN", {"default": False}), "style_preset": (["3d-model", "analog-film", "anime", "cinematic", "comic-book", "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly", "modeling-compound", "neon-punk", "origami", "photographic", "pixel-art", "tile-texture"],), "api_key_override": ("STRING", {"multiline": False}), @@ -216,8 +220,8 @@ class StabilityImageUltra(StabilityBase): "image": ("IMAGE",), "negative_prompt": ("STRING", {"multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), + "aspect_ratio": (["1:1(1024, 1024)", "5:4(1088, 896)", "3:2(1216, 832)", "16:9(1344, 768)", "21:9(1536, 640)", "4:5(896, 1088)", "2:3(832, 1216)", "9:16(768, 1344)", "9:21(640, 1536)"],), "output_format": (["png", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), }, @@ -361,8 +365,8 @@ class StabilitySD3(StabilityBase): "negative_prompt": ("STRING", {"multiline": True}), "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 10.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), + "aspect_ratio": (["1:1(1024, 1024)", "5:4(1088, 896)", "3:2(1216, 832)", "16:9(1344, 768)", "21:9(1536, 640)", "4:5(896, 1088)", "2:3(832, 1216)", "9:16(768, 1344)", "9:21(640, 1536)"],), "output_format": (["png", "jpeg"],), "api_key_override": ("STRING", {"multiline": False}), }, @@ -445,7 +449,7 @@ class StabilitySearchAndRecolor(StabilityBase): class StabilityControlStyle(StabilityBase): - API_ENDPOINT = "stable-image/control/structure" + API_ENDPOINT = "stable-image/control/style" ACCEPT = "image/*" INPUT_SPEC = { "required": { @@ -453,7 +457,7 @@ class StabilityControlStyle(StabilityBase): "prompt": ("STRING", {"multiline": True}), }, "optional": { - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "aspect_ratio": (["1:1(1024, 1024)", "5:4(1144, 915)", "3:2(1254, 836)", "16:9(1365, 768)", "21:9(1564, 670)", "4:5(915, 1144)", "2:3(836, 1254)", "9:16(768, 1365)", "9:21(670, 1564)"],), "fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "negative_prompt": ("STRING", {"multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), From 67db771e0e9d1c99a65c9c9872d4dad29e72bef3 Mon Sep 17 00:00:00 2001 From: Jake Chai <61836864+jakechai@users.noreply.github.com> Date: Sun, 22 Dec 2024 17:56:21 +0800 Subject: [PATCH 9/9] Fix mask issue for Erase and Inpaint Fix the mask missing issue for Erase and Inpaint. The mask is changed to optional. --- stability_api.py | 63 +++++++++++++++++++++++++----------------------- 1 file changed, 33 insertions(+), 30 deletions(-) diff --git a/stability_api.py b/stability_api.py index a87a5d8..9eb4760 100644 --- a/stability_api.py +++ b/stability_api.py @@ -66,7 +66,7 @@ class StabilityBase: buffered_sub = BytesIO() image_subject = ToPILImage()(image_subject.squeeze(0).permute(2,0,1)) image_subject.save(buffered_sub, format="PNG") - files = self._get_files_sub(buffered_sub, **kwargs) + files["subject_image"] = buffered_sub.getvalue() image_bg_ref = kwargs.get('background_reference', None) if image_bg_ref is not None: @@ -126,7 +126,7 @@ class StabilityBase: req.prepare_headers(headers) req.prepare_body(data=data, files=files) response = requests.Session().send(req) - + if response.status_code == 200: if self.POLL_ENDPOINT != "": id = response.json().get("id") @@ -181,13 +181,8 @@ class StabilityBase: "image": buffered.getvalue() } - def _get_files_sub(self, buffered, **kwargs): - return { - "subject_image": buffered.getvalue() - } - def _get_data(self, **kwargs): - return {k: v for k, v in kwargs.items() if k != "image"} + return {k: v for k, v in kwargs.items() if k != "image" and k != "mask"} class StabilityCore(StabilityBase): @@ -281,10 +276,10 @@ class StabilityInpainting(StabilityBase): INPUT_SPEC = { "required": { "image": ("IMAGE",), - "mask": ("MASK",), "prompt": ("STRING", {"multiline": True, "default": ""}), }, "optional": { + "mask": ("MASK",), "negative_prompt": ("STRING", {"multiline": True, "default": ""}), "grow_mask": ("INT", {"default": 5, "min": 0, "max": 20}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), @@ -293,16 +288,20 @@ class StabilityInpainting(StabilityBase): } } def _get_files(self, buffered, **kwargs): - mask = kwargs.get("mask") - to_pil = ToPILImage() - mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) - mask = to_pil(mask.squeeze(0).permute(2,0,1)) - buffered_mask = BytesIO() - mask.save(buffered_mask, format="PNG") - return { - "image": buffered.getvalue(), - "mask": buffered_mask.getvalue(), - } + mask = kwargs.get("mask", None) + if mask != None: + mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + mask = ToPILImage()(mask.squeeze(0).permute(2,0,1)) + buffered_mask = BytesIO() + mask.save(buffered_mask, format="PNG") + return { + "image": buffered.getvalue(), + "mask": buffered_mask.getvalue(), + } + else: + return { + "image": buffered.getvalue(), + } class StabilityErase(StabilityBase): @@ -311,9 +310,9 @@ class StabilityErase(StabilityBase): INPUT_SPEC = { "required": { "image": ("IMAGE",), - "mask": ("MASK",) }, "optional": { + "mask": ("MASK",), "grow_mask": ("INT", {"default": 5, "min": 0, "max": 20}), "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), "output_format": (["png", "webp", "jpeg"],), @@ -321,16 +320,20 @@ class StabilityErase(StabilityBase): } } def _get_files(self, buffered, **kwargs): - mask = kwargs.get("mask") - to_pil = ToPILImage() - mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) - mask = to_pil(mask.squeeze(0).permute(2,0,1)) - buffered_mask = BytesIO() - mask.save(buffered_mask, format="PNG") - return { - "image": buffered.getvalue(), - "mask": buffered_mask.getvalue(), - } + mask = kwargs.get("mask", None) + if mask != None: + mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + mask = ToPILImage()(mask.squeeze(0).permute(2,0,1)) + buffered_mask = BytesIO() + mask.save(buffered_mask, format="PNG") + return { + "image": buffered.getvalue(), + "mask": buffered_mask.getvalue(), + } + else: + return { + "image": buffered.getvalue(), + } class StabilitySearchAndReplace(StabilityBase):