Add Runway Gen 4 Text to Image Node (#193)

comfy_api_nodes/nodes_runway.py

@@ -26,6 +26,10 @@ from comfy_api_nodes.apis import (
     RunwayAspectRatioEnum as AspectRatio,
     RunwayPromptImageObject,
     RunwayPromptImageDetailedObject,
+    RunwayTextToImageRequest,
+    RunwayTextToImageResponse,
+    Model4,
+    ReferenceImage,
 )
 from comfy_api_nodes.apis.client import (
     ApiEndpoint,
@@ -39,16 +43,19 @@ from comfy_api_nodes.apinode_utils import (
     download_url_to_video_output,
     image_tensor_pair_to_batch,
     validate_string,
+    download_url_to_image_tensor,
 )
 from comfy_api_nodes.mapper_utils import model_field_to_node_input
 from comfy_api.input_impl import VideoFromFile
 from comfy.comfy_types.node_typing import IO, ComfyNodeABC
 
 PATH_IMAGE_TO_VIDEO = "/proxy/runway/image-to-video"
+PATH_TEXT_TO_IMAGE = "/proxy/runway/text_to_image"
 PATH_GET_TASK_STATUS = "/proxy/runway/tasks"
 
 AVERAGE_DURATION_I2V_SECONDS = 128
 AVERAGE_DURATION_FLF_SECONDS = 256
+AVERAGE_DURATION_T2I_SECONDS = 32
 
 
 class RunwayApiError(Exception):
@@ -68,6 +75,24 @@ class RunwayBasicAspectRatio(str, Enum):
     field_1584_672 = "1584:672"
 
 
+# TODO: replace with enum after it's added in comfy-api and generated by code gen
+class RunwayTextToImageRatio(str, Enum):
+    """Aspect ratios supported for Text to Image API."""
+
+    # 16:9 landscape
+    field_16_9 = "1920:1088"
+    # 9:16 portrait
+    field_9_16 = "1088:1920"
+    # 1:1 square
+    field_1_1 = "1088:1088"
+    # 4:3 landscape
+    field_4_3 = "1456:1088"
+    # 3:4 portrait
+    field_3_4 = "1088:1456"
+    # 21:9 ultrawide
+    field_21_9 = "2112:912"
+
+
 def get_video_url_from_task_status(response: TaskStatusResponse) -> Union[str, None]:
     """Returns the video URL from the task status response if it exists."""
     if response.output and len(response.output) > 0:
@@ -117,6 +142,13 @@ def extract_progress_from_task_status(
     return None
 
 
+def get_image_url_from_task_status(response: TaskStatusResponse) -> Union[str, None]:
+    """Returns the image URL from the task status response if it exists."""
+    if response.output and len(response.output) > 0:
+        return response.output[0]
+    return None
+
+
 class RunwayVideoGenNode(ComfyNodeABC):
     """Runway Video Node Base."""
 
@@ -383,12 +415,147 @@ class RunwayStartEndFrameNode(RunwayVideoGenNode):
         )
 
 
+class RunwayTextToImageNode(ComfyNodeABC):
+    """Runway Text to Image Node."""
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "api_call"
+    CATEGORY = "api node/image/Runway"
+    API_NODE = True
+    DESCRIPTION = "Generate an image from a text prompt using Runway's Gen 4 model. You can also include reference images to guide the generation."
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": model_field_to_node_input(
+                    IO.STRING, RunwayTextToImageRequest, "promptText", multiline=True
+                ),
+                "ratio": model_field_to_node_input(
+                    IO.COMBO,
+                    RunwayTextToImageRequest,
+                    "ratio",
+                    enum_type=RunwayTextToImageRatio,
+                ),
+            },
+            "optional": {
+                "reference_image": (
+                    IO.IMAGE,
+                    {"tooltip": "Optional reference image to guide the generation"},
+                )
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    def validate_task_created(self, response: RunwayTextToImageResponse) -> bool:
+        """
+        Validate the task creation response from the Runway API matches
+        expected format.
+        """
+        if not bool(response.id):
+            raise RunwayApiError("Invalid initial response from Runway API.")
+        return True
+
+    def validate_response(self, response: TaskStatusResponse) -> bool:
+        """
+        Validate the successful task status response from the Runway API
+        matches expected format.
+        """
+        if not response.output or len(response.output) == 0:
+            raise RunwayApiError(
+                "Runway task succeeded but no image data found in response."
+            )
+        return True
+
+    def get_response(
+        self, task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
+    ) -> TaskStatusResponse:
+        """Poll the task status until it is finished then get the response."""
+        return poll_until_finished(
+            auth_kwargs,
+            ApiEndpoint(
+                path=f"{PATH_GET_TASK_STATUS}/{task_id}",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=TaskStatusResponse,
+            ),
+            estimated_duration=AVERAGE_DURATION_T2I_SECONDS,
+            node_id=node_id,
+        )
+
+    def api_call(
+        self,
+        prompt: str,
+        ratio: str,
+        reference_image: Optional[torch.Tensor] = None,
+        unique_id: Optional[str] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor]:
+        # Validate inputs
+        validate_string(prompt, min_length=1)
+
+        # Prepare reference images if provided
+        reference_images = None
+        if reference_image is not None:
+            validate_input_image(reference_image)
+            download_urls = upload_images_to_comfyapi(
+                reference_image,
+                max_images=1,
+                mime_type="image/png",
+                auth_kwargs=kwargs,
+            )
+            if len(download_urls) != 1:
+                raise RunwayApiError("Failed to upload reference image to comfy api.")
+
+            reference_images = [ReferenceImage(uri=str(download_urls[0]))]
+
+        # Create request
+        request = RunwayTextToImageRequest(
+            promptText=prompt,
+            model=Model4.gen4_image,
+            ratio=ratio,
+            referenceImages=reference_images,
+        )
+
+        # Execute initial request
+        initial_operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=PATH_TEXT_TO_IMAGE,
+                method=HttpMethod.POST,
+                request_model=RunwayTextToImageRequest,
+                response_model=RunwayTextToImageResponse,
+            ),
+            request=request,
+            auth_kwargs=kwargs,
+        )
+
+        initial_response = initial_operation.execute()
+        self.validate_task_created(initial_response)
+        task_id = initial_response.id
+
+        # Poll for completion
+        final_response = self.get_response(
+            task_id, auth_kwargs=kwargs, node_id=unique_id
+        )
+        self.validate_response(final_response)
+
+        # Download and return image
+        image_url = get_image_url_from_task_status(final_response)
+        return (download_url_to_image_tensor(image_url),)
+
+
 NODE_CLASS_MAPPINGS = {
     "RunwayStartEndFrameNode": RunwayStartEndFrameNode,
     "RunwayImageToVideoNode": RunwayImageToVideoNode,
+    "RunwayTextToImageNode": RunwayTextToImageNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "RunwayStartEndFrameNode": "Runway Start-End Frame to Video",
     "RunwayImageToVideoNode": "Runway Image to Video",
+    "RunwayTextToImageNode": "Runway Text to Image",
 }