1 line
5.8 KiB
JSON
1 line
5.8 KiB
JSON
{"repo": "aryan185/comfyui-externalapi-helpers", "files": {"grok.py": {"1df85ba66dc674498b88d921ed360f5f56b70f67": {"t": 1785674551.3593502, "n": 5390, "h": "7af622488826e292", "c": "import os\nimport io\nimport base64\nimport requests\nimport torch\nimport numpy as np\nfrom PIL import Image\n\n\nXAI_GENERATIONS_URL = \"https://api.x.ai/v1/images/generations\"\nXAI_EDITS_URL = \"https://api.x.ai/v1/images/edits\"\n\n\ndef _tensor_to_base64(tensor) -> str:\n arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)\n buf = io.BytesIO()\n Image.fromarray(arr).save(buf, format=\"PNG\")\n return base64.b64encode(buf.getvalue()).decode(\"utf-8\")\n\n\ndef _image_obj(tensor) -> dict:\n frame = tensor[0] if tensor.dim() == 4 else tensor\n return {\"url\": f\"data:image/png;base64,{_tensor_to_base64(frame)}\", \"type\": \"image_url\"}\n\n\ndef _url_to_tensor(url: str) -> torch.Tensor:\n resp = requests.get(url, timeout=60)\n resp.raise_for_status()\n pil = Image.open(io.BytesIO(resp.content)).convert(\"RGB\")\n arr = np.array(pil).astype(np.float32) / 255.0\n return torch.from_numpy(arr).unsqueeze(0)\n\n\nclass GrokImageAPINode:\n\n @classmethod\n def INPUT_TYPES(cls):\n return {\n \"required\": {\n \"prompt\": (\"STRING\", {\"multiline\": True, \"default\": \"\"}),\n \"model\": (\n [\"grok-imagine-image\", \"grok-imagine-image-beta\"],\n {\"default\": \"grok-imagine-image\"},\n ),\n \"aspect_ratio\": (\n [\"auto\", \"1:1\", \"2:3\", \"3:2\", \"3:4\", \"4:3\", \"9:16\", \"16:9\",\n \"9:19.5\", \"19.5:9\", \"9:20\", \"20:9\", \"1:2\", \"2:1\"],\n {\"default\": \"1:1\",\n \"tooltip\": \"Applies to generation and multi-image edits. \"\n \"Ignored for single-image edits (output follows input aspect ratio).\"},\n ),\n \"resolution\": ([\"1k\", \"2k\"], {\"default\": \"1k\"}),\n \"n\": (\"INT\", {\"default\": 1, \"min\": 1, \"max\": 10, \"step\": 1,\n \"tooltip\": \"Number of images to generate or edited variations to produce\"}),\n \"seed\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 2147483647, \"step\": 1,\n \"tooltip\": \"Used to trigger re-runs in ComfyUI; not sent to the API\"}),\n \"api_key\": (\"STRING\", {\"multiline\": False, \"default\": \"\",\n \"tooltip\": \"Directly put xAI API key or .env variable name (XAI_API_KEY)\"}),\n },\n \"optional\": {\n \"image_1\": (\"IMAGE\",),\n \"image_2\": (\"IMAGE\",),\n \"image_3\": (\"IMAGE\",),\n \"image_4\": (\"IMAGE\",),\n \"image_5\": (\"IMAGE\",),\n },\n }\n\n RETURN_TYPES = (\"IMAGE\",)\n RETURN_NAMES = (\"image\",)\n FUNCTION = \"edit_image\"\n CATEGORY = \"image/generation\"\n\n def edit_image(self, prompt, model, aspect_ratio, resolution, n, seed, api_key,\n image_1=None, image_2=None, image_3=None, image_4=None, image_5=None):\n if not prompt.strip():\n raise ValueError(\"Grok: prompt cannot be empty.\")\n\n key = os.environ.get(api_key.strip(), api_key.strip()) or os.environ.get(\"XAI_API_KEY\")\n if not key:\n raise ValueError(\"Grok: no API key provided.\")\n\n input_images = [img for img in [image_1, image_2, image_3, image_4, image_5] if img is not None]\n num_images = len(input_images)\n\n if num_images == 0:\n endpoint = XAI_GENERATIONS_URL\n payload = {\n \"model\": model,\n \"prompt\": prompt.strip(),\n \"aspect_ratio\": aspect_ratio,\n \"resolution\": resolution,\n \"n\": n,\n }\n else:\n endpoint = XAI_EDITS_URL\n single = num_images == 1\n payload = {\n \"model\": model,\n \"image\" if single else \"images\": _image_obj(input_images[0]) if single else [_image_obj(img) for img in input_images],\n \"prompt\": prompt.strip(),\n \"resolution\": resolution,\n \"n\": n,\n }\n if single:\n print(\"Grok: aspect_ratio ignored for single-image edits (output follows input aspect ratio).\")\n else:\n payload[\"aspect_ratio\"] = aspect_ratio\n\n headers = {\"Authorization\": f\"Bearer {key}\", \"Content-Type\": \"application/json\"}\n\n resp = requests.post(endpoint, json=payload, headers=headers, timeout=120)\n if resp.status_code != 200:\n raise RuntimeError(f\"Grok API error {resp.status_code} – {resp.text}\")\n\n data = resp.json().get(\"data\", [])\n if not data:\n raise RuntimeError(\"Grok: API returned no image data.\")\n\n tensors = []\n for item in data:\n if item.get(\"url\"):\n tensors.append(_url_to_tensor(item[\"url\"]))\n elif item.get(\"b64_json\"):\n pil = Image.open(io.BytesIO(base64.b64decode(item[\"b64_json\"]))).convert(\"RGB\")\n tensors.append(torch.from_numpy(np.array(pil).astype(np.float32) / 255.0).unsqueeze(0))\n if not tensors:\n raise RuntimeError(\"Grok: could not retrieve any output images.\")\n\n return (torch.cat(tensors, dim=0),)\n\n @classmethod\n def IS_CHANGED(cls, prompt, seed, **kwargs):\n return f\"{prompt}-{seed}\"\n\n\nNODE_CLASS_MAPPINGS = {\"GrokImageAPINode\": GrokImageAPINode}\nNODE_DISPLAY_NAME_MAPPINGS = {\"GrokImageAPINode\": \"Grok Image API\"}"}}}} |