Skip to content

Commit cfdfede

Browse files
authored
feat(api-nodes): add price extractor feature; small fixes to Kling & Pika nodes (comfyanonymous#10284)
1 parent 26b3c5c commit cfdfede

File tree

3 files changed

+33
-17
lines changed

3 files changed

+33
-17
lines changed

comfy_api_nodes/apis/client.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -782,9 +782,11 @@ def __init__(
782782
poll_endpoint: ApiEndpoint[EmptyRequest, R],
783783
completed_statuses: list[str],
784784
failed_statuses: list[str],
785+
*,
785786
status_extractor: Callable[[R], Optional[str]],
786787
progress_extractor: Callable[[R], Optional[float]] | None = None,
787788
result_url_extractor: Callable[[R], Optional[str]] | None = None,
789+
price_extractor: Callable[[R], Optional[float]] | None = None,
788790
request: Optional[T] = None,
789791
api_base: str | None = None,
790792
auth_token: Optional[str] = None,
@@ -815,10 +817,12 @@ def __init__(
815817
self.status_extractor = status_extractor or (lambda x: getattr(x, "status", None))
816818
self.progress_extractor = progress_extractor
817819
self.result_url_extractor = result_url_extractor
820+
self.price_extractor = price_extractor
818821
self.node_id = node_id
819822
self.completed_statuses = completed_statuses
820823
self.failed_statuses = failed_statuses
821824
self.final_response: Optional[R] = None
825+
self.extracted_price: Optional[float] = None
822826

823827
async def execute(self, client: Optional[ApiClient] = None) -> R:
824828
owns_client = client is None
@@ -840,6 +844,8 @@ async def execute(self, client: Optional[ApiClient] = None) -> R:
840844
def _display_text_on_node(self, text: str):
841845
if not self.node_id:
842846
return
847+
if self.extracted_price is not None:
848+
text = f"Price: {self.extracted_price}$\n{text}"
843849
PromptServer.instance.send_progress_text(text, self.node_id)
844850

845851
def _display_time_progress_on_node(self, time_completed: int | float):
@@ -877,9 +883,7 @@ async def _poll_until_complete(self, client: ApiClient) -> R:
877883
try:
878884
logging.debug("[DEBUG] Polling attempt #%s", poll_count)
879885

880-
request_dict = (
881-
None if self.request is None else self.request.model_dump(exclude_none=True)
882-
)
886+
request_dict = None if self.request is None else self.request.model_dump(exclude_none=True)
883887

884888
if poll_count == 1:
885889
logging.debug(
@@ -912,6 +916,11 @@ async def _poll_until_complete(self, client: ApiClient) -> R:
912916
if new_progress is not None:
913917
progress.update_absolute(new_progress, total=PROGRESS_BAR_MAX)
914918

919+
if self.price_extractor:
920+
price = self.price_extractor(response_obj)
921+
if price is not None:
922+
self.extracted_price = price
923+
915924
if status == TaskStatus.COMPLETED:
916925
message = "Task completed successfully"
917926
if self.result_url_extractor:

comfy_api_nodes/nodes_kling.py

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
validate_video_dimensions,
7474
validate_video_duration,
7575
)
76+
from comfy_api.input_impl import VideoFromFile
7677
from comfy_api.input.basic_types import AudioInput
7778
from comfy_api.input.video_types import VideoInput
7879
from comfy_api.latest import ComfyExtension, io as comfy_io
@@ -511,7 +512,7 @@ async def execute_video_effect(
511512
image_1: torch.Tensor,
512513
image_2: Optional[torch.Tensor] = None,
513514
model_mode: Optional[KlingVideoGenMode] = None,
514-
) -> comfy_io.NodeOutput:
515+
) -> tuple[VideoFromFile, str, str]:
515516
if dual_character:
516517
request_input_field = KlingDualCharacterEffectInput(
517518
model_name=model_name,
@@ -562,7 +563,7 @@ async def execute_video_effect(
562563
validate_video_result_response(final_response)
563564

564565
video = get_video_from_response(final_response)
565-
return comfy_io.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
566+
return await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration)
566567

567568

568569
async def execute_lipsync(
@@ -1271,7 +1272,7 @@ async def execute(
12711272
image_1=image_left,
12721273
image_2=image_right,
12731274
)
1274-
return video, duration
1275+
return comfy_io.NodeOutput(video, duration)
12751276

12761277

12771278
class KlingSingleImageVideoEffectNode(comfy_io.ComfyNode):
@@ -1320,17 +1321,21 @@ async def execute(
13201321
model_name: KlingSingleImageEffectModelName,
13211322
duration: KlingVideoGenDuration,
13221323
) -> comfy_io.NodeOutput:
1323-
return await execute_video_effect(
1324-
auth_kwargs={
1325-
"auth_token": cls.hidden.auth_token_comfy_org,
1326-
"comfy_api_key": cls.hidden.api_key_comfy_org,
1327-
},
1328-
node_id=cls.hidden.unique_id,
1329-
dual_character=False,
1330-
effect_scene=effect_scene,
1331-
model_name=model_name,
1332-
duration=duration,
1333-
image_1=image,
1324+
return comfy_io.NodeOutput(
1325+
*(
1326+
await execute_video_effect(
1327+
auth_kwargs={
1328+
"auth_token": cls.hidden.auth_token_comfy_org,
1329+
"comfy_api_key": cls.hidden.api_key_comfy_org,
1330+
},
1331+
node_id=cls.hidden.unique_id,
1332+
dual_character=False,
1333+
effect_scene=effect_scene,
1334+
model_name=model_name,
1335+
duration=duration,
1336+
image_1=image,
1337+
)
1338+
)
13341339
)
13351340

13361341

comfy_api_nodes/nodes_pika.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from comfy_api_nodes.apinode_utils import (
1818
download_url_to_video_output,
1919
tensor_to_bytesio,
20+
validate_string,
2021
)
2122
from comfy_api_nodes.apis import pika_defs
2223
from comfy_api_nodes.apis.client import (
@@ -590,6 +591,7 @@ async def execute(
590591
resolution: str,
591592
duration: int,
592593
) -> comfy_io.NodeOutput:
594+
validate_string(prompt_text, field_name="prompt_text", min_length=1)
593595
pika_files = [
594596
("keyFrames", ("image_start.png", tensor_to_bytesio(image_start), "image/png")),
595597
("keyFrames", ("image_end.png", tensor_to_bytesio(image_end), "image/png")),

0 commit comments

Comments
 (0)