Fix runway image upload and progress polling (#39)

comfy_api_nodes/apis/__init__.py

@@ -1355,7 +1355,7 @@ class RunwayPromptImageDetailedObject(BaseModel):
         ...,
         description="The position of the image in the output video. 'last' is currently supported for gen3a_turbo only.",
     )
-    uri: AnyUrl = Field(
+    uri: str = Field(
         ..., description='A HTTPS URL or data URI containing an encoded image.'
     )
 

comfy_api_nodes/apis/client.py

@@ -105,7 +105,7 @@ from typing import (
     TypeVar,
     Generic,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, HttpUrl
 from enum import Enum
 import json
 import requests
@@ -127,11 +127,11 @@ class EmptyRequest(BaseModel):
 
 class UploadRequest(BaseModel):
     filename: str = Field(..., description="Filename to upload")
-    mime_type: str = Field(..., description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.")
+    content_type: str = Field(..., description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.")
 
 class UploadResponse(BaseModel):
-    download_url: str = Field(..., description='URL to GET uploaded file')
+    download_url: HttpUrl = Field(..., description='URL to GET uploaded file')
-    upload_url: str = Field(..., description='URL to PUT file to upload')
+    upload_url: HttpUrl = Field(..., description='URL to PUT file to upload')
 
 
 class HttpMethod(str, Enum):
@@ -297,7 +297,7 @@ class ApiClient:
     def upload_file(
         upload_url: str,
         file: io.BytesIO | str,
-        mime_type: str | None = None,
+        content_type: str | None = None,
     ):
         """Upload a file to the API. Make sure the file has a filename equal to what the url expects.
 
@@ -307,8 +307,8 @@ class ApiClient:
             mime_type: Optional mime type to set for the upload
         """
         headers = {}
-        if mime_type:
+        if content_type:
-            headers["Content-Type"] = mime_type
+            headers["Content-Type"] = content_type
 
         if isinstance(file, io.BytesIO):
             file.seek(0)  # Ensure we're at the start of the file
@@ -503,7 +503,6 @@ class PollingOperation(Generic[T, R]):
     def _poll_until_complete(self, client: ApiClient) -> R:
         """Poll until the task is complete"""
         poll_count = 0
-        progress = 0
         if self.progress_extractor:
             progress = utils.ProgressBar(100)
 
@@ -542,11 +541,15 @@ class PollingOperation(Generic[T, R]):
 
                 # If progress extractor is provided, extract progress
                 if self.progress_extractor:
-                    progress.update(self.progress_extractor(response_obj))
+                    new_progress = self.progress_extractor(response_obj)
+                    if new_progress is not None:
+                        progress.update(new_progress)
 
                 if status == TaskStatus.COMPLETED:
                     logging.debug("[DEBUG] Task completed successfully")
                     self.final_response = response_obj
+                    if self.progress_extractor:
+                        progress.update(100)
                     return self.final_response
                 elif status == TaskStatus.FAILED:
                     logging.debug(f"[DEBUG] Task failed: {json.dumps(resp)}")

comfy_api_nodes/nodes_api.py

@@ -327,13 +327,13 @@ def upload_images_to_comfyapi(
                 request_model=UploadRequest,
                 response_model=UploadResponse,
             ),
-            request=UploadRequest(filename=img_binary.name, mime_type=mime_type),
+            request=UploadRequest(filename=img_binary.name, content_type=mime_type),
             auth_token=auth_token,
         )
         response = operation.execute()
 
         upload_response = ApiClient.upload_file(
-            response.upload_url, img_binary, mime_type=mime_type
+            response.upload_url, img_binary, content_type=mime_type
         )
         # verify success
         try:

comfy_api_nodes/nodes_runway.py

@@ -30,6 +30,7 @@ from comfy_api.input_impl import VideoFromFile
 from comfy_api_nodes.mapper_utils import model_field_to_node_input
 
 PATH_IMAGE_TO_VIDEO = "/proxy/runway/image-to-video"
+PATH_GET_TASK_STATUS = "/proxy/runway/tasks"
 
 
 class RunwayApiError(Exception):
@@ -38,6 +39,12 @@ class RunwayApiError(Exception):
     pass
 
 
+def extract_progress_from_task_status(response: TaskStatusResponse) -> float:
+    if hasattr(response, "progress") and response.progress is not None:
+        return response.progress * 100
+    return None
+
+
 class RunwayImageToVideoNode(ComfyNodeABC):
     """
     Runway Image to Video Node.
@@ -86,7 +93,7 @@ class RunwayImageToVideoNode(ComfyNodeABC):
         """
         polling_operation = PollingOperation(
             poll_endpoint=ApiEndpoint(
-                path=f"{PATH_IMAGE_TO_VIDEO}/{task_id}",
+                path=f"{PATH_GET_TASK_STATUS}/{task_id}",
                 method=HttpMethod.GET,
                 request_model=EmptyRequest,
                 response_model=TaskStatusResponse,
@@ -98,7 +105,7 @@ class RunwayImageToVideoNode(ComfyNodeABC):
                 TaskStatus.FAILED.value,
                 TaskStatus.CANCELLED.value,
             ],
-            progress_extractor=lambda response: (response.progress * 100),
+            progress_extractor=extract_progress_from_task_status,
             status_extractor=lambda response: (response.status.value),
             auth_token=auth_token,
         )
@@ -193,7 +200,10 @@ class RunwayImageToVideoNode(ComfyNodeABC):
         prompt_images_tensor = torch.cat(prompt_images_tensors, dim=0)
 
         download_urls = upload_images_to_comfyapi(
-            prompt_images_tensor, max_images=2, auth_token=auth_token, mime_type="image/png"
+            prompt_images_tensor,
+            max_images=2,
+            auth_token=auth_token,
+            mime_type="image/png",
         )
 
         # Create a list of detailed image objects
@@ -202,15 +212,15 @@ class RunwayImageToVideoNode(ComfyNodeABC):
         ]
         if len(download_urls) > 1:
             prompt_image_details.append(
-                RunwayPromptImageDetailedObject(uri=str(download_urls[1]), position="last")
+                RunwayPromptImageDetailedObject(
+                    uri=str(download_urls[1]), position="last"
+                )
             )
 
         # Wrap the list in the main object if details exist
         prompt_image_object: Optional[RunwayPromptImageObject] = None
         if prompt_image_details:
-            prompt_image_object = RunwayPromptImageObject(
+            prompt_image_object = RunwayPromptImageObject(root=prompt_image_details)
-                root=prompt_image_details
-            )
 
         initial_operation = SynchronousOperation(
             endpoint=ApiEndpoint(