Skip to content

Commit 1094578

Browse files
bigcat88bigcat88
committed
feat(api-nodes): add Sora2 API node
1 parent 8aea746 commit 1094578

File tree

3 files changed

+192
-5
lines changed

3 files changed

+192
-5
lines changed

‎comfy_api_nodes/apinode_utils.py‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
UploadResponse,
1919
)
2020
from server import PromptServer
21-
21+
from comfy.cli_args import args
2222

2323
import numpy as np
2424
from PIL import Image
@@ -30,7 +30,9 @@
3030
import av
3131

3232

33-
async def download_url_to_video_output(video_url: str, timeout: int = None) -> VideoFromFile:
33+
async def download_url_to_video_output(
34+
video_url: str, timeout: int = None, auth_kwargs: Optional[dict[str, str]] = None
35+
) -> VideoFromFile:
3436
"""Downloads a video from a URL and returns a `VIDEO` output.
3537
3638
Args:
@@ -39,7 +41,7 @@ async def download_url_to_video_output(video_url: str, timeout: int = None) -> V
3941
Returns:
4042
A Comfy node `VIDEO` output.
4143
"""
42-
video_io = await download_url_to_bytesio(video_url, timeout)
44+
video_io = await download_url_to_bytesio(video_url, timeout, auth_kwargs=auth_kwargs)
4345
if video_io is None:
4446
error_msg = f"Failed to download video from {video_url}"
4547
logging.error(error_msg)
@@ -164,7 +166,9 @@ def mimetype_to_extension(mime_type: str) -> str:
164166
return mime_type.split("/")[-1].lower()
165167

166168

167-
async def download_url_to_bytesio(url: str, timeout: int = None) -> BytesIO:
169+
async def download_url_to_bytesio(
170+
url: str, timeout: int = None, auth_kwargs: Optional[dict[str, str]] = None
171+
) -> BytesIO:
168172
"""Downloads content from a URL using requests and returns it as BytesIO.
169173
170174
Args:
@@ -174,9 +178,18 @@ async def download_url_to_bytesio(url: str, timeout: int = None) -> BytesIO:
174178
Returns:
175179
BytesIO object containing the downloaded content.
176180
"""
181+
headers = {}
182+
if url.startswith("/proxy/"):
183+
url = str(args.comfy_api_base).rstrip("/") + url
184+
auth_token = auth_kwargs.get("auth_token")
185+
comfy_api_key = auth_kwargs.get("comfy_api_key")
186+
if auth_token:
187+
headers["Authorization"] = f"Bearer {auth_token}"
188+
elif comfy_api_key:
189+
headers["X-API-KEY"] = comfy_api_key
177190
timeout_cfg = aiohttp.ClientTimeout(total=timeout) if timeout else None
178191
async with aiohttp.ClientSession(timeout=timeout_cfg) as session:
179-
async with session.get(url) as resp:
192+
async with session.get(url, headers=headers) as resp:
180193
resp.raise_for_status() # Raises HTTPError for bad responses (4XX or 5XX)
181194
return BytesIO(await resp.read())
182195

‎comfy_api_nodes/nodes_sora.py‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
from typing import Optional
2+
from typing_extensions import override
3+
4+
import torch
5+
from pydantic import BaseModel, Field
6+
from comfy_api.latest import ComfyExtension, io as comfy_io
7+
from comfy_api_nodes.apis.client import (
8+
ApiEndpoint,
9+
HttpMethod,
10+
SynchronousOperation,
11+
PollingOperation,
12+
EmptyRequest,
13+
)
14+
from comfy_api_nodes.util.validation_utils import get_number_of_images
15+
16+
from comfy_api_nodes.apinode_utils import (
17+
download_url_to_video_output,
18+
tensor_to_bytesio,
19+
)
20+
21+
class Sora2GenerationRequest(BaseModel):
22+
prompt: str = Field(...)
23+
model: str = Field(...)
24+
seconds: str = Field(...)
25+
size: str = Field(...)
26+
27+
28+
class Sora2GenerationResponse(BaseModel):
29+
id: str = Field(...)
30+
error: Optional[dict] = Field(None)
31+
status: Optional[str] = Field(None)
32+
33+
34+
class OpenAIVideoSora2(comfy_io.ComfyNode):
35+
@classmethod
36+
def define_schema(cls):
37+
return comfy_io.Schema(
38+
node_id="OpenAIVideoSora2",
39+
display_name="OpenAI Sora - Video",
40+
category="api node/video/Sora",
41+
description="OpenAI video and audio generation.",
42+
inputs=[
43+
comfy_io.Combo.Input(
44+
"model",
45+
options=["sora-2", "sora-2-pro"],
46+
default="sora-2",
47+
),
48+
comfy_io.String.Input(
49+
"prompt",
50+
multiline=True,
51+
default="",
52+
tooltip="Guiding text; may be empty if an input image is present.",
53+
),
54+
comfy_io.Combo.Input(
55+
"size",
56+
options=[
57+
"720x1280",
58+
"1280x720",
59+
"1024x1792",
60+
"1792x1024",
61+
],
62+
default="1280x720",
63+
),
64+
comfy_io.Combo.Input(
65+
"duration",
66+
options=[4, 8, 12],
67+
default=8,
68+
),
69+
comfy_io.Image.Input(
70+
"image",
71+
optional=True,
72+
),
73+
comfy_io.Int.Input(
74+
"seed",
75+
default=0,
76+
min=0,
77+
max=2147483647,
78+
step=1,
79+
display_mode=comfy_io.NumberDisplay.number,
80+
control_after_generate=True,
81+
optional=True,
82+
tooltip="Seed to determine if node should re-run; "
83+
"actual results are nondeterministic regardless of seed.",
84+
),
85+
],
86+
outputs=[
87+
comfy_io.Video.Output(),
88+
],
89+
hidden=[
90+
comfy_io.Hidden.auth_token_comfy_org,
91+
comfy_io.Hidden.api_key_comfy_org,
92+
comfy_io.Hidden.unique_id,
93+
],
94+
is_api_node=True,
95+
)
96+
97+
@classmethod
98+
async def execute(
99+
cls,
100+
model: str,
101+
prompt: str,
102+
size: str = "1280x720",
103+
duration: int = 8,
104+
seed: int = 0,
105+
image: Optional[torch.Tensor] = None,
106+
):
107+
if model == "sora-2" and size not in ("720x1280", "1280x720"):
108+
raise ValueError("Invalid size for sora-2 model, only 720x1280 and 1280x720 are supported.")
109+
files_input = None
110+
if image is not None:
111+
if get_number_of_images(image) != 1:
112+
raise ValueError("Currently only one input image is supported.")
113+
files_input = {"input_reference": tensor_to_bytesio(image)}
114+
auth = {
115+
"auth_token": cls.hidden.auth_token_comfy_org,
116+
"comfy_api_key": cls.hidden.api_key_comfy_org,
117+
}
118+
payload = Sora2GenerationRequest(
119+
model=model,
120+
prompt=prompt,
121+
seconds=str(duration),
122+
size=size,
123+
)
124+
initial_operation = SynchronousOperation(
125+
endpoint=ApiEndpoint(
126+
path="/proxy/openai/v1/videos",
127+
method=HttpMethod.POST,
128+
request_model=Sora2GenerationRequest,
129+
response_model=Sora2GenerationResponse
130+
),
131+
request=payload,
132+
files=files_input,
133+
auth_kwargs=auth,
134+
)
135+
initial_response = await initial_operation.execute()
136+
if initial_response.error:
137+
raise Exception(initial_response.error.message)
138+
139+
model_time_multiplier = 1 if model == "sora-2" else 2
140+
poll_operation = PollingOperation(
141+
poll_endpoint=ApiEndpoint(
142+
path=f"/proxy/openai/v1/videos/{initial_response.id}",
143+
method=HttpMethod.GET,
144+
request_model=EmptyRequest,
145+
response_model=Sora2GenerationResponse
146+
),
147+
completed_statuses=["completed"],
148+
failed_statuses=["failed"],
149+
status_extractor=lambda x: x.status,
150+
auth_kwargs=auth,
151+
poll_interval=5.0,
152+
node_id=cls.hidden.unique_id,
153+
estimated_duration=45 * (duration / 4) * model_time_multiplier,
154+
)
155+
await poll_operation.execute()
156+
return comfy_io.NodeOutput(
157+
await download_url_to_video_output(
158+
f"/proxy/openai/v1/videos/{initial_response.id}/content",
159+
auth_kwargs=auth,
160+
)
161+
)
162+
163+
164+
class OpenAISoraExtension(ComfyExtension):
165+
@override
166+
async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
167+
return [
168+
OpenAIVideoSora2,
169+
]
170+
171+
172+
async def comfy_entrypoint() -> OpenAISoraExtension:
173+
return OpenAISoraExtension()

‎nodes.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2357,6 +2357,7 @@ async def init_builtin_api_nodes():
23572357
"nodes_stability.py",
23582358
"nodes_pika.py",
23592359
"nodes_runway.py",
2360+
"nodes_sora.py",
23602361
"nodes_tripo.py",
23612362
"nodes_moonvalley.py",
23622363
"nodes_rodin.py",

0 commit comments

Comments
 (0)