GrokImageEditNode: add optional aspect_ratio parameter

Author: bigcat88

comfy_api_nodes/apis/grok.py

@@ -23,6 +23,7 @@ class ImageEditRequest(BaseModel):
     n: int = Field(...)
     seed: int = Field(...)
     response_format: str = Field("url")
+    aspect_ratio: str | None = Field(...)
 
 
 class VideoGenerationRequest(BaseModel):

comfy_api_nodes/nodes_grok.py

@@ -187,6 +187,27 @@ def define_schema(cls):
                     tooltip="Seed to determine if node should re-run; "
                     "actual results are nondeterministic regardless of seed.",
                 ),
+                IO.Combo.Input(
+                    "aspect_ratio",
+                    options=[
+                        "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",
+                    ],
+                    optional=True,
+                    tooltip="Only allowed when multiple images are connected to the image input.",
+                ),
             ],
             outputs=[
                 IO.Image.Output(),
@@ -217,13 +238,18 @@ async def execute(
         resolution: str,
         number_of_images: int,
         seed: int,
+        aspect_ratio: str = "auto",
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         if model == "grok-imagine-image-pro":
             if get_number_of_images(image) > 1:
                 raise ValueError("The pro model supports only 1 input image.")
         elif get_number_of_images(image) > 3:
             raise ValueError("A maximum of 3 input images is supported.")
+        if aspect_ratio != "auto" and get_number_of_images(image) == 1:
+            raise ValueError(
+                "Custom aspect ratio is only allowed when multiple images are connected to the image input."
+            )
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/xai/v1/images/edits", method="POST"),
@@ -234,6 +260,7 @@ async def execute(
                 resolution=resolution.lower(),
                 n=number_of_images,
                 seed=seed,
+                aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
             ),
             response_model=ImageGenerationResponse,
             price_extractor=_extract_grok_price,