Merge pull request #6 from jakechai/main

fixed missing parameters and naming and added missing APIs
This commit is contained in:
Mohamad-Stability
2025-02-26 22:39:10 +02:00
committed by GitHub
2 changed files with 127 additions and 45 deletions
+3 -3
View File
@@ -1,4 +1,4 @@
from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure, StabilityFastUpscale, StabilityStyle, StabilitySearchAndRecolor
from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilityImageUltra, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure, StabilityFastUpscale, StabilityControlStyle, StabilitySearchAndRecolor
NODE_CLASS_MAPPINGS = {
"Stability Conservative Upscale": StabilityConservativeUpscale,
@@ -14,6 +14,6 @@ NODE_CLASS_MAPPINGS = {
"Stability Control Sketch": StabilityControlSketch,
"Stability Control Structure": StabilityControlStructure,
"Stability Fast Upscale": StabilityFastUpscale,
"Stability Style": StabilityStyle,
"Stability Style": StabilityControlStyle,
"Stability Search and Recolor": StabilitySearchAndRecolor,
}
}
+124 -42
View File
@@ -60,7 +60,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["subject_image"] = buffered_sub.getvalue()
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)
@@ -77,14 +98,35 @@ class StabilityBase:
headers["Accept"] = self.ACCEPT
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']
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)
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")
@@ -93,7 +135,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)
@@ -138,9 +180,9 @@ class StabilityBase:
return {
"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):
@@ -154,7 +196,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}),
@@ -170,9 +212,11 @@ class StabilityImageUltra(StabilityBase):
"prompt": ("STRING", {"multiline": True}),
},
"optional": {
"image": ("IMAGE",),
"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"],),
"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}),
},
@@ -199,8 +243,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",),
@@ -209,7 +253,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}),
}
@@ -236,27 +280,32 @@ 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}),
"output_format": (["png", "webp", "jpeg"],),
"api_key_override": ("STRING", {"multiline": False}),
}
}
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):
@@ -265,25 +314,30 @@ 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"],),
"api_key_override": ("STRING", {"multiline": False}),
}
}
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):
@@ -292,11 +346,12 @@ 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": {
"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"],),
@@ -315,9 +370,10 @@ class StabilitySD3(StabilityBase):
"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"],),
"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}),
},
@@ -333,10 +389,11 @@ class StabilityOutpainting(StabilityBase):
"left": ("INT", {"default": 0, "min": 0, "max": 512}),
"right": ("INT", {"default": 0, "min": 0, "max": 512}),
"up": ("INT", {"default": 0, "min": 0, "max": 512}),
"down": ("INT", {"default": 0, "min": 0, "max": 512}),\
"down": ("INT", {"default": 0, "min": 0, "max": 512}),
},
"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}),
@@ -391,7 +448,7 @@ class StabilityFastUpscale(StabilityBase):
}
}
class StabilityStyle(StabilityBase):
class StabilityControlStyle(StabilityBase):
API_ENDPOINT = "stable-image/control/style"
ACCEPT = "image/*"
INPUT_SPEC = {
@@ -400,13 +457,13 @@ class StabilityStyle(StabilityBase):
"prompt": ("STRING", {"multiline": True}),
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True}),
"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}),
"output_format": (["png", "webp", "jpeg"],),
"api_key_override": ("STRING", {"multiline": False}),
}
},
}
class StabilitySearchAndRecolor(StabilityBase):
@@ -420,7 +477,32 @@ class StabilitySearchAndRecolor(StabilityBase):
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True}),
"grow_mask": ("INT", {"default": 3, "min": 0, "max": 20}),
"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 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}),