Add Grok Imagine image nodes with task chaining

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).

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

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

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)

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",

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.

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`.
 

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.

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`.
 

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)