Grok Imagine Image 2.0 Segment Edit

What it is

Region-targeted editing for Grok Imagine Image 2.0, sitting at KIE Next / Image / xAI / Grok Imagine. You don't draw a mask. You refer to regions by index - the numbers produced by the Segment Map node - and describe the change. "Change the background to a sunset beach while preserving the subject" is the author's own placeholder prompt, and it's a decent summary of the whole idea: one part of the frame changes, the rest is meant to survive.

Why it's here at all, given that a local inpaint with a hand-painted mask is free and precise: because the thing most beginners can't do locally is good regional work. A mask tells your local model what to redraw but not how to blend it, and the result is a visible seam or a subject that quietly changed shape. Telling a closed model "this region, but a sunset beach" hands it the whole judgement call. That's worth a credit on a keepable image and not worth one on a rough draft.

How it works

The node submits a task that references a source task ID plus your region indices and prompt, polls until KIE reports completion, downloads the result and returns it as an IMAGE. Everything is server-side; the only thing your machine does is upload nothing, wait, and decode a PNG.

The task-ID requirement is the mechanism worth understanding. This model doesn't accept an image upload. It accepts a pointer to something KIE already generated, which is why it pairs with Text To Image and Segment Map and not with your local LoadImage. Per KIE's docs, the valid source is a task from grok-imagine-image-2-0/text-to-image, or a task from segment-map - both of which are nodes in this folder whose task_id output you wire directly in.

The inputs

Required:

• task_id - the source task. Wire it from the Text To Image or Segment Map node in the same graph. It's a string, and pasting the wrong one is the single most likely way to get a confusing result.
• prompt - the change you want, in words. Keep it scoped to the region you selected; the model can't read your intent, only your prompt.

Optional:

• mask_indexs - a JSON array of segment indices, defaulting to [1, 2]. This is your mask, as numbers. Get the numbers from looking at the segment map image, and note that the JSON is a string widget in ComfyUI, so it must parse - [1, 2] works, 1,2 doesn't.
• timeout_seconds (default 1200) and callback_url (leave empty for normal ComfyUI use).

Outputs

image is the edited result. url and all_urls_json are the hosted copies, raw_json is the provider's response - the thing you read when a task fails and you want the actual reason - and credits_consumed / credits_left report cost and remaining balance, also displayed on the node after each successful run. The task_id output is the task for this edit, which you can feed into another Segment Edit if you're doing a multi-pass pass over the same image.

Install

ComfyUI Manager → KIE.ai Nodes Next, or the two-liner:

cd ComfyUI/custom_nodes
git clone https://github.com/felipederosilva/ComfyUI-KIE-Nodes-Next

Then install the new folder's requirements.txt into the same Python environment ComfyUI runs in. The list is small (requests, Pillow, numpy, PyYAML) and there are no weights to download, because the model is hosted. Restart ComfyUI, then complete the one-time key step at Settings → KIE.ai Nodes Next → Connection - the key is validated against KIE before it's saved and only a masked fingerprint stays visible afterwards. On headless boxes, export KIE_API_KEY instead.

The pack generates its menu from KIE's API docs index on startup and re-checks about daily, reloading the frontend once during the first sync. Empty folder? Refresh the catalog, don't reinstall.

Gotchas

• Indices are not stable across images. Each segment map numbers its own regions; reusing [1, 2] after regenerating the source is how you edit a lamp instead of a shirt.
• mask_indexs is JSON-as-string. A malformed array throws a parse error from the node before anything is sent - which is at least free.
• Three calls, one picture. Text-to-image, segment map, segment edit. Budget like it's a pipeline, because it is.
• Prompt scope discipline. Ask for the region's change only; a prompt describing the whole scene gives the model permission to re-imagine the whole scene.
• No key → No KIE API key configured. Settings step.
• Local masking is still the right tool for pixel-precision work - this is for "make it look intentional", not "remove this exact hair".