feat(api-nodes): add price extractor feature; small fixes to Kling & Pika nodes (#10284)

This commit is contained in:
Alexander Piskun 2025-10-11 02:21:40 +03:00 committed by GitHub
parent aa895db7e8
commit 14d642acd6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 33 additions and 17 deletions

View File

@ -782,9 +782,11 @@ class PollingOperation(Generic[T, R]):
poll_endpoint: ApiEndpoint[EmptyRequest, R], poll_endpoint: ApiEndpoint[EmptyRequest, R],
completed_statuses: list[str], completed_statuses: list[str],
failed_statuses: list[str], failed_statuses: list[str],
*,
status_extractor: Callable[[R], Optional[str]], status_extractor: Callable[[R], Optional[str]],
progress_extractor: Callable[[R], Optional[float]] | None = None, progress_extractor: Callable[[R], Optional[float]] | None = None,
result_url_extractor: Callable[[R], Optional[str]] | None = None, result_url_extractor: Callable[[R], Optional[str]] | None = None,
price_extractor: Callable[[R], Optional[float]] | None = None,
request: Optional[T] = None, request: Optional[T] = None,
api_base: str | None = None, api_base: str | None = None,
auth_token: Optional[str] = None, auth_token: Optional[str] = None,
@ -815,10 +817,12 @@ class PollingOperation(Generic[T, R]):
self.status_extractor = status_extractor or (lambda x: getattr(x, "status", None)) self.status_extractor = status_extractor or (lambda x: getattr(x, "status", None))
self.progress_extractor = progress_extractor self.progress_extractor = progress_extractor
self.result_url_extractor = result_url_extractor self.result_url_extractor = result_url_extractor
self.price_extractor = price_extractor
self.node_id = node_id self.node_id = node_id
self.completed_statuses = completed_statuses self.completed_statuses = completed_statuses
self.failed_statuses = failed_statuses self.failed_statuses = failed_statuses
self.final_response: Optional[R] = None self.final_response: Optional[R] = None
self.extracted_price: Optional[float] = None
async def execute(self, client: Optional[ApiClient] = None) -> R: async def execute(self, client: Optional[ApiClient] = None) -> R:
owns_client = client is None owns_client = client is None
@ -840,6 +844,8 @@ class PollingOperation(Generic[T, R]):
def _display_text_on_node(self, text: str): def _display_text_on_node(self, text: str):
if not self.node_id: if not self.node_id:
return return
if self.extracted_price is not None:
text = f"Price: {self.extracted_price}$\n{text}"
PromptServer.instance.send_progress_text(text, self.node_id) PromptServer.instance.send_progress_text(text, self.node_id)
def _display_time_progress_on_node(self, time_completed: int | float): def _display_time_progress_on_node(self, time_completed: int | float):
@ -877,9 +883,7 @@ class PollingOperation(Generic[T, R]):
try: try:
logging.debug("[DEBUG] Polling attempt #%s", poll_count) logging.debug("[DEBUG] Polling attempt #%s", poll_count)
request_dict = ( request_dict = None if self.request is None else self.request.model_dump(exclude_none=True)
None if self.request is None else self.request.model_dump(exclude_none=True)
)
if poll_count == 1: if poll_count == 1:
logging.debug( logging.debug(
@ -912,6 +916,11 @@ class PollingOperation(Generic[T, R]):
if new_progress is not None: if new_progress is not None:
progress.update_absolute(new_progress, total=PROGRESS_BAR_MAX) progress.update_absolute(new_progress, total=PROGRESS_BAR_MAX)
if self.price_extractor:
price = self.price_extractor(response_obj)
if price is not None:
self.extracted_price = price
if status == TaskStatus.COMPLETED: if status == TaskStatus.COMPLETED:
message = "Task completed successfully" message = "Task completed successfully"
if self.result_url_extractor: if self.result_url_extractor:

View File

@ -73,6 +73,7 @@ from comfy_api_nodes.util.validation_utils import (
validate_video_dimensions, validate_video_dimensions,
validate_video_duration, validate_video_duration,
) )
from comfy_api.input_impl import VideoFromFile
from comfy_api.input.basic_types import AudioInput from comfy_api.input.basic_types import AudioInput
from comfy_api.input.video_types import VideoInput from comfy_api.input.video_types import VideoInput
from comfy_api.latest import ComfyExtension, io as comfy_io from comfy_api.latest import ComfyExtension, io as comfy_io
@ -511,7 +512,7 @@ async def execute_video_effect(
image_1: torch.Tensor, image_1: torch.Tensor,
image_2: Optional[torch.Tensor] = None, image_2: Optional[torch.Tensor] = None,
model_mode: Optional[KlingVideoGenMode] = None, model_mode: Optional[KlingVideoGenMode] = None,
) -> comfy_io.NodeOutput: ) -> tuple[VideoFromFile, str, str]:
if dual_character: if dual_character:
request_input_field = KlingDualCharacterEffectInput( request_input_field = KlingDualCharacterEffectInput(
model_name=model_name, model_name=model_name,
@ -562,7 +563,7 @@ async def execute_video_effect(
validate_video_result_response(final_response) validate_video_result_response(final_response)
video = get_video_from_response(final_response) video = get_video_from_response(final_response)
return comfy_io.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration)) return await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration)
async def execute_lipsync( async def execute_lipsync(
@ -1271,7 +1272,7 @@ class KlingDualCharacterVideoEffectNode(comfy_io.ComfyNode):
image_1=image_left, image_1=image_left,
image_2=image_right, image_2=image_right,
) )
return video, duration return comfy_io.NodeOutput(video, duration)
class KlingSingleImageVideoEffectNode(comfy_io.ComfyNode): class KlingSingleImageVideoEffectNode(comfy_io.ComfyNode):
@ -1320,17 +1321,21 @@ class KlingSingleImageVideoEffectNode(comfy_io.ComfyNode):
model_name: KlingSingleImageEffectModelName, model_name: KlingSingleImageEffectModelName,
duration: KlingVideoGenDuration, duration: KlingVideoGenDuration,
) -> comfy_io.NodeOutput: ) -> comfy_io.NodeOutput:
return await execute_video_effect( return comfy_io.NodeOutput(
auth_kwargs={ *(
"auth_token": cls.hidden.auth_token_comfy_org, await execute_video_effect(
"comfy_api_key": cls.hidden.api_key_comfy_org, auth_kwargs={
}, "auth_token": cls.hidden.auth_token_comfy_org,
node_id=cls.hidden.unique_id, "comfy_api_key": cls.hidden.api_key_comfy_org,
dual_character=False, },
effect_scene=effect_scene, node_id=cls.hidden.unique_id,
model_name=model_name, dual_character=False,
duration=duration, effect_scene=effect_scene,
image_1=image, model_name=model_name,
duration=duration,
image_1=image,
)
)
) )

View File

@ -17,6 +17,7 @@ from comfy_api.input_impl.video_types import VideoCodec, VideoContainer, VideoIn
from comfy_api_nodes.apinode_utils import ( from comfy_api_nodes.apinode_utils import (
download_url_to_video_output, download_url_to_video_output,
tensor_to_bytesio, tensor_to_bytesio,
validate_string,
) )
from comfy_api_nodes.apis import pika_defs from comfy_api_nodes.apis import pika_defs
from comfy_api_nodes.apis.client import ( from comfy_api_nodes.apis.client import (
@ -590,6 +591,7 @@ class PikaStartEndFrameNode(comfy_io.ComfyNode):
resolution: str, resolution: str,
duration: int, duration: int,
) -> comfy_io.NodeOutput: ) -> comfy_io.NodeOutput:
validate_string(prompt_text, field_name="prompt_text", min_length=1)
pika_files = [ pika_files = [
("keyFrames", ("image_start.png", tensor_to_bytesio(image_start), "image/png")), ("keyFrames", ("image_start.png", tensor_to_bytesio(image_start), "image/png")),
("keyFrames", ("image_end.png", tensor_to_bytesio(image_end), "image/png")), ("keyFrames", ("image_end.png", tensor_to_bytesio(image_end), "image/png")),