From 2eaf09f50d1f8c2cfd382302bed36396a3633331 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:38:08 +0300 Subject: [PATCH] [Partner Nodes] feat(Grok): add Grok Imagine Image 2.0 model (#15496) Signed-off-by: Alexander Piskun --- comfy_api_nodes/apis/grok.py | 2 + comfy_api_nodes/nodes_grok.py | 95 ++++++++++++++++++++++++++++------- 2 files changed, 78 insertions(+), 19 deletions(-) diff --git a/comfy_api_nodes/apis/grok.py b/comfy_api_nodes/apis/grok.py index 526d8c8ab..82dfc8b82 100644 --- a/comfy_api_nodes/apis/grok.py +++ b/comfy_api_nodes/apis/grok.py @@ -9,6 +9,7 @@ class ImageGenerationRequest(BaseModel): seed: int = Field(...) response_format: str = Field("url") resolution: str = Field(...) + quality: str | None = Field(None) class InputUrlObject(BaseModel): @@ -28,6 +29,7 @@ class ImageEditRequest(BaseModel): seed: int = Field(...) response_format: str = Field("url") aspect_ratio: str | None = Field(...) + quality: str | None = Field(None) class VideoGenerationRequest(BaseModel): diff --git a/comfy_api_nodes/nodes_grok.py b/comfy_api_nodes/nodes_grok.py index 672a3e537..c8f0f20d5 100644 --- a/comfy_api_nodes/nodes_grok.py +++ b/comfy_api_nodes/nodes_grok.py @@ -36,6 +36,26 @@ _GROK_VIDEO_MODEL_API_IDS = { "grok-imagine-video-1.5": "grok-imagine-video-1.5", } +_GROK_IMAGE_MODEL_API_IDS = { + "grok-imagine-image-2.0": "grok-imagine-image-2.0", +} + +_GROK_IMAGE_QUALITY_MODELS = {"grok-imagine-image-2.0"} + +_GROK_IMAGE_QUALITY_OPTIONS = ["medium", "low"] + +_GROK_IMAGE_EDIT_MAX_IMAGES = { + "grok-imagine-image-2.0": 3, + "grok-imagine-image-pro": 1, + "grok-imagine-image-quality": 3, + "grok-imagine-image": 3, +} + +_GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE = { + "grok-imagine-image-quality", + "grok-imagine-image", +} + _GROK_VOICE_OPTIONS = [ "none", "ara", @@ -132,6 +152,7 @@ class GrokImageNode(IO.ComfyNode): IO.Combo.Input( "model", options=[ + "grok-imagine-image-2.0", "grok-imagine-image-quality", "grok-imagine-image-pro", "grok-imagine-image", @@ -181,6 +202,12 @@ class GrokImageNode(IO.ComfyNode): "actual results are nondeterministic regardless of seed.", ), IO.Combo.Input("resolution", options=["1K", "2K"], optional=True), + IO.Combo.Input( + "quality", + options=_GROK_IMAGE_QUALITY_OPTIONS, + optional=True, + tooltip="Quality level, supported only by the grok-imagine-image-2.0 model.", + ), ], outputs=[ IO.Image.Output(), @@ -192,12 +219,15 @@ class GrokImageNode(IO.ComfyNode): ], is_api_node=True, price_badge=IO.PriceBadge( - depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution"]), + depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution", "quality"]), expr=""" ( - $rate := widgets.model = "grok-imagine-image-quality" - ? (widgets.resolution = "1k" ? 0.05 : 0.07) - : ($contains(widgets.model, "pro") ? 0.07 : 0.02); + $is1k := widgets.resolution = "1k"; + $rate := widgets.model = "grok-imagine-image-2.0" + ? (widgets.quality = "low" ? ($is1k ? 0.04 : 0.06) : ($is1k ? 0.06 : 0.08)) + : (widgets.model = "grok-imagine-image-quality" + ? ($is1k ? 0.05 : 0.07) + : ($contains(widgets.model, "pro") ? 0.07 : 0.02)); {"type":"usd","usd": $rate * widgets.number_of_images} ) """, @@ -213,18 +243,20 @@ class GrokImageNode(IO.ComfyNode): number_of_images: int, seed: int, resolution: str = "1K", + quality: str = "medium", ) -> IO.NodeOutput: validate_string(prompt, strip_whitespace=True, min_length=1) response = await sync_op( cls, ApiEndpoint(path="/proxy/xai/v1/images/generations", method="POST"), data=ImageGenerationRequest( - model=model, + model=_GROK_IMAGE_MODEL_API_IDS.get(model, model), prompt=prompt, aspect_ratio=aspect_ratio, n=number_of_images, seed=seed, resolution=resolution.lower(), + quality=quality if model in _GROK_IMAGE_QUALITY_MODELS else None, ), response_model=ImageGenerationResponse, ) @@ -255,7 +287,9 @@ _GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS = [ ] -def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: bool): +def _grok_image_edit_model_inputs( + *, max_ref_images: int, with_aspect_ratio: bool, with_quality: bool = False, aspect_ratio_needs_multiple: bool = True +): inputs = [ IO.Autogrow.Input( "images", @@ -281,12 +315,18 @@ def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: boo display_mode=IO.NumberDisplay.number, ), ] + if with_quality: + inputs.append(IO.Combo.Input("quality", options=_GROK_IMAGE_QUALITY_OPTIONS)) if with_aspect_ratio: inputs.append( IO.Combo.Input( "aspect_ratio", options=_GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS, - tooltip="Only allowed when multiple images are connected.", + tooltip=( + "Only allowed when multiple images are connected." + if aspect_ratio_needs_multiple + else "Aspect ratio of the edited image." + ), ) ) return inputs @@ -451,6 +491,15 @@ class GrokImageEditNodeV2(IO.ComfyNode): IO.DynamicCombo.Input( "model", options=[ + IO.DynamicCombo.Option( + "grok-imagine-image-2.0", + _grok_image_edit_model_inputs( + max_ref_images=3, + with_aspect_ratio=True, + with_quality=True, + aspect_ratio_needs_multiple=False, + ), + ), IO.DynamicCombo.Option( "grok-imagine-image-quality", _grok_image_edit_model_inputs(max_ref_images=3, with_aspect_ratio=True), @@ -488,18 +537,23 @@ class GrokImageEditNodeV2(IO.ComfyNode): is_api_node=True, price_badge=IO.PriceBadge( depends_on=IO.PriceBadgeDepends( - widgets=["model", "model.resolution", "model.number_of_images"], + widgets=["model", "model.resolution", "model.number_of_images", "model.quality"], ), expr=""" ( - $isQualityModel := widgets.model = "grok-imagine-image-quality"; + $is20 := widgets.model = "grok-imagine-image-2.0"; $isPro := $contains(widgets.model, "pro"); $res := $lookup(widgets, "model.resolution"); $n := $lookup(widgets, "model.number_of_images"); - $rate := $isQualityModel - ? ($res = "1k" ? 0.05 : 0.07) - : ($isPro ? 0.07 : 0.02); - $base := $isQualityModel ? 0.01 : 0.002; + $is1k := $res = "1k"; + $rate := $is20 + ? ($lookup(widgets, "model.quality") = "low" + ? ($is1k ? 0.04 : 0.06) + : ($is1k ? 0.06 : 0.08)) + : (widgets.model = "grok-imagine-image-quality" + ? ($is1k ? 0.05 : 0.07) + : ($isPro ? 0.07 : 0.02)); + $base := ($is20 or widgets.model = "grok-imagine-image-quality") ? 0.01 : 0.002; $output := $rate * $n; $isPro ? {"type":"usd","usd": $base + $output} @@ -525,13 +579,15 @@ class GrokImageEditNodeV2(IO.ComfyNode): image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None] n_images = sum(get_number_of_images(t) for t in image_tensors) + max_images = _GROK_IMAGE_EDIT_MAX_IMAGES.get(model_id, 3) if n_images < 1: raise ValueError("At least one image is required for editing.") - if model_id == "grok-imagine-image-pro" and n_images > 1: - raise ValueError("The pro model supports only 1 input image.") - if model_id != "grok-imagine-image-pro" and n_images > 3: - raise ValueError("A maximum of 3 input images is supported.") - if aspect_ratio != "auto" and n_images == 1: + if n_images > max_images: + raise ValueError( + f"The {model_id} model supports at most {max_images} input " + f"image{'s' if max_images > 1 else ''}; {n_images} are connected." + ) + if aspect_ratio != "auto" and model_id in _GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE and n_images == 1: raise ValueError( "Custom aspect ratio is only allowed when multiple images are connected to the image input." ) @@ -547,7 +603,7 @@ class GrokImageEditNodeV2(IO.ComfyNode): cls, ApiEndpoint(path="/proxy/xai/v1/images/edits", method="POST"), data=ImageEditRequest( - model=model_id, + model=_GROK_IMAGE_MODEL_API_IDS.get(model_id, model_id), images=[ InputUrlObject(url=f"data:image/png;base64,{tensor_to_base64_string(i)}") for i in flat_tensors ], @@ -556,6 +612,7 @@ class GrokImageEditNodeV2(IO.ComfyNode): n=number_of_images, seed=seed, aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio, + quality=model.get("quality") if model_id in _GROK_IMAGE_QUALITY_MODELS else None, ), response_model=ImageGenerationResponse, )