Conservative Upscale and Erase nodes

This commit is contained in:
Alex "mcmonkey" Goodwin
2024-05-22 22:09:31 -07:00
parent 30d125b22a
commit e6c35940f9
2 changed files with 51 additions and 3 deletions
+3 -1
View File
@@ -1,8 +1,10 @@
from .stability_api import StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure
from .stability_api import StabilityConservativeUpscale, StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityErase, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3, StabilityControlSketch, StabilityControlStructure
NODE_CLASS_MAPPINGS = {
"Stability Conservative Upscale": StabilityConservativeUpscale,
"Stability Creative Upscale": StabilityCreativeUpscale,
"Stability Remove Background": StabilityRemoveBackground,
"Stability Erase": StabilityErase,
"Stability Inpainting": StabilityInpainting,
"Stability Image Core": StabilityCore,
"Stability Search and Replace": StabilitySearchAndReplace,
+48 -2
View File
@@ -161,6 +161,24 @@ class StabilityCore(StabilityBase):
}
class StabilityConservativeUpscale(StabilityBase):
API_ENDPOINT = "stable-image/upscale/conservative"
ACCEPT = "image/*"
INPUT_SPEC = {
"required": {
"image": ("IMAGE",),
"prompt": ("STRING", {"multiline": True}),
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True}),
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
"creativity": ("FLOAT", {"default": 0.35, "min": 0.2, "max": 0.5, "step": 0.01}),
"output_format": (["png", "webp", "jpeg"],),
"api_key_override": ("STRING", {"multiline": False}),
}
}
class StabilityCreativeUpscale(StabilityBase):
API_ENDPOINT = "stable-image/upscale/creative"
POLL_ENDPOINT = "stable-image/upscale/creative/result/"
@@ -189,6 +207,7 @@ class StabilityRemoveBackground(StabilityBase):
},
}
class StabilityInpainting(StabilityBase):
API_ENDPOINT = "stable-image/edit/inpaint"
ACCEPT = "image/*"
@@ -196,10 +215,37 @@ class StabilityInpainting(StabilityBase):
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"prompt": ("STRING", {"multiline": True, "default": ""}),\
"prompt": ("STRING", {"multiline": True, "default": ""}),
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
"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(),
}
class StabilityErase(StabilityBase):
API_ENDPOINT = "stable-image/edit/erase"
ACCEPT = "image/*"
INPUT_SPEC = {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",)
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
"output_format": (["png", "webp", "jpeg"],),
"api_key_override": ("STRING", {"multiline": False}),