Skip to content

Commit f9ec85f

Browse files
bigcat88bigcat88
authored
feat(api-nodes): update xAI Grok nodes (Comfy-Org#13140)
1 parent 2d5fd3f commit f9ec85f

2 files changed

Lines changed: 260 additions & 1 deletion

File tree

‎comfy_api_nodes/apis/grok.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,21 @@ class ImageEditRequest(BaseModel):
2929
class VideoGenerationRequest(BaseModel):
3030
model: str = Field(...)
3131
prompt: str = Field(...)
32-
image: InputUrlObject | None = Field(...)
32+
image: InputUrlObject | None = Field(None)
33+
reference_images: list[InputUrlObject] | None = Field(None)
3334
duration: int = Field(...)
3435
aspect_ratio: str | None = Field(...)
3536
resolution: str = Field(...)
3637
seed: int = Field(...)
3738

3839

40+
class VideoExtensionRequest(BaseModel):
41+
prompt: str = Field(...)
42+
video: InputUrlObject = Field(...)
43+
duration: int = Field(default=6)
44+
model: str | None = Field(default=None)
45+
46+
3947
class VideoEditRequest(BaseModel):
4048
model: str = Field(...)
4149
prompt: str = Field(...)

‎comfy_api_nodes/nodes_grok.py‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
ImageGenerationResponse,
99
InputUrlObject,
1010
VideoEditRequest,
11+
VideoExtensionRequest,
1112
VideoGenerationRequest,
1213
VideoGenerationResponse,
1314
VideoStatusResponse,
@@ -21,6 +22,7 @@
2122
poll_op,
2223
sync_op,
2324
tensor_to_base64_string,
25+
upload_images_to_comfyapi,
2426
upload_video_to_comfyapi,
2527
validate_string,
2628
validate_video_duration,
@@ -33,6 +35,13 @@ def _extract_grok_price(response) -> float | None:
3335
return None
3436

3537

38+
def _extract_grok_video_price(response) -> float | None:
39+
price = _extract_grok_price(response)
40+
if price is not None:
41+
return price * 1.43
42+
return None
43+
44+
3645
class GrokImageNode(IO.ComfyNode):
3746

3847
@classmethod
@@ -354,6 +363,8 @@ async def execute(
354363
seed: int,
355364
image: Input.Image | None = None,
356365
) -> IO.NodeOutput:
366+
if model == "grok-imagine-video-beta":
367+
model = "grok-imagine-video"
357368
image_url = None
358369
if image is not None:
359370
if get_number_of_images(image) != 1:
@@ -462,14 +473,254 @@ async def execute(
462473
return IO.NodeOutput(await download_url_to_video_output(response.video.url))
463474

464475

476+
class GrokVideoReferenceNode(IO.ComfyNode):
477+
478+
@classmethod
479+
def define_schema(cls):
480+
return IO.Schema(
481+
node_id="GrokVideoReferenceNode",
482+
display_name="Grok Reference-to-Video",
483+
category="api node/video/Grok",
484+
description="Generate video guided by reference images as style and content references.",
485+
inputs=[
486+
IO.String.Input(
487+
"prompt",
488+
multiline=True,
489+
tooltip="Text description of the desired video.",
490+
),
491+
IO.DynamicCombo.Input(
492+
"model",
493+
options=[
494+
IO.DynamicCombo.Option(
495+
"grok-imagine-video",
496+
[
497+
IO.Autogrow.Input(
498+
"reference_images",
499+
template=IO.Autogrow.TemplatePrefix(
500+
IO.Image.Input("image"),
501+
prefix="reference_",
502+
min=1,
503+
max=7,
504+
),
505+
tooltip="Up to 7 reference images to guide the video generation.",
506+
),
507+
IO.Combo.Input(
508+
"resolution",
509+
options=["480p", "720p"],
510+
tooltip="The resolution of the output video.",
511+
),
512+
IO.Combo.Input(
513+
"aspect_ratio",
514+
options=["16:9", "4:3", "3:2", "1:1", "2:3", "3:4", "9:16"],
515+
tooltip="The aspect ratio of the output video.",
516+
),
517+
IO.Int.Input(
518+
"duration",
519+
default=6,
520+
min=2,
521+
max=10,
522+
step=1,
523+
tooltip="The duration of the output video in seconds.",
524+
display_mode=IO.NumberDisplay.slider,
525+
),
526+
],
527+
),
528+
],
529+
tooltip="The model to use for video generation.",
530+
),
531+
IO.Int.Input(
532+
"seed",
533+
default=0,
534+
min=0,
535+
max=2147483647,
536+
step=1,
537+
display_mode=IO.NumberDisplay.number,
538+
control_after_generate=True,
539+
tooltip="Seed to determine if node should re-run; "
540+
"actual results are nondeterministic regardless of seed.",
541+
),
542+
],
543+
outputs=[
544+
IO.Video.Output(),
545+
],
546+
hidden=[
547+
IO.Hidden.auth_token_comfy_org,
548+
IO.Hidden.api_key_comfy_org,
549+
IO.Hidden.unique_id,
550+
],
551+
is_api_node=True,
552+
price_badge=IO.PriceBadge(
553+
depends_on=IO.PriceBadgeDepends(
554+
widgets=["model.duration", "model.resolution"],
555+
input_groups=["model.reference_images"],
556+
),
557+
expr="""
558+
(
559+
$res := $lookup(widgets, "model.resolution");
560+
$dur := $lookup(widgets, "model.duration");
561+
$refs := inputGroups["model.reference_images"];
562+
$rate := $res = "720p" ? 0.07 : 0.05;
563+
$price := ($rate * $dur + 0.002 * $refs) * 1.43;
564+
{"type":"usd","usd": $price}
565+
)
566+
""",
567+
),
568+
)
569+
570+
@classmethod
571+
async def execute(
572+
cls,
573+
prompt: str,
574+
model: dict,
575+
seed: int,
576+
) -> IO.NodeOutput:
577+
validate_string(prompt, strip_whitespace=True, min_length=1)
578+
ref_image_urls = await upload_images_to_comfyapi(
579+
cls,
580+
list(model["reference_images"].values()),
581+
mime_type="image/png",
582+
wait_label="Uploading base images",
583+
max_images=7,
584+
)
585+
initial_response = await sync_op(
586+
cls,
587+
ApiEndpoint(path="/proxy/xai/v1/videos/generations", method="POST"),
588+
data=VideoGenerationRequest(
589+
model=model["model"],
590+
reference_images=[InputUrlObject(url=i) for i in ref_image_urls],
591+
prompt=prompt,
592+
resolution=model["resolution"],
593+
duration=model["duration"],
594+
aspect_ratio=model["aspect_ratio"],
595+
seed=seed,
596+
),
597+
response_model=VideoGenerationResponse,
598+
)
599+
response = await poll_op(
600+
cls,
601+
ApiEndpoint(path=f"/proxy/xai/v1/videos/{initial_response.request_id}"),
602+
status_extractor=lambda r: r.status if r.status is not None else "complete",
603+
response_model=VideoStatusResponse,
604+
price_extractor=_extract_grok_video_price,
605+
)
606+
return IO.NodeOutput(await download_url_to_video_output(response.video.url))
607+
608+
609+
class GrokVideoExtendNode(IO.ComfyNode):
610+
611+
@classmethod
612+
def define_schema(cls):
613+
return IO.Schema(
614+
node_id="GrokVideoExtendNode",
615+
display_name="Grok Video Extend",
616+
category="api node/video/Grok",
617+
description="Extend an existing video with a seamless continuation based on a text prompt.",
618+
inputs=[
619+
IO.String.Input(
620+
"prompt",
621+
multiline=True,
622+
tooltip="Text description of what should happen next in the video.",
623+
),
624+
IO.Video.Input("video", tooltip="Source video to extend. MP4 format, 2-15 seconds."),
625+
IO.DynamicCombo.Input(
626+
"model",
627+
options=[
628+
IO.DynamicCombo.Option(
629+
"grok-imagine-video",
630+
[
631+
IO.Int.Input(
632+
"duration",
633+
default=8,
634+
min=2,
635+
max=10,
636+
step=1,
637+
tooltip="Length of the extension in seconds.",
638+
display_mode=IO.NumberDisplay.slider,
639+
),
640+
],
641+
),
642+
],
643+
tooltip="The model to use for video extension.",
644+
),
645+
IO.Int.Input(
646+
"seed",
647+
default=0,
648+
min=0,
649+
max=2147483647,
650+
step=1,
651+
display_mode=IO.NumberDisplay.number,
652+
control_after_generate=True,
653+
tooltip="Seed to determine if node should re-run; "
654+
"actual results are nondeterministic regardless of seed.",
655+
),
656+
],
657+
outputs=[
658+
IO.Video.Output(),
659+
],
660+
hidden=[
661+
IO.Hidden.auth_token_comfy_org,
662+
IO.Hidden.api_key_comfy_org,
663+
IO.Hidden.unique_id,
664+
],
665+
is_api_node=True,
666+
price_badge=IO.PriceBadge(
667+
depends_on=IO.PriceBadgeDepends(widgets=["model.duration"]),
668+
expr="""
669+
(
670+
$dur := $lookup(widgets, "model.duration");
671+
{
672+
"type": "range_usd",
673+
"min_usd": (0.02 + 0.05 * $dur) * 1.43,
674+
"max_usd": (0.15 + 0.05 * $dur) * 1.43
675+
}
676+
)
677+
""",
678+
),
679+
)
680+
681+
@classmethod
682+
async def execute(
683+
cls,
684+
prompt: str,
685+
video: Input.Video,
686+
model: dict,
687+
seed: int,
688+
) -> IO.NodeOutput:
689+
validate_string(prompt, strip_whitespace=True, min_length=1)
690+
validate_video_duration(video, min_duration=2, max_duration=15)
691+
video_size = get_fs_object_size(video.get_stream_source())
692+
if video_size > 50 * 1024 * 1024:
693+
raise ValueError(f"Video size ({video_size / 1024 / 1024:.1f}MB) exceeds 50MB limit.")
694+
initial_response = await sync_op(
695+
cls,
696+
ApiEndpoint(path="/proxy/xai/v1/videos/extensions", method="POST"),
697+
data=VideoExtensionRequest(
698+
prompt=prompt,
699+
video=InputUrlObject(url=await upload_video_to_comfyapi(cls, video)),
700+
duration=model["duration"],
701+
),
702+
response_model=VideoGenerationResponse,
703+
)
704+
response = await poll_op(
705+
cls,
706+
ApiEndpoint(path=f"/proxy/xai/v1/videos/{initial_response.request_id}"),
707+
status_extractor=lambda r: r.status if r.status is not None else "complete",
708+
response_model=VideoStatusResponse,
709+
price_extractor=_extract_grok_video_price,
710+
)
711+
return IO.NodeOutput(await download_url_to_video_output(response.video.url))
712+
713+
465714
class GrokExtension(ComfyExtension):
466715
@override
467716
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
468717
return [
469718
GrokImageNode,
470719
GrokImageEditNode,
471720
GrokVideoNode,
721+
GrokVideoReferenceNode,
472722
GrokVideoEditNode,
723+
GrokVideoExtendNode,
473724
]
474725

475726

0 commit comments

Comments
 (0)