diff --git a/comfy_api_nodes/apis/stability_api.py b/comfy_api_nodes/apis/stability_api.py new file mode 100644 index 000000000..410884c6e --- /dev/null +++ b/comfy_api_nodes/apis/stability_api.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from pydantic import BaseModel, Field, confloat + + +class StabilityFormat(str, Enum): + png = 'png' + jpeg = 'jpeg' + webp = 'webp' + + +class StabilityAspectRatio(str, Enum): + ratio_1_1 = "1:1" + ratio_16_9 = "16:9" + ratio_9_16 = "9:16" + ratio_3_2 = "3:2" + ratio_2_3 = "2:3" + ratio_5_4 = "5:4" + ratio_4_5 = "4:5" + ratio_21_9 = "21:9" + ratio_9_21 = "9:21" + + +def get_stability_style_presets(include_none=True): + presets = [] + if include_none: + presets.append("None") + return presets + [x.value for x in StabilityStylePreset] + + +class StabilityStylePreset(str, Enum): + _3d_model = "3d-model" + analog_film = "analog-film" + anime = "anime" + cinematic = "cinematic" + comic_book = "comic-book" + digital_art = "digital-art" + enhance = "enhance" + fantasy_art = "fantasy-art" + isometric = "isometric" + line_art = "line-art" + low_poly = "low-poly" + modeling_compound = "modeling-compound" + neon_punk = "neon-punk" + origami = "origami" + photographic = "photographic" + pixel_art = "pixel-art" + tile_texture = "tile-texture" + + +class StabilityStableUltraRequest(BaseModel): + prompt: str = Field(...) + negative_prompt: Optional[str] = Field(None) + aspect_ratio: Optional[str] = Field(None) + seed: Optional[int] = Field(None) + output_format: Optional[str] = Field(StabilityFormat.png.value) + image: Optional[str] = Field(None) + style_preset: Optional[str] = Field(None) + strength: Optional[confloat(ge=0.0, le=1.0)] = Field(None) + + +class StabilityStableUltraResponse(BaseModel): + image: Optional[str] = Field(None) + finish_reason: Optional[str] = Field(None) + seed: Optional[int] = Field(None) + diff --git a/comfy_api_nodes/nodes_stability.py b/comfy_api_nodes/nodes_stability.py new file mode 100644 index 000000000..44d8aa17a --- /dev/null +++ b/comfy_api_nodes/nodes_stability.py @@ -0,0 +1,160 @@ +from inspect import cleandoc +from comfy.comfy_types.node_typing import IO +from comfy_api_nodes.apis.stability_api import ( + StabilityStableUltraRequest, + StabilityStableUltraResponse, + StabilityAspectRatio, + get_stability_style_presets, +) +from comfy_api_nodes.apis.client import ( + ApiEndpoint, + HttpMethod, + SynchronousOperation, +) +from comfy_api_nodes.apinode_utils import ( + bytesio_to_image_tensor, + tensor_to_bytesio, +) + +import torch +import base64 +from io import BytesIO + + + +class StabilityStableImageUltraNode: + """ + Generates images synchronously based on prompt and resolution. + """ + + RETURN_TYPES = (IO.IMAGE,) + DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value + FUNCTION = "api_call" + API_NODE = True + CATEGORY = "api node/image/stability" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "prompt": ( + IO.STRING, + { + "multiline": True, + "default": "", + "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines" + + "What you wish to see in the output image. A strong, descriptive prompt that clearly defines" + + "elements, colors, and subjects will lead to better results. " + + "To control the weight of a given word use the format `(word:weight)`," + + "where `word` is the word you'd like to control the weight of and `weight`" + + "is a value between 0 and 1. For example: `The sky was a crisp (blue:0.3) and (green:0.8)`" + + "would convey a sky that was blue and green, but more green than blue." + }, + ), + "aspect_ratio": ([x.value for x in StabilityAspectRatio], + { + "default": StabilityAspectRatio.ratio_1_1, + "tooltip": "Aspect ratio of generated image.", + }, + ), + "style_preset": (get_stability_style_presets(), + { + "tooltip": "Optional desired style of generated image.", + }, + ), + "seed": ( + IO.INT, + { + "default": 0, + "min": 0, + "max": 4294967294, + "control_after_generate": True, + "tooltip": "The random seed used for creating the noise.", + }, + ), + }, + "optional": { + "image": (IO.IMAGE,), + "negative_prompt": ( + IO.STRING, + { + "default": "", + "forceInput": True, + "tooltip": "A blurb of text describing what you do not wish to see in the output image. This is an advanced feature." + }, + ), + "image_denoise": ( + IO.FLOAT, + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.", + }, + ), + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + }, + } + + def api_call(self, prompt: str, aspect_ratio: str, style_preset: str, seed: int, + negative_prompt: str=None, image: torch.Tensor = None, image_denoise: float=None, + auth_token=None): + # prepare image binary if image present + image_binary = None + if image is not None: + image_binary = tensor_to_bytesio(image, 1504 * 1504).read() + else: + image_denoise = None + + if not negative_prompt: + negative_prompt = None + if style_preset == "None": + style_preset = None + + files = { + "image": image_binary + } + + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path="/proxy/stability/v2beta/stable-image/generate/ultra", + method=HttpMethod.POST, + request_model=StabilityStableUltraRequest, + response_model=StabilityStableUltraResponse, + ), + request=StabilityStableUltraRequest( + prompt=prompt, + negative_prompt=negative_prompt, + aspect_ratio=aspect_ratio, + seed=seed, + strength=image_denoise, + style_preset=style_preset, + ), + files=files, + content_type="multipart/form-data", + auth_token=auth_token, + ) + response_api = operation.execute() + + if response_api.finish_reason != "SUCCESS": + raise Exception(f"Stable Image Ultra generation failed: {response_api.finish_reason}.") + + image_data = base64.b64decode(response_api.image) + returned_image = bytesio_to_image_tensor(BytesIO(image_data)) + + return (returned_image,) + + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "StabilityStableImageUltraNode": StabilityStableImageUltraNode, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "StabilityStableImageUltraNode": "Stabiliy Stable Image Ultra", +} diff --git a/nodes.py b/nodes.py index 9ac9be365..9d5ebd5c0 100644 --- a/nodes.py +++ b/nodes.py @@ -2272,6 +2272,7 @@ def init_builtin_extra_nodes(): "nodes_luma.py", "nodes_recraft.py", "nodes_pixverse.py", + "nodes_stability.py", "nodes_pika.py", ]