Skip to content

Commit 3cdc0d5

Browse files
authored
[Partner Nodes] GPTImage: fix price badges, add new resolutions (#13519)
* fix(api-nodes): fixed price badges, add new resolutions Signed-off-by: bigcat88 <bigcat88@icloud.com> * proper calculate the total run cost when "n > 1" Signed-off-by: bigcat88 <bigcat88@icloud.com> --------- Signed-off-by: bigcat88 <bigcat88@icloud.com>
1 parent 749d5b4 commit 3cdc0d5

1 file changed

Lines changed: 46 additions & 13 deletions

File tree

comfy_api_nodes/nodes_openai.py

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,10 @@ def calculate_tokens_price_image_1_5(response: OpenAIImageGenerationResponse) ->
357357
return ((response.usage.input_tokens * 8.0) + (response.usage.output_tokens * 32.0)) / 1_000_000.0
358358

359359

360+
def calculate_tokens_price_image_2_0(response: OpenAIImageGenerationResponse) -> float | None:
361+
return ((response.usage.input_tokens * 8.0) + (response.usage.output_tokens * 30.0)) / 1_000_000.0
362+
363+
360364
class OpenAIGPTImage1(IO.ComfyNode):
361365

362366
@classmethod
@@ -401,7 +405,17 @@ def define_schema(cls):
401405
IO.Combo.Input(
402406
"size",
403407
default="auto",
404-
options=["auto", "1024x1024", "1024x1536", "1536x1024"],
408+
options=[
409+
"auto",
410+
"1024x1024",
411+
"1024x1536",
412+
"1536x1024",
413+
"2048x2048",
414+
"2048x1152",
415+
"1152x2048",
416+
"3840x2160",
417+
"2160x3840",
418+
],
405419
tooltip="Image size",
406420
optional=True,
407421
),
@@ -427,7 +441,7 @@ def define_schema(cls):
427441
),
428442
IO.Combo.Input(
429443
"model",
430-
options=["gpt-image-1", "gpt-image-1.5", 'gpt-image-2'],
444+
options=["gpt-image-1", "gpt-image-1.5", "gpt-image-2"],
431445
default="gpt-image-2",
432446
optional=True,
433447
),
@@ -442,23 +456,36 @@ def define_schema(cls):
442456
],
443457
is_api_node=True,
444458
price_badge=IO.PriceBadge(
445-
depends_on=IO.PriceBadgeDepends(widgets=["quality", "n"]),
459+
depends_on=IO.PriceBadgeDepends(widgets=["quality", "n", "model"]),
446460
expr="""
447461
(
448462
$ranges := {
449-
"low": [0.011, 0.02],
450-
"medium": [0.046, 0.07],
451-
"high": [0.167, 0.3]
463+
"gpt-image-1": {
464+
"low": [0.011, 0.02],
465+
"medium": [0.042, 0.07],
466+
"high": [0.167, 0.25]
467+
},
468+
"gpt-image-1.5": {
469+
"low": [0.009, 0.02],
470+
"medium": [0.034, 0.062],
471+
"high": [0.133, 0.22]
472+
},
473+
"gpt-image-2": {
474+
"low": [0.0048, 0.012],
475+
"medium": [0.041, 0.112],
476+
"high": [0.165, 0.43]
477+
}
452478
};
453-
$range := $lookup($ranges, widgets.quality);
454-
$n := widgets.n;
479+
$range := $lookup($lookup($ranges, widgets.model), widgets.quality);
480+
$nRaw := widgets.n;
481+
$n := ($nRaw != null and $nRaw != 0) ? $nRaw : 1;
455482
($n = 1)
456-
? {"type":"range_usd","min_usd": $range[0], "max_usd": $range[1]}
483+
? {"type":"range_usd","min_usd": $range[0], "max_usd": $range[1], "format": {"approximate": true}}
457484
: {
458485
"type":"range_usd",
459-
"min_usd": $range[0],
460-
"max_usd": $range[1],
461-
"format": { "suffix": " x " & $string($n) & "/Run" }
486+
"min_usd": $range[0] * $n,
487+
"max_usd": $range[1] * $n,
488+
"format": { "suffix": "/Run", "approximate": true }
462489
}
463490
)
464491
""",
@@ -483,12 +510,18 @@ async def execute(
483510
if mask is not None and image is None:
484511
raise ValueError("Cannot use a mask without an input image")
485512

513+
if model in ("gpt-image-1", "gpt-image-1.5"):
514+
if size not in ("auto", "1024x1024", "1024x1536", "1536x1024"):
515+
raise ValueError(f"Resolution {size} is only supported by GPT Image 2 model")
516+
486517
if model == "gpt-image-1":
487518
price_extractor = calculate_tokens_price_image_1
488519
elif model == "gpt-image-1.5":
489520
price_extractor = calculate_tokens_price_image_1_5
490521
elif model == "gpt-image-2":
491-
price_extractor = calculate_tokens_price_image_1_5
522+
price_extractor = calculate_tokens_price_image_2_0
523+
if background == "transparent":
524+
raise ValueError("Transparent background is not supported for GPT Image 2 model")
492525
else:
493526
raise ValueError(f"Unknown model: {model}")
494527

0 commit comments

Comments
 (0)