From aab6db2a6e5de4a0f8d0a5e07f7b4932bac4341f Mon Sep 17 00:00:00 2001 From: Dreaming Computers Date: Sat, 7 Mar 2026 14:18:43 +0700 Subject: [PATCH] Add Grok Imagine image nodes with task chaining --- README.md | 10 +- kie_api/grok_imagine_i2i.py | 88 +++++++++++++++++ kie_api/grok_imagine_t2i.py | 66 +++++++++++++ kie_api/images.py | 25 +++++ nodes.py | 141 +++++++++++++++++++++++++++ web/docs/KIE_GrokImagine_I2I.md | 17 ++++ web/docs/KIE_GrokImagine_I2I_Spec.md | 8 +- web/docs/KIE_GrokImagine_T2I.md | 17 ++++ web/docs/KIE_GrokImagine_T2I_Spec.md | 6 +- web/docs/README.md | 2 + 10 files changed, 372 insertions(+), 8 deletions(-) create mode 100644 kie_api/grok_imagine_i2i.py create mode 100644 kie_api/grok_imagine_t2i.py create mode 100644 web/docs/KIE_GrokImagine_I2I.md create mode 100644 web/docs/KIE_GrokImagine_T2I.md diff --git a/README.md b/README.md index 64d9993..4d20e17 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,12 @@ This node pack currently includes the following nodes: - **Flux 2 Image-to-Image (Pro/Flex)** - Image-to-image node with a model dropdown for Pro or Flex. - Accepts 1–8 input images via ComfyUI batch. +- **Grok Imagine Text-to-Image** + - Generates one or more images from a text prompt. + - Returns the image batch plus a `task_id` for downstream Grok I2V chaining. +- **Grok Imagine Image-to-Image** + - Generates one or more edited images from a single uploaded source image. + - Returns the image batch plus a `task_id` for downstream Grok I2V chaining. ## Video Nodes - **Kling 2.6 Image-to-Video** @@ -202,7 +208,9 @@ This is useful for quickly evaluating model preference by prompt type, fidelity, ## Changelog - 2026-03-07: Added the Grok Imagine text-to-video node and synced docs/spec references. - 2026-03-07: Added the Grok Imagine image-to-video node and synced docs/spec references. -- 2026-03-07: Added Grok Imagine text-to-image and image-to-image API spec docs for future node implementation. +- 2026-03-07: Added the Grok Imagine text-to-image node and synced docs/spec references. +- 2026-03-07: Added the Grok Imagine image-to-image node and synced docs/spec references. +- 2026-03-07: Added Grok Imagine text-to-image and image-to-image API spec docs. - 2026-02-28: Added new example workflow `KIE-AI-Banana-Pro-Banana-Pro.json` to repository docs. - 2026-02-28: Documentation clarified Nano Banana Pro (up to 8 images) and Nano Banana 2 (up to 14 images + optional Google web search). Bumped version to 0.1.8. - 2026-02-28: Updated Nano Banana Pro + Nano Banana 2 payload behavior to always send `image_input` (empty list when no images are connected). diff --git a/kie_api/grok_imagine_i2i.py b/kie_api/grok_imagine_i2i.py new file mode 100644 index 0000000..70cec2c --- /dev/null +++ b/kie_api/grok_imagine_i2i.py @@ -0,0 +1,88 @@ +"""Grok Imagine image-to-image helper.""" + +import time + +import torch + +from .auth import _load_api_key +from .credits import _log_remaining_credits +from .images import _download_images_as_batch +from .jobs import _create_task, _poll_task_until_complete +from .log import _log +from .results import _extract_result_urls +from .upload import _image_tensor_to_png_bytes, _truncate_url, _upload_image +from .validation import _validate_image_tensor_batch + + +MODEL_NAME = "grok-imagine/image-to-image" +PROMPT_MAX_LENGTH = 390000 + + +def _validate_optional_prompt(prompt: str) -> str: + if not isinstance(prompt, str): + raise RuntimeError("prompt must be a string.") + + prompt_value = prompt or "" + if len(prompt_value) > PROMPT_MAX_LENGTH: + raise RuntimeError(f"Prompt exceeds the maximum length of {PROMPT_MAX_LENGTH} characters.") + + return prompt_value + + +def run_grok_imagine_i2i( + images: torch.Tensor, + prompt: str, + poll_interval_s: float, + timeout_s: int, + log: bool, +) -> tuple[torch.Tensor, str]: + images = _validate_image_tensor_batch(images) + prompt_value = _validate_optional_prompt(prompt) + + if images.shape[0] > 1: + _log(log, f"More than 1 image provided ({images.shape[0]}); only the first will be used.") + + api_key = _load_api_key() + + _log(log, "Uploading source image for Grok Imagine I2I...") + png_bytes = _image_tensor_to_png_bytes(images[0]) + image_url = _upload_image(api_key, png_bytes) + _log(log, f"Image upload success: {_truncate_url(image_url)}") + + input_payload: dict[str, object] = { + "image_urls": [image_url], + } + if prompt_value.strip(): + input_payload["prompt"] = prompt_value + + payload = { + "model": MODEL_NAME, + "input": input_payload, + } + + _log(log, "Creating Grok Imagine I2I task...") + start_time = time.time() + task_id, create_response_text = _create_task(api_key, payload) + _log( + log, + f"createTask response (elapsed={time.time() - start_time:.1f}s): {create_response_text}", + ) + _log(log, f"Task created with ID {task_id}. Polling for completion...") + + record_data = _poll_task_until_complete( + api_key, + task_id, + poll_interval_s, + timeout_s, + log, + start_time, + ) + + result_urls = _extract_result_urls(record_data) + _log(log, f"Result URLs: {result_urls}") + _log(log, f"Downloading {len(result_urls)} Grok Imagine I2I result image(s)...") + image_batch = _download_images_as_batch(result_urls) + _log(log, "Images downloaded and decoded.") + + _log_remaining_credits(log, record_data, api_key, _log) + return image_batch, task_id diff --git a/kie_api/grok_imagine_t2i.py b/kie_api/grok_imagine_t2i.py new file mode 100644 index 0000000..08e1e89 --- /dev/null +++ b/kie_api/grok_imagine_t2i.py @@ -0,0 +1,66 @@ +"""Grok Imagine text-to-image helper.""" + +import time + +import torch + +from .auth import _load_api_key +from .credits import _log_remaining_credits +from .images import _download_images_as_batch +from .jobs import _create_task, _poll_task_until_complete +from .log import _log +from .results import _extract_result_urls +from .validation import _validate_prompt + + +MODEL_NAME = "grok-imagine/text-to-image" +PROMPT_MAX_LENGTH = 5000 +ASPECT_RATIO_OPTIONS = ["2:3", "3:2", "1:1", "9:16", "16:9"] + + +def run_grok_imagine_t2i( + prompt: str, + aspect_ratio: str, + poll_interval_s: float, + timeout_s: int, + log: bool, +) -> tuple[torch.Tensor, str]: + _validate_prompt(prompt, max_length=PROMPT_MAX_LENGTH) + if aspect_ratio not in ASPECT_RATIO_OPTIONS: + raise RuntimeError("Invalid aspect_ratio. Use the pinned enum options.") + + api_key = _load_api_key() + payload = { + "model": MODEL_NAME, + "input": { + "prompt": prompt, + "aspect_ratio": aspect_ratio, + }, + } + + _log(log, "Creating Grok Imagine T2I task...") + start_time = time.time() + task_id, create_response_text = _create_task(api_key, payload) + _log( + log, + f"createTask response (elapsed={time.time() - start_time:.1f}s): {create_response_text}", + ) + _log(log, f"Task created with ID {task_id}. Polling for completion...") + + record_data = _poll_task_until_complete( + api_key, + task_id, + poll_interval_s, + timeout_s, + log, + start_time, + ) + + result_urls = _extract_result_urls(record_data) + _log(log, f"Result URLs: {result_urls}") + _log(log, f"Downloading {len(result_urls)} Grok Imagine T2I result image(s)...") + image_batch = _download_images_as_batch(result_urls) + _log(log, "Images downloaded and decoded.") + + _log_remaining_credits(log, record_data, api_key, _log) + return image_batch, task_id diff --git a/kie_api/images.py b/kie_api/images.py index bacfb8e..a555299 100644 --- a/kie_api/images.py +++ b/kie_api/images.py @@ -40,3 +40,28 @@ def _download_image(url: str) -> bytes: raise RuntimeError(f"Failed to download result image (status code {response.status_code}).") return response.content + + +def _stack_image_tensors(image_tensors: list[torch.Tensor]) -> torch.Tensor: + """Combine decoded image tensors into one ComfyUI IMAGE batch.""" + if not image_tensors: + raise RuntimeError("No result images were returned.") + + base_shape = tuple(image_tensors[0].shape[1:]) + for idx, tensor in enumerate(image_tensors[1:], start=1): + if tuple(tensor.shape[1:]) != base_shape: + raise RuntimeError( + "Result images returned inconsistent sizes; cannot combine them into one IMAGE batch " + f"(mismatch at index {idx})." + ) + + return torch.cat(image_tensors, dim=0) + + +def _download_images_as_batch(urls: list[str]) -> torch.Tensor: + """Download multiple image URLs and return a single IMAGE batch tensor.""" + if not urls: + raise RuntimeError("No result image URLs were returned.") + + image_tensors = [_image_bytes_to_tensor(_download_image(url)) for url in urls] + return _stack_image_tensors(image_tensors) diff --git a/nodes.py b/nodes.py index 846bfaa..f1a7ea7 100644 --- a/nodes.py +++ b/nodes.py @@ -72,6 +72,10 @@ from .kie_api.flux2_i2i import ( RESOLUTION_OPTIONS as FLUX2_RESOLUTION_OPTIONS, run_flux2_i2i, ) +from .kie_api.grok_imagine_t2i import ( + ASPECT_RATIO_OPTIONS as GROK_T2I_ASPECT_RATIO_OPTIONS, + run_grok_imagine_t2i, +) from .kie_api.grok_imagine_t2v import ( ASPECT_RATIO_OPTIONS as GROK_T2V_ASPECT_RATIO_OPTIONS, DURATION_OPTIONS as GROK_T2V_DURATION_OPTIONS, @@ -79,6 +83,7 @@ from .kie_api.grok_imagine_t2v import ( RESOLUTION_OPTIONS as GROK_T2V_RESOLUTION_OPTIONS, run_grok_imagine_t2v_video, ) +from .kie_api.grok_imagine_i2i import run_grok_imagine_i2i from .kie_api.grok_imagine_i2v import ( DURATION_OPTIONS as GROK_I2V_DURATION_OPTIONS, MODE_OPTIONS as GROK_I2V_MODE_OPTIONS, @@ -426,6 +431,138 @@ Outputs: return (image_tensor,) +class KIE_GrokImagine_T2I: + HELP = """ +KIE Grok Imagine (Text-to-Image) + +Generate one or more images from a text prompt using Grok Imagine. + +Inputs: +- prompt: Text prompt (required, up to 5000 chars) +- aspect_ratio: 2:3, 3:2, 1:1, 9:16, or 16:9 +- poll_interval_s / timeout_s / log +- retry_on_fail / max_retries / retry_backoff_s + +Outputs: +- IMAGE: ComfyUI image batch (BHWC float32 0-1) +- STRING: task_id for Grok image-reference chaining +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + }, + "optional": { + "aspect_ratio": ("COMBO", {"options": GROK_T2I_ASPECT_RATIO_OPTIONS, "default": "1:1"}), + "log": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("image", "task_id") + FUNCTION = "generate" + CATEGORY = "kie/api" + + def generate( + self, + prompt: str, + aspect_ratio: str = "1:1", + log: bool = True, + poll_interval_s: float = 10.0, + timeout_s: int = 300, + retry_on_fail: bool = True, + max_retries: int = 2, + retry_backoff_s: float = 3.0, + ): + attempts = max_retries + 1 if retry_on_fail else 1 + attempts = max(attempts, 1) + backoff = retry_backoff_s if retry_backoff_s >= 0 else 0.0 + + for attempt in range(1, attempts + 1): + try: + image_output, task_id = run_grok_imagine_t2i( + prompt=prompt, + aspect_ratio=aspect_ratio, + poll_interval_s=poll_interval_s, + timeout_s=timeout_s, + log=log, + ) + return (image_output, task_id) + except TransientKieError: + if not retry_on_fail or attempt >= attempts: + raise + _log(log, f"Retrying (attempt {attempt + 1}/{attempts}) after {backoff}s") + time.sleep(backoff) + + +class KIE_GrokImagine_I2I: + HELP = """ +KIE Grok Imagine (Image-to-Image) + +Generate one or more images from a source image using Grok Imagine. + +Inputs: +- images: Source image batch (first image used) +- prompt: Optional prompt (up to 390000 chars) +- poll_interval_s / timeout_s / log +- retry_on_fail / max_retries / retry_backoff_s + +Outputs: +- IMAGE: ComfyUI image batch (BHWC float32 0-1) +- STRING: task_id for Grok image-reference chaining +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + }, + "optional": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + "log": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("image", "task_id") + FUNCTION = "generate" + CATEGORY = "kie/api" + + def generate( + self, + images: torch.Tensor, + prompt: str = "", + log: bool = True, + poll_interval_s: float = 10.0, + timeout_s: int = 300, + retry_on_fail: bool = True, + max_retries: int = 2, + retry_backoff_s: float = 3.0, + ): + attempts = max_retries + 1 if retry_on_fail else 1 + attempts = max(attempts, 1) + backoff = retry_backoff_s if retry_backoff_s >= 0 else 0.0 + + for attempt in range(1, attempts + 1): + try: + image_output, task_id = run_grok_imagine_i2i( + images=images, + prompt=prompt, + poll_interval_s=poll_interval_s, + timeout_s=timeout_s, + log=log, + ) + return (image_output, task_id) + except TransientKieError: + if not retry_on_fail or attempt >= attempts: + raise + _log(log, f"Retrying (attempt {attempt + 1}/{attempts}) after {backoff}s") + time.sleep(backoff) + + class KIE_Kling25_I2V_Pro: HELP = """ KIE Kling 2.5 I2V Pro @@ -1827,6 +1964,8 @@ NODE_CLASS_MAPPINGS = { "KIE_NanoBanana2_Image": KIE_NanoBanana2_Image, "KIE_Seedream45_TextToImage": KIE_Seedream45_TextToImage, "KIE_Seedream45_Edit": KIE_Seedream45_Edit, + "KIE_GrokImagine_T2I": KIE_GrokImagine_T2I, + "KIE_GrokImagine_I2I": KIE_GrokImagine_I2I, "KIE_SeedanceV1Pro_Fast_I2V": KIE_SeedanceV1Pro_Fast_I2V, "KIE_Seedance15Pro_I2V": KIE_Seedance15Pro_I2V, "KIE_Kling25_I2V_Pro": KIE_Kling25_I2V_Pro, @@ -1853,6 +1992,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "KIE_NanoBanana2_Image": "Nano Banana 2", "KIE_Seedream45_TextToImage": "KIE Seedream 4.5 Text-To-Image", "KIE_Seedream45_Edit": "KIE Seedream 4.5 Edit", + "KIE_GrokImagine_T2I": "KIE Grok Imagine (T2I)", + "KIE_GrokImagine_I2I": "KIE Grok Imagine (I2I)", "KIE_SeedanceV1Pro_Fast_I2V": "KIE Seedance V1 Pro Fast (I2V)", "KIE_Seedance15Pro_I2V": "KIE Seedance 1.5 Pro (I2V/T2V)", "KIE_Kling25_I2V_Pro": "KIE Kling 2.5 I2V Pro", diff --git a/web/docs/KIE_GrokImagine_I2I.md b/web/docs/KIE_GrokImagine_I2I.md new file mode 100644 index 0000000..025e31b --- /dev/null +++ b/web/docs/KIE_GrokImagine_I2I.md @@ -0,0 +1,17 @@ +# KIE Grok Imagine (I2I) + +Generate one or more images from a source image using Grok Imagine. + +## Inputs +- `images` (IMAGE, required): Source image batch. Only the first image is uploaded. +- `prompt` (STRING, optional): Edit/style prompt text. Max length: `390000`. +- `log` (BOOLEAN, optional): Enable helper logging (default: `true`). + +## Outputs +- `image` (IMAGE): ComfyUI image batch (BHWC float32, 0..1). +- `task_id` (STRING): KIE task id from `createTask`, for chaining into Grok I2V. + +## Notes +- If multiple input images are connected, the node uploads only the first image and logs that choice. +- The node downloads all returned `resultUrls` and returns them as one `IMAGE` batch. +- This is the intended upstream source for Grok I2V `task_id + index` chaining. diff --git a/web/docs/KIE_GrokImagine_I2I_Spec.md b/web/docs/KIE_GrokImagine_I2I_Spec.md index 7c224cd..9bd6582 100644 --- a/web/docs/KIE_GrokImagine_I2I_Spec.md +++ b/web/docs/KIE_GrokImagine_I2I_Spec.md @@ -1,7 +1,7 @@ # KIE Grok Imagine (I2I) API Spec ## Status -Reference spec for future ComfyUI node work. This repository does not expose a Grok Imagine image-to-image node yet. +Reference spec for the implemented Grok Imagine image-to-image node. See [`KIE_GrokImagine_I2I.md`](KIE_GrokImagine_I2I.md) for the ComfyUI node surface. ## Endpoint - Method: `POST` @@ -75,7 +75,7 @@ If `callBackUrl` is present, KIE posts task completion payloads to that URL for - `msg` ## ComfyUI Mapping Notes -Recommended node shape for this repo: +Implemented node shape in this repo: - `images`: `IMAGE`, required. Upload the first image only. - `prompt`: `STRING` multiline, default empty string. - `log`: `BOOLEAN`, optional. @@ -83,11 +83,11 @@ Recommended node shape for this repo: - `IMAGE` - `task_id` (`STRING`) -Recommended validation behavior: +Validation behavior: - Require at least one input image. - If multiple ComfyUI images are connected, upload the first image only. -Recommended output behavior: +Output behavior: - Download all `resultUrls` and return them as a ComfyUI `IMAGE` batch if the endpoint returns multiple images. - Also return the `task_id` from task creation so Grok I2V can reference one of the generated images via `task_id + index`. diff --git a/web/docs/KIE_GrokImagine_T2I.md b/web/docs/KIE_GrokImagine_T2I.md new file mode 100644 index 0000000..938f2c8 --- /dev/null +++ b/web/docs/KIE_GrokImagine_T2I.md @@ -0,0 +1,17 @@ +# KIE Grok Imagine (T2I) + +Generate one or more images from a text prompt using Grok Imagine. + +## Inputs +- `prompt` (STRING, required): Generation prompt text. Max length: `5000`. +- `aspect_ratio` (COMBO, optional): `2:3`, `3:2`, `1:1`, `9:16`, `16:9` (default: `1:1`). +- `log` (BOOLEAN, optional): Enable helper logging (default: `true`). + +## Outputs +- `image` (IMAGE): ComfyUI image batch (BHWC float32, 0..1). +- `task_id` (STRING): KIE task id from `createTask`, for chaining into Grok I2V. + +## Notes +- The node downloads all returned `resultUrls` and returns them as one `IMAGE` batch. +- This is the intended upstream source for Grok I2V `task_id + index` chaining. +- Polling and timeout are handled internally by shared helpers. diff --git a/web/docs/KIE_GrokImagine_T2I_Spec.md b/web/docs/KIE_GrokImagine_T2I_Spec.md index 34266cf..6daf32b 100644 --- a/web/docs/KIE_GrokImagine_T2I_Spec.md +++ b/web/docs/KIE_GrokImagine_T2I_Spec.md @@ -1,7 +1,7 @@ # KIE Grok Imagine (T2I) API Spec ## Status -Reference spec for future ComfyUI node work. This repository does not expose a Grok Imagine text-to-image node yet. +Reference spec for the implemented Grok Imagine text-to-image node. See [`KIE_GrokImagine_T2I.md`](KIE_GrokImagine_T2I.md) for the ComfyUI node surface. ## Endpoint - Method: `POST` @@ -73,7 +73,7 @@ If `callBackUrl` is present, KIE posts task completion payloads to that URL for - `msg` ## ComfyUI Mapping Notes -Recommended node shape for this repo: +Implemented node shape in this repo: - `prompt`: `STRING` multiline, required. - `aspect_ratio`: `COMBO`, optional. - `log`: `BOOLEAN`, optional. @@ -81,7 +81,7 @@ Recommended node shape for this repo: - `IMAGE` - `task_id` (`STRING`) -Recommended output behavior: +Output behavior: - Download all `resultUrls` and return them as a ComfyUI `IMAGE` batch if the endpoint returns multiple images. - Also return the `task_id` from task creation so Grok I2V can reference one of the generated images via `task_id + index`. diff --git a/web/docs/README.md b/web/docs/README.md index 13b65d0..d75e8b4 100644 --- a/web/docs/README.md +++ b/web/docs/README.md @@ -28,6 +28,8 @@ This folder contains node documentation for the ComfyUI KIE API pack. - Nano Banana Pro: [`KIE_NanoBananaPro_Image.md`](KIE_NanoBananaPro_Image.md) - Nano Banana 2: [`KIE_NanoBanana2_Image.md`](KIE_NanoBanana2_Image.md) - Flux 2 I2I: [`KIE_Flux2_I2I.md`](KIE_Flux2_I2I.md) +- Grok Imagine T2I: [`KIE_GrokImagine_T2I.md`](KIE_GrokImagine_T2I.md) +- Grok Imagine I2I: [`KIE_GrokImagine_I2I.md`](KIE_GrokImagine_I2I.md) - Seedream 4.5 Text-to-Image: [`KIE_Seedream45_TextToImage.md`](KIE_Seedream45_TextToImage.md) - Seedream 4.5 Edit: [`KIE_Seedream45_Edit.md`](KIE_Seedream45_Edit.md)