from inspect import cleandoc from comfy.comfy_types.node_typing import IO from comfy_api_nodes.apis.stability_api import ( StabilityUpscaleConservativeRequest, StabilityUpscaleCreativeRequest, StabilityAsyncResponse, StabilityResultsGetResponse, StabilityStable3_5Request, StabilityStableUltraRequest, StabilityStableUltraResponse, StabilityAspectRatio, Stability_SD3_5_Model, Stability_SD3_5_GenerationMode, get_stability_style_presets, ) from comfy_api_nodes.apis.client import ( ApiEndpoint, HttpMethod, SynchronousOperation, PollingOperation, EmptyRequest, ) from comfy_api_nodes.apinode_utils import ( bytesio_to_image_tensor, tensor_to_bytesio, validate_string, ) import torch import base64 from io import BytesIO from enum import Enum class StabilityPollStatus(str, Enum): finished = "finished" in_progress = "in_progress" failed = "failed" def get_async_dummy_status(x: StabilityResultsGetResponse): if x.name is not None or x.errors is not None: return StabilityPollStatus.failed elif x.finish_reason is not None: return StabilityPollStatus.finished return StabilityPollStatus.in_progress 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 AI" @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", "comfy_api_key": "API_KEY_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, **kwargs): validate_string(prompt, strip_whitespace=False) # prepare image binary if image present image_binary = None if image is not None: image_binary = tensor_to_bytesio(image, total_pixels=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_kwargs=kwargs, ) 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,) class StabilityStableImageSD_3_5Node: """ 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 AI" @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 elements, colors, and subjects will lead to better results." }, ), "model": ([x.value for x in Stability_SD3_5_Model],), "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.", }, ), "cfg_scale": ( IO.FLOAT, { "default": 4.0, "min": 1.0, "max": 10.0, "step": 0.1, "tooltip": "How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)", }, ), "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": "Keywords of 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", "comfy_api_key": "API_KEY_COMFY_ORG", }, } def api_call(self, model: str, prompt: str, aspect_ratio: str, style_preset: str, seed: int, cfg_scale: float, negative_prompt: str=None, image: torch.Tensor = None, image_denoise: float=None, **kwargs): validate_string(prompt, strip_whitespace=False) # prepare image binary if image present image_binary = None mode = Stability_SD3_5_GenerationMode.text_to_image if image is not None: image_binary = tensor_to_bytesio(image, total_pixels=1504*1504).read() mode = Stability_SD3_5_GenerationMode.image_to_image aspect_ratio = None 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/sd3", method=HttpMethod.POST, request_model=StabilityStable3_5Request, response_model=StabilityStableUltraResponse, ), request=StabilityStable3_5Request( prompt=prompt, negative_prompt=negative_prompt, aspect_ratio=aspect_ratio, seed=seed, strength=image_denoise, style_preset=style_preset, cfg_scale=cfg_scale, model=model, mode=mode, ), files=files, content_type="multipart/form-data", auth_kwargs=kwargs, ) response_api = operation.execute() if response_api.finish_reason != "SUCCESS": raise Exception(f"Stable Diffusion 3.5 Image 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,) class StabilityUpscaleConservativeNode: """ Upscale image with minimal alterations to 4K resolution. """ RETURN_TYPES = (IO.IMAGE,) DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value FUNCTION = "api_call" API_NODE = True CATEGORY = "api node/image/Stability AI" @classmethod def INPUT_TYPES(s): return { "required": { "image": (IO.IMAGE,), "prompt": ( IO.STRING, { "multiline": True, "default": "", "tooltip": "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." }, ), "creativity": ( IO.FLOAT, { "default": 0.35, "min": 0.2, "max": 0.5, "step": 0.01, "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init image.", }, ), "seed": ( IO.INT, { "default": 0, "min": 0, "max": 4294967294, "control_after_generate": True, "tooltip": "The random seed used for creating the noise.", }, ), }, "optional": { "negative_prompt": ( IO.STRING, { "default": "", "forceInput": True, "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature." }, ), }, "hidden": { "auth_token": "AUTH_TOKEN_COMFY_ORG", "comfy_api_key": "API_KEY_COMFY_ORG", }, } def api_call(self, image: torch.Tensor, prompt: str, creativity: float, seed: int, negative_prompt: str=None, **kwargs): validate_string(prompt, strip_whitespace=False) image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read() if not negative_prompt: negative_prompt = None files = { "image": image_binary } operation = SynchronousOperation( endpoint=ApiEndpoint( path="/proxy/stability/v2beta/stable-image/upscale/conservative", method=HttpMethod.POST, request_model=StabilityUpscaleConservativeRequest, response_model=StabilityStableUltraResponse, ), request=StabilityUpscaleConservativeRequest( prompt=prompt, negative_prompt=negative_prompt, creativity=round(creativity,2), seed=seed, ), files=files, content_type="multipart/form-data", auth_kwargs=kwargs, ) response_api = operation.execute() if response_api.finish_reason != "SUCCESS": raise Exception(f"Stability Upscale Conservative 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,) class StabilityUpscaleCreativeNode: """ Upscale image with minimal alterations to 4K resolution. """ RETURN_TYPES = (IO.IMAGE,) DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value FUNCTION = "api_call" API_NODE = True CATEGORY = "api node/image/Stability AI" @classmethod def INPUT_TYPES(s): return { "required": { "image": (IO.IMAGE,), "prompt": ( IO.STRING, { "multiline": True, "default": "", "tooltip": "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." }, ), "creativity": ( IO.FLOAT, { "default": 0.3, "min": 0.1, "max": 0.5, "step": 0.01, "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init 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": { "negative_prompt": ( IO.STRING, { "default": "", "forceInput": True, "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature." }, ), }, "hidden": { "auth_token": "AUTH_TOKEN_COMFY_ORG", "comfy_api_key": "API_KEY_COMFY_ORG", }, } def api_call(self, image: torch.Tensor, prompt: str, creativity: float, style_preset: str, seed: int, negative_prompt: str=None, **kwargs): validate_string(prompt, strip_whitespace=False) image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read() 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/upscale/creative", method=HttpMethod.POST, request_model=StabilityUpscaleCreativeRequest, response_model=StabilityAsyncResponse, ), request=StabilityUpscaleCreativeRequest( prompt=prompt, negative_prompt=negative_prompt, creativity=round(creativity,2), style_preset=style_preset, seed=seed, ), files=files, content_type="multipart/form-data", auth_kwargs=kwargs, ) response_api = operation.execute() operation = PollingOperation( poll_endpoint=ApiEndpoint( path=f"/proxy/stability/v2beta/results/{response_api.id}", method=HttpMethod.GET, request_model=EmptyRequest, response_model=StabilityResultsGetResponse, ), poll_interval=3, completed_statuses=[StabilityPollStatus.finished], failed_statuses=[StabilityPollStatus.failed], status_extractor=lambda x: get_async_dummy_status(x), auth_kwargs=kwargs, ) response_poll: StabilityResultsGetResponse = operation.execute() if response_poll.finish_reason != "SUCCESS": raise Exception(f"Stability Upscale Creative generation failed: {response_poll.finish_reason}.") image_data = base64.b64decode(response_poll.result) returned_image = bytesio_to_image_tensor(BytesIO(image_data)) return (returned_image,) class StabilityUpscaleFastNode: """ Quickly upscales an image via Stability API call to 4x its original size; intended for upscaling low-quality/compressed images. """ RETURN_TYPES = (IO.IMAGE,) DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value FUNCTION = "api_call" API_NODE = True CATEGORY = "api node/image/Stability AI" @classmethod def INPUT_TYPES(s): return { "required": { "image": (IO.IMAGE,), }, "optional": { }, "hidden": { "auth_token": "AUTH_TOKEN_COMFY_ORG", "comfy_api_key": "API_KEY_COMFY_ORG", }, } def api_call(self, image: torch.Tensor, **kwargs): image_binary = tensor_to_bytesio(image, total_pixels=4096*4096).read() files = { "image": image_binary } operation = SynchronousOperation( endpoint=ApiEndpoint( path="/proxy/stability/v2beta/stable-image/upscale/fast", method=HttpMethod.POST, request_model=EmptyRequest, response_model=StabilityStableUltraResponse, ), request=EmptyRequest(), files=files, content_type="multipart/form-data", auth_kwargs=kwargs, ) response_api = operation.execute() if response_api.finish_reason != "SUCCESS": raise Exception(f"Stability Upscale Fast 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, "StabilityStableImageSD_3_5Node": StabilityStableImageSD_3_5Node, "StabilityUpscaleConservativeNode": StabilityUpscaleConservativeNode, "StabilityUpscaleCreativeNode": StabilityUpscaleCreativeNode, "StabilityUpscaleFastNode": StabilityUpscaleFastNode, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "StabilityStableImageUltraNode": "Stability AI Stable Image Ultra", "StabilityStableImageSD_3_5Node": "Stability AI Stable Diffusion 3.5 Image", "StabilityUpscaleConservativeNode": "Stability AI Upscale Conservative", "StabilityUpscaleCreativeNode": "Stability AI Upscale Creative", "StabilityUpscaleFastNode": "Stability AI Upscale Fast", }