1- from io import BytesIO
1+ import base64
22import os
33from enum import Enum
4- from inspect import cleandoc
4+ from io import BytesIO
5+
56import numpy as np
67import torch
78from PIL import Image
8- import folder_paths
9- import base64
10- from comfy_api .latest import IO , ComfyExtension
119from typing_extensions import override
1210
13-
11+ import folder_paths
12+ from comfy_api .latest import IO , ComfyExtension , Input
1413from comfy_api_nodes .apis import (
15- OpenAIImageGenerationRequest ,
16- OpenAIImageEditRequest ,
17- OpenAIImageGenerationResponse ,
18- OpenAICreateResponse ,
19- OpenAIResponse ,
2014 CreateModelResponseProperties ,
21- Item ,
22- OutputContent ,
23- InputImageContent ,
2415 Detail ,
25- InputTextContent ,
26- InputMessage ,
27- InputMessageContentList ,
2816 InputContent ,
2917 InputFileContent ,
18+ InputImageContent ,
19+ InputMessage ,
20+ InputMessageContentList ,
21+ InputTextContent ,
22+ Item ,
23+ OpenAICreateResponse ,
24+ OpenAIResponse ,
25+ OutputContent ,
26+ )
27+ from comfy_api_nodes .apis .openai_api import (
28+ OpenAIImageEditRequest ,
29+ OpenAIImageGenerationRequest ,
30+ OpenAIImageGenerationResponse ,
3031)
31-
3232from comfy_api_nodes .util import (
33- downscale_image_tensor ,
34- download_url_to_bytesio ,
35- validate_string ,
36- tensor_to_base64_string ,
3733 ApiEndpoint ,
38- sync_op ,
34+ download_url_to_bytesio ,
35+ downscale_image_tensor ,
3936 poll_op ,
37+ sync_op ,
38+ tensor_to_base64_string ,
4039 text_filepath_to_data_uri ,
40+ validate_string ,
4141)
4242
43-
4443RESPONSES_ENDPOINT = "/proxy/openai/v1/responses"
4544STARTING_POINT_ID_PATTERN = r"<starting_point_id:(.*)>"
4645
@@ -98,17 +97,14 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten
9897
9998
10099class OpenAIDalle2 (IO .ComfyNode ):
101- """
102- Generates images synchronously via OpenAI's DALL·E 2 endpoint.
103- """
104100
105101 @classmethod
106102 def define_schema (cls ):
107103 return IO .Schema (
108104 node_id = "OpenAIDalle2" ,
109105 display_name = "OpenAI DALL·E 2" ,
110106 category = "api node/image/OpenAI" ,
111- description = cleandoc ( cls . __doc__ or "" ) ,
107+ description = "Generates images synchronously via OpenAI's DALL·E 2 endpoint." ,
112108 inputs = [
113109 IO .String .Input (
114110 "prompt" ,
@@ -234,17 +230,14 @@ async def execute(
234230
235231
236232class OpenAIDalle3 (IO .ComfyNode ):
237- """
238- Generates images synchronously via OpenAI's DALL·E 3 endpoint.
239- """
240233
241234 @classmethod
242235 def define_schema (cls ):
243236 return IO .Schema (
244237 node_id = "OpenAIDalle3" ,
245238 display_name = "OpenAI DALL·E 3" ,
246239 category = "api node/image/OpenAI" ,
247- description = cleandoc ( cls . __doc__ or "" ) ,
240+ description = "Generates images synchronously via OpenAI's DALL·E 3 endpoint." ,
248241 inputs = [
249242 IO .String .Input (
250243 "prompt" ,
@@ -327,23 +320,20 @@ async def execute(
327320
328321
329322class OpenAIGPTImage1 (IO .ComfyNode ):
330- """
331- Generates images synchronously via OpenAI's GPT Image 1 endpoint.
332- """
333323
334324 @classmethod
335325 def define_schema (cls ):
336326 return IO .Schema (
337327 node_id = "OpenAIGPTImage1" ,
338328 display_name = "OpenAI GPT Image 1" ,
339329 category = "api node/image/OpenAI" ,
340- description = cleandoc ( cls . __doc__ or "" ) ,
330+ description = "Generates images synchronously via OpenAI's GPT Image 1 endpoint." ,
341331 inputs = [
342332 IO .String .Input (
343333 "prompt" ,
344334 default = "" ,
345335 multiline = True ,
346- tooltip = "Text prompt for GPT Image 1 " ,
336+ tooltip = "Text prompt for GPT Image" ,
347337 ),
348338 IO .Int .Input (
349339 "seed" ,
@@ -365,8 +355,8 @@ def define_schema(cls):
365355 ),
366356 IO .Combo .Input (
367357 "background" ,
368- default = "opaque " ,
369- options = ["opaque" , "transparent" ],
358+ default = "auto " ,
359+ options = ["auto" , " opaque" , "transparent" ],
370360 tooltip = "Return image with or without background" ,
371361 optional = True ,
372362 ),
@@ -397,6 +387,11 @@ def define_schema(cls):
397387 tooltip = "Optional mask for inpainting (white areas will be replaced)" ,
398388 optional = True ,
399389 ),
390+ IO .Combo .Input (
391+ "model" ,
392+ options = ["gpt-image-1" , "gpt-image-1.5" ],
393+ optional = True ,
394+ ),
400395 ],
401396 outputs = [
402397 IO .Image .Output (),
@@ -412,32 +407,27 @@ def define_schema(cls):
412407 @classmethod
413408 async def execute (
414409 cls ,
415- prompt ,
416- seed = 0 ,
417- quality = "low" ,
418- background = "opaque" ,
419- image = None ,
420- mask = None ,
421- n = 1 ,
422- size = "1024x1024" ,
410+ prompt : str ,
411+ seed : int = 0 ,
412+ quality : str = "low" ,
413+ background : str = "opaque" ,
414+ image : Input .Image | None = None ,
415+ mask : Input .Image | None = None ,
416+ n : int = 1 ,
417+ size : str = "1024x1024" ,
418+ model : str = "gpt-image-1" ,
423419 ) -> IO .NodeOutput :
424420 validate_string (prompt , strip_whitespace = False )
425- model = "gpt-image-1"
426- path = "/proxy/openai/images/generations"
427- content_type = "application/json"
428- request_class = OpenAIImageGenerationRequest
429- files = []
430421
431- if image is not None :
432- path = "/proxy/openai/images/edits"
433- request_class = OpenAIImageEditRequest
434- content_type = "multipart/form-data"
422+ if mask is not None and image is None :
423+ raise ValueError ("Cannot use a mask without an input image" )
435424
425+ if image is not None :
426+ files = []
436427 batch_size = image .shape [0 ]
437-
438428 for i in range (batch_size ):
439- single_image = image [i : i + 1 ]
440- scaled_image = downscale_image_tensor (single_image ).squeeze ()
429+ single_image = image [i : i + 1 ]
430+ scaled_image = downscale_image_tensor (single_image , total_pixels = 2048 * 2048 ).squeeze ()
441431
442432 image_np = (scaled_image .numpy () * 255 ).astype (np .uint8 )
443433 img = Image .fromarray (image_np )
@@ -450,44 +440,56 @@ async def execute(
450440 else :
451441 files .append (("image[]" , (f"image_{ i } .png" , img_byte_arr , "image/png" )))
452442
453- if mask is not None :
454- if image is None :
455- raise Exception ("Cannot use a mask without an input image" )
456- if image .shape [0 ] != 1 :
457- raise Exception ("Cannot use a mask with multiple image" )
458- if mask .shape [1 :] != image .shape [1 :- 1 ]:
459- raise Exception ("Mask and Image must be the same size" )
460- batch , height , width = mask .shape
461- rgba_mask = torch .zeros (height , width , 4 , device = "cpu" )
462- rgba_mask [:, :, 3 ] = 1 - mask .squeeze ().cpu ()
443+ if mask is not None :
444+ if mask .shape [1 :] != image .shape [1 :- 1 ]:
445+ raise Exception ("Mask and Image must be the same size" )
446+ _ , height , width = mask .shape
447+ # Build RGBA PNG mask where ALPHA=0 means "edit here" (transparent), ALPHA=1 means "keep" (opaque).
448+ rgba_mask = torch .zeros (height , width , 4 , device = "cpu" )
449+ rgba_mask [:, :, 3 ] = 1 - mask .squeeze ().cpu ()
463450
464- scaled_mask = downscale_image_tensor (rgba_mask .unsqueeze (0 )).squeeze ()
451+ scaled_mask = downscale_image_tensor (rgba_mask .unsqueeze (0 ), total_pixels = 2048 * 2048 ).squeeze ()
465452
466- mask_np = (scaled_mask .numpy () * 255 ).astype (np .uint8 )
467- mask_img = Image .fromarray (mask_np )
468- mask_img_byte_arr = BytesIO ()
469- mask_img .save (mask_img_byte_arr , format = "PNG" )
470- mask_img_byte_arr .seek (0 )
471- files .append (("mask" , ("mask.png" , mask_img_byte_arr , "image/png" )))
472-
473- # Build the operation
474- response = await sync_op (
475- cls ,
476- ApiEndpoint (path = path , method = "POST" ),
477- response_model = OpenAIImageGenerationResponse ,
478- data = request_class (
479- model = model ,
480- prompt = prompt ,
481- quality = quality ,
482- background = background ,
483- n = n ,
484- seed = seed ,
485- size = size ,
486- ),
487- files = files if files else None ,
488- content_type = content_type ,
489- )
453+ mask_np = (scaled_mask .numpy () * 255 ).astype (np .uint8 )
454+ mask_img = Image .fromarray (mask_np )
455+ mask_img_byte_arr = BytesIO ()
456+ mask_img .save (mask_img_byte_arr , format = "PNG" )
457+ mask_img_byte_arr .seek (0 )
458+ files .append (("mask" , ("mask.png" , mask_img_byte_arr , "image/png" )))
490459
460+ response = await sync_op (
461+ cls ,
462+ ApiEndpoint (path = "/proxy/openai/images/edits" , method = "POST" ),
463+ response_model = OpenAIImageGenerationResponse ,
464+ data = OpenAIImageEditRequest (
465+ model = model ,
466+ prompt = prompt ,
467+ quality = quality ,
468+ background = background ,
469+ n = n ,
470+ seed = seed ,
471+ size = size ,
472+ moderation = "low" ,
473+ ),
474+ content_type = "multipart/form-data" ,
475+ files = files ,
476+ )
477+ else :
478+ response = await sync_op (
479+ cls ,
480+ ApiEndpoint (path = "/proxy/openai/images/generations" , method = "POST" ),
481+ response_model = OpenAIImageGenerationResponse ,
482+ data = OpenAIImageGenerationRequest (
483+ model = model ,
484+ prompt = prompt ,
485+ quality = quality ,
486+ background = background ,
487+ n = n ,
488+ seed = seed ,
489+ size = size ,
490+ moderation = "low" ,
491+ ),
492+ )
491493 return IO .NodeOutput (await validate_and_cast_response (response ))
492494
493495
0 commit comments