separate out nodes per openai model

comfy_api_nodes/apis/PixverseController.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
-#   filename:  http://localhost:8080/openapi
+#   filename:  https://api.comfy.org/openapi
-#   timestamp: 2025-04-22T20:42:39+00:00
+#   timestamp: 2025-04-23T15:56:33+00:00
 
 from __future__ import annotations
 

comfy_api_nodes/apis/PixverseDto.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
-#   filename:  http://localhost:8080/openapi
+#   filename:  https://api.comfy.org/openapi
-#   timestamp: 2025-04-22T20:42:39+00:00
+#   timestamp: 2025-04-23T15:56:33+00:00
 
 from __future__ import annotations
 

comfy_api_nodes/apis/__init__.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
-#   filename:  http://localhost:8080/openapi
+#   filename:  https://api.comfy.org/openapi
-#   timestamp: 2025-04-22T20:42:39+00:00
+#   timestamp: 2025-04-23T15:56:33+00:00
 
 from __future__ import annotations
 
@@ -874,9 +874,8 @@ class GenerationType3(str, Enum):
 
 
 class LumaVideoModel(str, Enum):
-    ray_1_6 = 'ray-1-6'
     ray_2 = 'ray-2'
-    ray_flash_2 = 'ray-flash-2'
+    ray_2_flash = 'ray-2-flash'
 
 
 class LumaVideoModelOutputDuration1(str, Enum):
@@ -1058,34 +1057,6 @@ class NodeVersionUpdateRequest(BaseModel):
     )
 
 
-class Quality(str, Enum):
-    low = 'low'
-    medium = 'medium'
-    high = 'high'
-
-
-class OpenAIImageEditRequest(BaseModel):
-    image: Union[bytes_aliased, List[bytes_aliased]] = Field(
-        ...,
-        description='Image(s) to edit. For DALL-E 2, only a single image is supported. For gpt-image-1, multiple images can be provided using image[] notation in form data.',
-    )
-    mask: Optional[bytes_aliased] = Field(
-        None,
-        description='An additional image whose fully transparent areas (e.g. where alpha is zero) indicate where image should be edited',
-    )
-    model: str = Field(
-        ..., description='The model to use for image editing', examples=['gpt-image-1']
-    )
-    prompt: str = Field(
-        ...,
-        description='A text description of the desired edit',
-        examples=['Give the rocketship rainbow coloring'],
-    )
-    quality: Optional[Quality] = Field(
-        None, description='The quality of the edited image', examples=['low']
-    )
-
-
 class Background(str, Enum):
     transparent = 'transparent'
     opaque = 'opaque'
@@ -1102,6 +1073,57 @@ class OutputFormat(str, Enum):
     jpeg = 'jpeg'
 
 
+class Quality(str, Enum):
+    low = 'low'
+    medium = 'medium'
+    high = 'high'
+
+
+class OpenAIImageEditRequest(BaseModel):
+    background: Optional[Background] = Field(
+        None, description='Background transparency', examples=['opaque']
+    )
+    model: str = Field(
+        ..., description='The model to use for image editing', examples=['gpt-image-1']
+    )
+    moderation: Optional[Moderation] = Field(
+        None, description='Content moderation setting', examples=['auto']
+    )
+    n: Optional[int] = Field(
+        None, description='The number of images to generate', examples=[1]
+    )
+    output_compression: Optional[int] = Field(
+        None, description='Compression level for JPEG or WebP (0-100)', examples=[100]
+    )
+    output_format: Optional[OutputFormat] = Field(
+        None, description='Format of the output image', examples=['png']
+    )
+    prompt: str = Field(
+        ...,
+        description='A text description of the desired edit',
+        examples=['Give the rocketship rainbow coloring'],
+    )
+    quality: Optional[Quality] = Field(
+        None, description='The quality of the edited image', examples=['low']
+    )
+    size: Optional[str] = Field(
+        None, description='Size of the output image', examples=['1024x1024']
+    )
+    user: Optional[str] = Field(
+        None,
+        description='A unique identifier for end-user monitoring',
+        examples=['user-1234'],
+    )
+
+
+class Quality1(str, Enum):
+    low = 'low'
+    medium = 'medium'
+    high = 'high'
+    standard = 'standard'
+    hd = 'hd'
+
+
 class ResponseFormat(str, Enum):
     url = 'url'
     b64_json = 'b64_json'
@@ -1117,9 +1139,7 @@ class OpenAIImageGenerationRequest(BaseModel):
         None, description='Background transparency', examples=['opaque']
     )
     model: Optional[str] = Field(
-        None,
+        None, description='The model to use for image generation', examples=['dall-e-3']
-        description='The model to use for image generation',
-        examples=['gpt-image-1'],
     )
     moderation: Optional[Moderation] = Field(
         None, description='Content moderation setting', examples=['auto']
@@ -1140,7 +1160,7 @@ class OpenAIImageGenerationRequest(BaseModel):
         description='A text description of the desired image',
         examples=['Draw a rocket in front of a blackhole in deep space'],
     )
-    quality: Optional[Quality] = Field(
+    quality: Optional[Quality1] = Field(
         None, description='The quality of the generated image', examples=['high']
     )
     response_format: Optional[ResponseFormat] = Field(
@@ -1500,18 +1520,18 @@ class KlingAuthenticationError(KlingErrorResponse):
 
 
 class LumaGenerationRequest(BaseModel):
-    aspect_ratio: Optional[LumaAspectRatio] = '16:9'
+    aspect_ratio: LumaAspectRatio
     callback_url: Optional[AnyUrl] = Field(
         None,
         description='The callback URL of the generation, a POST request with Generation object will be sent to the callback URL when the generation is dreaming, completed, or failed',
     )
-    duration: Optional[LumaVideoModelOutputDuration] = None
+    duration: LumaVideoModelOutputDuration
     generation_type: Optional[GenerationType1] = 'video'
     keyframes: Optional[LumaKeyframes] = None
     loop: Optional[bool] = Field(None, description='Whether to loop the video')
-    model: Optional[LumaVideoModel] = 'ray-1-6'
+    model: LumaVideoModel
-    prompt: Optional[str] = Field(None, description='The prompt of the generation')
+    prompt: str = Field(..., description='The prompt of the generation')
-    resolution: Optional[LumaVideoModelOutputResolution] = None
+    resolution: LumaVideoModelOutputResolution
 
 
 class CharacterRef(BaseModel):

comfy_api_nodes/nodes_api.py

@@ -1,16 +1,60 @@
 import io
 from inspect import cleandoc
 
+from comfy.utils import common_upscale
 from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
 from comfy_api_nodes.apis import (
     IdeogramGenerateRequest,
     IdeogramGenerateResponse,
     ImageRequest,
     OpenAIImageGenerationRequest,
+    OpenAIImageEditRequest,
     OpenAIImageGenerationResponse
 )
 from comfy_api_nodes.apis.client import ApiEndpoint, HttpMethod, SynchronousOperation
 
+import numpy as np
+from PIL import Image
+import requests
+import torch
+import math
+
+def downscale_input(image):
+    samples = image.movedim(-1,1)
+    #downscaling input images to roughly the same size as the outputs
+    total = int(1024 * 1024)
+    scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
+    if scale_by >= 1:
+        return (image,)
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+
+    s = common_upscale(samples, width, height, "lanczos", "disabled")
+    s = s.movedim(1,-1)
+    return s
+
+def validate_and_cast_response (response):
+    # validate raw JSON response
+    data = response.data
+    if not data or len(data) == 0:
+        raise Exception("No images returned from API endpoint")
+
+    # Get base64 image data
+    image_url = data[0].url
+    if not image_url:
+        raise Exception("No image URL was generated in the response")
+    img_response = requests.get(image_url)
+    if img_response.status_code != 200:
+        raise Exception("Failed to download the image")
+
+    img = Image.open(io.BytesIO(img_response.content))
+    img = img.convert("RGB")  # Ensure RGB format
+
+    # Convert to numpy array, normalize to float32 between 0 and 1
+    img_array = np.array(img).astype(np.float32) / 255.0
+
+    # Convert to torch tensor and add batch dimension
+    return torch.from_numpy(img_array)[None,]
 
 class IdeogramTextToImage(ComfyNodeABC):
     """
@@ -165,11 +209,11 @@ class IdeogramTextToImage(ComfyNodeABC):
     #def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
     #    return ""
 
-class OpenAITextToImage(ComfyNodeABC):
+class OpenAIDalle2(ComfyNodeABC):
     """
-    Generates images synchronously via OpenAI's DALL·E 3 endpoint.
+    Generates images synchronously via OpenAI's DALL·E 2 endpoint.
 
-    Uses the proxy at /proxy/dalle-3/generate. Returned URLs are short‑lived,
+    Uses the proxy at /proxy/openai/images/generations. Returned URLs are short‑lived,
     so download or cache results if you need to keep them.
     """
     def __init__(self):
@@ -184,14 +228,21 @@ class OpenAITextToImage(ComfyNodeABC):
                     "default": "",
                     "tooltip": "Text prompt for DALL·E",
                 }),
-                # TODO: add NEW MODEL
-                "model": (IO.COMBO, {
-                    "options": ["dall-e-3", "dall-e-2"],
-                    "default": "dall-e-3",
-                    "tooltip": "OpenAI model name",
-                }),
             },
             "optional": {
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 2**31-1,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "not implemented yet in backend",
+                }),
+                "size": (IO.COMBO, {
+                    "options": ["256x256", "512x512", "1024x1024"],
+                    "default": "1024x1024",
+                    "tooltip": "Image size",
+                }),
                 "n": (IO.INT, {
                     "default": 1,
                     "min": 1,
@@ -200,18 +251,13 @@ class OpenAITextToImage(ComfyNodeABC):
                     "display": "number",
                     "tooltip": "How many images to generate",
                 }),
-                "size": (IO.COMBO, {
+                "image": (IO.IMAGE, {
-                    "options": ["256x256", "512x512", "1024x1792", "1792x1024", "1024x1024", "1536x1024", "1024x1536", "auto"],
+                    "default": None,
-                    "default": "auto",
+                    "tooltip": "Optional reference image for image editing.",
-                    "tooltip": "Image size",
                 }),
-                "seed": (IO.INT, {
+                "mask": (IO.MASK, {
-                    "default": 0,
+                    "default": None,
-                    "min": 0,
+                    "tooltip": "Optional mask for inpainting (white areas will be replaced)",
-                    "max": 2**31-1,
-                    "step": 1,
-                    "display": "number",
-                    "tooltip": "Optional random seed",
                 }),
             },
             "hidden": {
@@ -225,28 +271,121 @@ class OpenAITextToImage(ComfyNodeABC):
     DESCRIPTION = cleandoc(__doc__ or "")
     API_NODE = True
 
-    def api_call(self, prompt, model, n=1, size="1024x1024", seed=0, auth_token=None):
+    def api_call(self, prompt, seed=0, image=None, mask=None, n=1, size="1024x1024", auth_token=None):
-        # Validate size based on model
+        model = "dall-e-2"
-        if model == "dall-e-2":
+        path = "/proxy/openai/images/generations"
-            if size == "auto":
+        request_class = OpenAIImageGenerationRequest
-                size = "1024x1024"
+        img_binary = None
-            valid_sizes = ["256x256", "512x512", "1024x1024"]
-            if size not in valid_sizes:
-                raise ValueError(f"Size {size} not valid for dall-e-2. Must be one of: {', '.join(valid_sizes)}")
-        elif model == "dall-e-3":
-            if size == "auto":
-                size = "1024x1024"
-            valid_sizes = ["1024x1024", "1792x1024", "1024x1792"]
-            if size not in valid_sizes:
-                raise ValueError(f"Size {size} not valid for dall-e-3. Must be one of: {', '.join(valid_sizes)}")
-        # TODO: add NEW MODEL
 
+        if image is not None and mask is not None:
+            path = "/proxy/openai/images/edits"
+            request_class = OpenAIImageEditRequest
 
+            input_tensor = image.squeeze().cpu()
+            height, width, channels = input_tensor.shape
+            rgba_tensor = torch.ones(height, width, 4, device="cpu")
+            rgba_tensor[:, :, :channels] = input_tensor
 
-        import numpy as np
+            if mask.shape[1:] != image.shape[1:-1]:
-        import torch
+                raise Exception("Mask and Image must be the same size")
-        from PIL import Image
+            rgba_tensor[:,:,3] = (1-mask.squeeze().cpu())
-        import requests
+
+            rgba_tensor = downscale_input(rgba_tensor.unsqueeze(0)).squeeze()
+
+            image_np = (rgba_tensor.numpy() * 255).astype(np.uint8)
+            img = Image.fromarray(image_np)
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            img_byte_arr.seek(0)
+            img_binary = img_byte_arr#.getvalue()
+            img_binary.name = "image.png"
+        elif image is not None or mask is not None:
+            raise Exception("Dall-E 2 image editing requires an image AND a mask")
+
+        # Build the operation
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=path,
+                method=HttpMethod.POST,
+                request_model=request_class,
+                response_model=OpenAIImageGenerationResponse
+            ),
+            request=request_class(
+                model=model,
+                prompt=prompt,
+                n=n,
+                size=size,
+                seed=seed,
+            ),
+            files={
+                "image": img_binary,
+            } if img_binary else None,
+            auth_token=auth_token
+        )
+
+        response = operation.execute()
+
+        img_tensor = validate_and_cast_response(response)
+        return (img_tensor,)
+
+class OpenAIDalle3(ComfyNodeABC):
+    """
+    Generates images synchronously via OpenAI's DALL·E 3 endpoint.
+
+    Uses the proxy at /proxy/openai/images/generations. Returned URLs are short‑lived,
+    so download or cache results if you need to keep them.
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls) -> InputTypeDict:
+        return {
+            "required": {
+                "prompt": (IO.STRING, {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Text prompt for DALL·E",
+                }),
+            },
+            "optional": {
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 2**31-1,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "not implemented yet in backend",
+                }),
+                "quality" : (IO.COMBO, {
+                    "options": ["standard","hd"],
+                    "default": "standard",
+                    "tooltip": "Image quality",
+                }),
+                "style": (IO.COMBO, {
+                    "options": ["natural","vivid"],
+                    "default": "natural",
+                    "tooltip": "Vivid causes the model to lean towards generating hyper-real and dramatic images. Natural causes the model to produce more natural, less hyper-real looking images.",
+                }),
+                "size": (IO.COMBO, {
+                    "options": ["1024x1024", "1024x1792", "1792x1024"],
+                    "default": "1024x1024",
+                    "tooltip": "Image size",
+                }),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG"
+            }
+        }
+
+    RETURN_TYPES = (IO.IMAGE,)
+    FUNCTION = "api_call"
+    CATEGORY = "Example"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    API_NODE = True
+
+    def api_call(self, prompt, seed=0, style="natural", quality="standard", size="1024x1024", auth_token=None):
+        model = "dall-e-3"
 
         # build the operation
         operation = SynchronousOperation(
@@ -259,49 +398,175 @@ class OpenAITextToImage(ComfyNodeABC):
             request=OpenAIImageGenerationRequest(
                 model=model,
                 prompt=prompt,
-                n=n,
+                quality=quality,
                 size=size,
-                seed=seed if seed != 0 else None
+                style=style,
+                seed=seed,
             ),
             auth_token=auth_token
         )
 
         response = operation.execute()
 
-        # validate raw JSON response
+        img_tensor = validate_and_cast_response(response)
-
-        data = response.data
-        if not data or len(data) == 0:
-            raise Exception("No images returned from OpenAI endpoint")
-
-        # Get base64 image data
-        image_url = data[0].url
-        if not image_url:
-            raise Exception("No image URL was generated in the response")
-        img_response = requests.get(image_url)
-        if img_response.status_code != 200:
-            raise Exception("Failed to download the image")
-
-        img = Image.open(io.BytesIO(img_response.content))
-        img = img.convert("RGB")  # Ensure RGB format
-
-        # Convert to numpy array, normalize to float32 between 0 and 1
-        img_array = np.array(img).astype(np.float32) / 255.0
-
-        # Convert to torch tensor and add batch dimension
-        img_tensor = torch.from_numpy(img_array)[None,]
-
         return (img_tensor,)
 
+class OpenAIXXX(ComfyNodeABC):
+    """
+    Generates images synchronously via OpenAI's DALL·E 2 endpoint.
+
+    Uses the proxy at /proxy/openai/images/generations. Returned URLs are short‑lived,
+    so download or cache results if you need to keep them.
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls) -> InputTypeDict:
+        return {
+            "required": {
+                "prompt": (IO.STRING, {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Text prompt for XXX",
+                }),
+            },
+            "optional": {
+                "seed": (IO.INT, {
+                    "default": 0,
+                    "min": 0,
+                    "max": 2**31-1,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "not implemented yet in backend",
+                }),
+                "quality": (IO.COMBO, {
+                    "options": ["low","medium","high"],
+                    "default": "low",
+                    "tooltip": "Image quality, affects cost and generation time.",
+                }),
+                "background": (IO.COMBO, {
+                    "options": ["opaque","transparent"],
+                    "default": "opaque",
+                    "tooltip": "Return image with or without background",
+                }),
+                "size": (IO.COMBO, {
+                    "options": ["auto", "1024x1024", "1024x1536", "1536x1024"],
+                    "default": "auto",
+                    "tooltip": "Image size",
+                }),
+                "n": (IO.INT, {
+                    "default": 1,
+                    "min": 1,
+                    "max": 8,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "How many images to generate",
+                }),
+                "image": (IO.IMAGE, {
+                    "default": None,
+                    "tooltip": "Optional reference image for image editing.",
+                }),
+                "mask": (IO.MASK, {
+                    "default": None,
+                    "tooltip": "Optional mask for inpainting (white areas will be replaced)",
+                }),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG"
+            }
+        }
+
+    RETURN_TYPES = (IO.IMAGE,)
+    FUNCTION = "api_call"
+    CATEGORY = "Example"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    API_NODE = True
+
+    def api_call(self, prompt, seed=0, quality="low", background="opaque", image=None, mask=None, n=1, size="1024x1024", auth_token=None):
+        model = "xxx"
+        path = "/proxy/openai/images/generations"
+        request_class = OpenAIImageGenerationRequest
+        img_binary = None
+        mask_binary = None
+
+
+        if image is not None:
+            path = "/proxy/openai/images/edits"
+            request_class = OpenAIImageEditRequest
+
+            scaled_image = downscale_input(image).squeeze()
+
+            image_np = (scaled_image.numpy() * 255).astype(np.uint8)
+            img = Image.fromarray(image_np)
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            img_byte_arr.seek(0)
+            img_binary = img_byte_arr#.getvalue()
+            img_binary.name = "image.png"
+
+        if mask is not None:
+            if image is None:
+                raise Exception("Cannot use a mask without an input image")
+            if mask.shape[1:] != image.shape[1:-1]:
+                raise Exception("Mask and Image must be the same size")
+            batch, height, width = mask.shape
+            rgba_mask = torch.zeros(height, width, 4, device="cpu")
+            rgba_mask[:,:,3] = (1-mask.squeeze().cpu())
+            mask_np = (rgba_mask.numpy() * 255).astype(np.uint8)
+            mask_img = Image.fromarray(mask_np)
+            mask_img_byte_arr = io.BytesIO()
+            mask_img.save(mask_img_byte_arr, format='PNG')
+            mask_img_byte_arr.seek(0)
+            mask_binary = mask_img_byte_arr#.getvalue()
+            mask_binary.name = "mask.png"
+
+        files = {}
+        if img_binary:
+            files["image"] = img_binary
+        if mask_binary:
+            files["mask"] = mask_binary
+
+        # Build the operation
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=path,
+                method=HttpMethod.POST,
+                request_model=request_class,
+                response_model=OpenAIImageGenerationResponse
+            ),
+            request=request_class(
+                model=model,
+                prompt=prompt,
+                quality=quality,
+                background=background,
+                n=n,
+                seed=seed,
+                size=size,
+            ),
+            files=files if files else None,
+            auth_token=auth_token
+        )
+
+        response = operation.execute()
+
+        img_tensor = validate_and_cast_response(response)
+        return (img_tensor,)
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "IdeogramTextToImage": IdeogramTextToImage,
-    "OpenAIDalleTextToImage": OpenAITextToImage,
+    "OpenAIDalle2": OpenAIDalle2,
+    "OpenAIDalle3": OpenAIDalle3,
+    "OpenAIXXX": OpenAIXXX,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "IdeogramTextToImage": "Ideogram Text to Image",
-    "OpenAIDalleTextToImage": "OpenAI DALL·E 3 Text to Image",
+    "OpenAIDalle2": "OpenAI DALL·E 2",
+    "OpenAIDalle3": "OpenAI DALL·E 3",
+    "OpenAIXXX": "XXX",
 }