Skip to content

Commit 0695702

Browse files
authored
fix(api-nodes): add support for "thought_image" in Nano Banana 2 and corrected price badges (Comfy-Org#13038)
1 parent b941913 commit 0695702

File tree

2 files changed

+15
-3
lines changed

2 files changed

+15
-3
lines changed

comfy_api_nodes/apis/gemini.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ class GeminiPart(BaseModel):
6767
inlineData: GeminiInlineData | None = Field(None)
6868
fileData: GeminiFileData | None = Field(None)
6969
text: str | None = Field(None)
70+
thought: bool | None = Field(None)
7071

7172

7273
class GeminiTextPart(BaseModel):

comfy_api_nodes/nodes_gemini.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
$m := widgets.model;
6464
$r := widgets.resolution;
6565
$isFlash := $contains($m, "nano banana 2");
66-
$flashPrices := {"1k": 0.0696, "2k": 0.0696, "4k": 0.123};
66+
$flashPrices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
6767
$proPrices := {"1k": 0.134, "2k": 0.134, "4k": 0.24};
6868
$prices := $isFlash ? $flashPrices : $proPrices;
6969
{"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
@@ -188,10 +188,12 @@ def get_text_from_response(response: GeminiGenerateContentResponse) -> str:
188188
return "\n".join([part.text for part in parts])
189189

190190

191-
async def get_image_from_response(response: GeminiGenerateContentResponse) -> Input.Image:
191+
async def get_image_from_response(response: GeminiGenerateContentResponse, thought: bool = False) -> Input.Image:
192192
image_tensors: list[Input.Image] = []
193193
parts = get_parts_by_type(response, "image/*")
194194
for part in parts:
195+
if (part.thought is True) != thought:
196+
continue
195197
if part.inlineData:
196198
image_data = base64.b64decode(part.inlineData.data)
197199
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
@@ -931,6 +933,11 @@ def define_schema(cls):
931933
outputs=[
932934
IO.Image.Output(),
933935
IO.String.Output(),
936+
IO.Image.Output(
937+
display_name="thought_image",
938+
tooltip="First image from the model's thinking process. "
939+
"Only available with thinking_level HIGH and IMAGE+TEXT modality.",
940+
),
934941
],
935942
hidden=[
936943
IO.Hidden.auth_token_comfy_org,
@@ -992,7 +999,11 @@ async def execute(
992999
response_model=GeminiGenerateContentResponse,
9931000
price_extractor=calculate_tokens_price,
9941001
)
995-
return IO.NodeOutput(await get_image_from_response(response), get_text_from_response(response))
1002+
return IO.NodeOutput(
1003+
await get_image_from_response(response),
1004+
get_text_from_response(response),
1005+
await get_image_from_response(response, thought=True),
1006+
)
9961007

9971008

9981009
class GeminiExtension(ComfyExtension):

0 commit comments

Comments
 (0)