From d6baab76b7b82c10a12d501ca9af1575bb022f5b Mon Sep 17 00:00:00 2001 From: Aryan185 Date: Fri, 20 Mar 2026 09:55:11 +0000 Subject: [PATCH] added grok image node --- .env.example | 3 +- __init__.py | 5 +- grok.py | 138 +++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 4 files changed, 144 insertions(+), 4 deletions(-) create mode 100644 grok.py diff --git a/.env.example b/.env.example index fced9e9..b9853a2 100644 --- a/.env.example +++ b/.env.example @@ -5,4 +5,5 @@ GEMINI_API_KEY="your_gemini_api_key" OPENAI_API_KEY="your_openai_api_key" XI_API_KEY="your_elevenlabs_api_key" REPLICATE_API_TOKEN="your_replicate_api_token" -TRIPO_API_KEY="your_tripo_api_key" \ No newline at end of file +TRIPO_API_KEY="your_tripo_api_key" +XAI_API_KEY="your_xai_api_key" \ No newline at end of file diff --git a/__init__.py b/__init__.py index a125dd5..59b81d3 100644 --- a/__init__.py +++ b/__init__.py @@ -20,7 +20,8 @@ from .sora import NODE_CLASS_MAPPINGS as SORA_MAPPINGS, NODE_DISPLAY_NAME_MAPPIN from .cleanup import NODE_CLASS_MAPPINGS as CLEANUP_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CLEANUP_DISPLAY from .tripoTextToModel import NODE_CLASS_MAPPINGS as TRIPO_TEXT_TO_MODEL_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as TRIPO_TEXT_TO_MODEL_DISPLAY from .tripoImageToModel import NODE_CLASS_MAPPINGS as TRIPO_IMAGE_TO_MODEL_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as TRIPO_IMAGE_TO_MODEL_DISPLAY +from .grok import NODE_CLASS_MAPPINGS as GROK_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GROK_DISPLAY -NODE_CLASS_MAPPINGS = {**PRO_MAPPINGS, **GEMINI_MAPPINGS, **GEMINI_DIAR_MAPPINGS, **GPT_IMAGE_MAPPINGS, **IMAGEN_IMAGE_MAPPINGS, **IMAGEN_EDIT_MAPPINGS, **VEO_MAPPINGS, **VEO_GEMINI_MAPPINGS, **GEMINI_SEGMENT_MAPPINGS, **NANO_BANANA_MAPPINGS, **GEMINI_TTS_MAPPINGS, **ELEVENLABS_TTS_MAPPINGS, **FLUX2_MAPPINGS, **OPENAI_TTS_MAPPINGS, **OPENAI_LLM_MAPPINGS, **SORA_MAPPINGS, **CLEANUP_MAPPINGS, **TRIPO_TEXT_TO_MODEL_MAPPINGS, **TRIPO_IMAGE_TO_MODEL_MAPPINGS} -NODE_DISPLAY_NAME_MAPPINGS = {**PRO_DISPLAY, **GEMINI_DISPLAY, **GEMINI_DIAR_DISPLAY, **GPT_IMAGE_DISPLAY, **IMAGEN_IMAGE_DISPLAY, **IMAGEN_EDIT_DISPLAY, **VEO_DISPLAY, **VEO_GEMINI_DISPLAY, **GEMINI_SEGMENT_DISPLAY, **NANO_BANANA_DISPLAY, **GEMINI_TTS_DISPLAY, **ELEVENLABS_TTS_DISPLAY, **FLUX2_DISPLAY, **OPENAI_TTS_DISPLAY, **OPENAI_LLM_DISPLAY, **SORA_DISPLAY, **CLEANUP_DISPLAY, **TRIPO_TEXT_TO_MODEL_DISPLAY, **TRIPO_IMAGE_TO_MODEL_DISPLAY} +NODE_CLASS_MAPPINGS = {**PRO_MAPPINGS, **GEMINI_MAPPINGS, **GEMINI_DIAR_MAPPINGS, **GPT_IMAGE_MAPPINGS, **IMAGEN_IMAGE_MAPPINGS, **IMAGEN_EDIT_MAPPINGS, **VEO_MAPPINGS, **VEO_GEMINI_MAPPINGS, **GEMINI_SEGMENT_MAPPINGS, **NANO_BANANA_MAPPINGS, **GEMINI_TTS_MAPPINGS, **ELEVENLABS_TTS_MAPPINGS, **FLUX2_MAPPINGS, **OPENAI_TTS_MAPPINGS, **OPENAI_LLM_MAPPINGS, **SORA_MAPPINGS, **CLEANUP_MAPPINGS, **TRIPO_TEXT_TO_MODEL_MAPPINGS, **TRIPO_IMAGE_TO_MODEL_MAPPINGS, **GROK_MAPPINGS} +NODE_DISPLAY_NAME_MAPPINGS = {**PRO_DISPLAY, **GEMINI_DISPLAY, **GEMINI_DIAR_DISPLAY, **GPT_IMAGE_DISPLAY, **IMAGEN_IMAGE_DISPLAY, **IMAGEN_EDIT_DISPLAY, **VEO_DISPLAY, **VEO_GEMINI_DISPLAY, **GEMINI_SEGMENT_DISPLAY, **NANO_BANANA_DISPLAY, **GEMINI_TTS_DISPLAY, **ELEVENLABS_TTS_DISPLAY, **FLUX2_DISPLAY, **OPENAI_TTS_DISPLAY, **OPENAI_LLM_DISPLAY, **SORA_DISPLAY, **CLEANUP_DISPLAY, **TRIPO_TEXT_TO_MODEL_DISPLAY, **TRIPO_IMAGE_TO_MODEL_DISPLAY, **GROK_DISPLAY} __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/grok.py b/grok.py new file mode 100644 index 0000000..12b4b3d --- /dev/null +++ b/grok.py @@ -0,0 +1,138 @@ +import os +import io +import base64 +import requests +import torch +import numpy as np +from PIL import Image + + +XAI_GENERATIONS_URL = "https://api.x.ai/v1/images/generations" +XAI_EDITS_URL = "https://api.x.ai/v1/images/edits" + + +def _tensor_to_base64(tensor) -> str: + arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8) + buf = io.BytesIO() + Image.fromarray(arr).save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("utf-8") + + +def _image_obj(tensor) -> dict: + frame = tensor[0] if tensor.dim() == 4 else tensor + return {"url": f"data:image/png;base64,{_tensor_to_base64(frame)}", "type": "image_url"} + + +def _url_to_tensor(url: str) -> torch.Tensor: + resp = requests.get(url, timeout=60) + resp.raise_for_status() + pil = Image.open(io.BytesIO(resp.content)).convert("RGB") + arr = np.array(pil).astype(np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) + + +class GrokImageAPINode: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": ""}), + "model": ( + ["grok-imagine-image", "grok-imagine-image-beta"], + {"default": "grok-imagine-image"}, + ), + "aspect_ratio": ( + ["auto", "1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", + "9:19.5", "19.5:9", "9:20", "20:9", "1:2", "2:1"], + {"default": "1:1", + "tooltip": "Applies to generation and multi-image edits. " + "Ignored for single-image edits (output follows input aspect ratio)."}, + ), + "resolution": (["1k", "2k"], {"default": "1k"}), + "n": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1, + "tooltip": "Number of images to generate or edited variations to produce"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 2147483647, "step": 1, + "tooltip": "Used to trigger re-runs in ComfyUI; not sent to the API"}), + "api_key": ("STRING", {"multiline": False, "default": "", + "tooltip": "Directly put xAI API key or .env variable name (XAI_API_KEY)"}), + }, + "optional": { + "image_1": ("IMAGE",), + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "image_4": ("IMAGE",), + "image_5": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "edit_image" + CATEGORY = "image/generation" + + def edit_image(self, prompt, model, aspect_ratio, resolution, n, seed, api_key, + image_1=None, image_2=None, image_3=None, image_4=None, image_5=None): + if not prompt.strip(): + raise ValueError("Grok: prompt cannot be empty.") + + key = os.environ.get(api_key.strip(), api_key.strip()) or os.environ.get("XAI_API_KEY") + if not key: + raise ValueError("Grok: no API key provided.") + + input_images = [img for img in [image_1, image_2, image_3, image_4, image_5] if img is not None] + num_images = len(input_images) + + if num_images == 0: + endpoint = XAI_GENERATIONS_URL + payload = { + "model": model, + "prompt": prompt.strip(), + "aspect_ratio": aspect_ratio, + "resolution": resolution, + "n": n, + } + else: + endpoint = XAI_EDITS_URL + single = num_images == 1 + payload = { + "model": model, + "image" if single else "images": _image_obj(input_images[0]) if single else [_image_obj(img) for img in input_images], + "prompt": prompt.strip(), + "resolution": resolution, + "n": n, + } + if single: + print("Grok: aspect_ratio ignored for single-image edits (output follows input aspect ratio).") + else: + payload["aspect_ratio"] = aspect_ratio + + headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"} + + resp = requests.post(endpoint, json=payload, headers=headers, timeout=120) + if resp.status_code != 200: + raise RuntimeError(f"Grok API error {resp.status_code} – {resp.text}") + + data = resp.json().get("data", []) + if not data: + raise RuntimeError("Grok: API returned no image data.") + + tensors = [] + for item in data: + if item.get("url"): + tensors.append(_url_to_tensor(item["url"])) + elif item.get("b64_json"): + pil = Image.open(io.BytesIO(base64.b64decode(item["b64_json"]))).convert("RGB") + tensors.append(torch.from_numpy(np.array(pil).astype(np.float32) / 255.0).unsqueeze(0)) + if not tensors: + raise RuntimeError("Grok: could not retrieve any output images.") + + return (torch.cat(tensors, dim=0),) + + @classmethod + def IS_CHANGED(cls, prompt, seed, **kwargs): + return f"{prompt}-{seed}" + + +NODE_CLASS_MAPPINGS = {"GrokImageAPINode": GrokImageAPINode} +NODE_DISPLAY_NAME_MAPPINGS = {"GrokImageAPINode": "Grok Image API"} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 70517cd..25fc4df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "externalapi-helpers" description = "Various ComfyUI nodes for Gemini, Replicate and OpenAI" -version = "1.0.9" +version = "1.0.10" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems)