RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[Partner Nodes] feat(Grok): add Grok Imagine Image 2.0 model (#15496) 

Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
This commit is contained in:
Alexander Piskun 2026-08-11 22:38:08 +03:00 committed by GitHub
parent eb4a7b4fcf
commit 2eaf09f50d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 78 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy_api_nodes/apis/grok.py
View File

@ -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):

95
comfy_api_nodes/nodes_grok.py
View File

@ -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,
)