848 lines
31 KiB
Python
848 lines
31 KiB
Python
"""
|
|
ComfyUI custom nodes for interacting with a Sora-compatible API.
|
|
|
|
This module avoids external dependencies so it can run inside a vanilla
|
|
ComfyUI installation. Configure the Sora API key through the node input,
|
|
the `SORA_API_KEY` environment variable, or `OPENAI_API_KEY`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional, Union, Tuple
|
|
|
|
# Try to import PIL and numpy for image handling (optional dependencies)
|
|
try:
|
|
from PIL import Image
|
|
PIL_AVAILABLE = True
|
|
except ImportError:
|
|
PIL_AVAILABLE = False
|
|
Image = None
|
|
|
|
try:
|
|
import numpy as np
|
|
NUMPY_AVAILABLE = True
|
|
except ImportError:
|
|
NUMPY_AVAILABLE = False
|
|
np = None
|
|
|
|
|
|
DEFAULT_SORA_BASE_URL = "https://api.openai.com/v1/sora"
|
|
|
|
|
|
class SoraAPIError(RuntimeError):
|
|
"""Raised when the upstream Sora service reports an error."""
|
|
|
|
|
|
@dataclass
|
|
class SoraJobResult:
|
|
job_id: str
|
|
status: str
|
|
response_json: Dict[str, Any]
|
|
downloaded_file: Optional[str]
|
|
|
|
|
|
class SoraClient:
|
|
"""Minimal HTTP client for the Sora REST endpoints."""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str,
|
|
base_url: str = DEFAULT_SORA_BASE_URL,
|
|
endpoint: str = "/videos",
|
|
timeout: float = 60.0,
|
|
) -> None:
|
|
if not api_key:
|
|
raise ValueError("Missing Sora API key")
|
|
self.api_key = api_key
|
|
self.base_url = base_url.rstrip("/")
|
|
self.endpoint = endpoint if endpoint.startswith("/") else f"/{endpoint}"
|
|
self.timeout = timeout
|
|
|
|
def create_job(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
return self._request("POST", self.endpoint, payload)
|
|
|
|
def create_job_with_files(
|
|
self,
|
|
payload: Dict[str, Any],
|
|
files: Optional[Dict[str, Tuple[bytes, str]]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
创建任务并上传文件（multipart/form-data）。
|
|
|
|
Args:
|
|
payload: JSON 数据
|
|
files: 文件字典，格式为 {field_name: (file_bytes, mime_type)}
|
|
"""
|
|
return self._request_multipart("POST", self.endpoint, payload, files)
|
|
|
|
def fetch_job(self, job_id: str) -> Dict[str, Any]:
|
|
return self._request("GET", f"{self.endpoint}/{job_id}")
|
|
|
|
def download_asset(self, url: str, target_path: str) -> str:
|
|
request = urllib.request.Request(url, method="GET")
|
|
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
|
data = response.read()
|
|
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
|
with open(target_path, "wb") as file:
|
|
file.write(data)
|
|
return target_path
|
|
|
|
def _request(self, method: str, path: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
url = f"{self.base_url}{path}"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data: Optional[bytes] = None
|
|
if payload is not None:
|
|
data = json.dumps(payload).encode("utf-8")
|
|
request = urllib.request.Request(url, data=data, headers=headers, method=method)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
|
body = response.read().decode("utf-8")
|
|
except urllib.error.HTTPError as http_error:
|
|
error_body = http_error.read().decode("utf-8")
|
|
raise SoraAPIError(f"Sora API error ({http_error.code}): {error_body}") from http_error
|
|
except urllib.error.URLError as url_error:
|
|
raise SoraAPIError(f"Sora request failed: {url_error.reason}") from url_error
|
|
return json.loads(body or "{}")
|
|
|
|
def _request_multipart(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
payload: Dict[str, Any],
|
|
files: Optional[Dict[str, Tuple[bytes, str]]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
发送 multipart/form-data 请求（用于文件上传）。
|
|
"""
|
|
import uuid
|
|
|
|
url = f"{self.base_url}{path}"
|
|
boundary = f"----WebKitFormBoundary{uuid.uuid4().hex}"
|
|
|
|
# 构建 multipart 数据
|
|
parts = []
|
|
|
|
# 添加 JSON 字段
|
|
for key, value in payload.items():
|
|
if value is not None:
|
|
parts.append(f'--{boundary}\r\n'.encode())
|
|
parts.append(f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode())
|
|
if isinstance(value, (dict, list)):
|
|
parts.append(json.dumps(value).encode("utf-8"))
|
|
else:
|
|
parts.append(str(value).encode("utf-8"))
|
|
parts.append(b'\r\n')
|
|
|
|
# 添加文件字段
|
|
if files:
|
|
for field_name, (file_bytes, mime_type) in files.items():
|
|
parts.append(f'--{boundary}\r\n'.encode())
|
|
parts.append(
|
|
f'Content-Disposition: form-data; name="{field_name}"; filename="file"\r\n'.encode()
|
|
)
|
|
parts.append(f'Content-Type: {mime_type}\r\n\r\n'.encode())
|
|
parts.append(file_bytes)
|
|
parts.append(b'\r\n')
|
|
|
|
parts.append(f'--{boundary}--\r\n'.encode())
|
|
data = b''.join(parts)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": f"multipart/form-data; boundary={boundary}",
|
|
}
|
|
|
|
request = urllib.request.Request(url, data=data, headers=headers, method=method)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
|
body = response.read().decode("utf-8")
|
|
except urllib.error.HTTPError as http_error:
|
|
error_body = http_error.read().decode("utf-8")
|
|
raise SoraAPIError(f"Sora API error ({http_error.code}): {error_body}") from http_error
|
|
except urllib.error.URLError as url_error:
|
|
raise SoraAPIError(f"Sora request failed: {url_error.reason}") from url_error
|
|
return json.loads(body or "{}")
|
|
|
|
|
|
def _resolve_api_key(explicit_key: Optional[str]) -> str:
|
|
if explicit_key:
|
|
return explicit_key.strip()
|
|
env_key = os.getenv("SORA_API_KEY") or os.getenv("OPENAI_API_KEY")
|
|
if not env_key:
|
|
raise ValueError("Provide a Sora API key via the node input or SORA_API_KEY/OPENAI_API_KEY env vars.")
|
|
return env_key.strip()
|
|
|
|
|
|
def _process_image_input(image_input: Any) -> Tuple[bytes, str]:
|
|
"""
|
|
处理 ComfyUI 的图像输入，转换为字节数据。
|
|
|
|
支持：
|
|
- PIL.Image 对象
|
|
- numpy 数组
|
|
- 文件路径字符串
|
|
- base64 编码的字符串
|
|
|
|
返回: (image_bytes, mime_type)
|
|
"""
|
|
# 如果是字符串，可能是文件路径或 base64
|
|
if isinstance(image_input, str):
|
|
# 检查是否是 base64
|
|
if image_input.startswith("data:image") or (len(image_input) > 100 and not os.path.exists(image_input)):
|
|
# 可能是 base64，尝试解码
|
|
if "," in image_input:
|
|
# data:image/png;base64,xxx 格式
|
|
header, data = image_input.split(",", 1)
|
|
image_bytes = base64.b64decode(data)
|
|
mime_type = "image/png" # 默认
|
|
if "image/" in header:
|
|
mime_type = header.split("image/")[1].split(";")[0]
|
|
return image_bytes, mime_type
|
|
else:
|
|
# 纯 base64
|
|
image_bytes = base64.b64decode(image_input)
|
|
return image_bytes, "image/png"
|
|
elif os.path.exists(image_input):
|
|
# 文件路径
|
|
with open(image_input, "rb") as f:
|
|
image_bytes = f.read()
|
|
ext = os.path.splitext(image_input)[1].lower()
|
|
mime_types = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp"}
|
|
mime_type = mime_types.get(ext, "image/png")
|
|
return image_bytes, mime_type
|
|
|
|
# 如果是 PIL Image
|
|
if PIL_AVAILABLE and isinstance(image_input, Image.Image):
|
|
buffer = io.BytesIO()
|
|
# 转换为 RGB 如果必要
|
|
if image_input.mode != "RGB":
|
|
image_input = image_input.convert("RGB")
|
|
image_input.save(buffer, format="PNG")
|
|
return buffer.getvalue(), "image/png"
|
|
|
|
# 如果是 numpy 数组
|
|
if NUMPY_AVAILABLE and isinstance(image_input, np.ndarray):
|
|
# 处理不同的数组形状
|
|
img_array = image_input
|
|
if len(img_array.shape) == 4: # [batch, H, W, C]
|
|
img_array = img_array[0]
|
|
elif len(img_array.shape) == 3 and img_array.shape[0] < img_array.shape[2]: # [C, H, W]
|
|
img_array = np.transpose(img_array, (1, 2, 0))
|
|
|
|
# 确保值在 0-255 范围
|
|
if img_array.max() <= 1.0:
|
|
img_array = (img_array * 255).astype(np.uint8)
|
|
else:
|
|
img_array = img_array.astype(np.uint8)
|
|
|
|
# 转换为 PIL Image 然后编码
|
|
if PIL_AVAILABLE:
|
|
pil_image = Image.fromarray(img_array)
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
return buffer.getvalue(), "image/png"
|
|
else:
|
|
raise ValueError("PIL is required to process numpy arrays. Install Pillow.")
|
|
|
|
# 如果已经是字节数据
|
|
if isinstance(image_input, bytes):
|
|
return image_input, "image/png"
|
|
|
|
raise ValueError(f"Unsupported image input type: {type(image_input)}")
|
|
|
|
|
|
def _process_video_input(video_input: Any) -> Tuple[bytes, str]:
|
|
"""
|
|
处理视频输入，转换为字节数据。
|
|
|
|
支持：
|
|
- 文件路径字符串
|
|
- base64 编码的字符串
|
|
- 字节数据
|
|
|
|
返回: (video_bytes, mime_type)
|
|
"""
|
|
if isinstance(video_input, str):
|
|
# 检查是否是 base64
|
|
if video_input.startswith("data:video") or (len(video_input) > 100 and not os.path.exists(video_input)):
|
|
if "," in video_input:
|
|
header, data = video_input.split(",", 1)
|
|
video_bytes = base64.b64decode(data)
|
|
mime_type = "video/mp4"
|
|
if "video/" in header:
|
|
mime_type = header.split("video/")[1].split(";")[0]
|
|
return video_bytes, mime_type
|
|
else:
|
|
video_bytes = base64.b64decode(video_input)
|
|
return video_bytes, "video/mp4"
|
|
elif os.path.exists(video_input):
|
|
with open(video_input, "rb") as f:
|
|
video_bytes = f.read()
|
|
ext = os.path.splitext(video_input)[1].lower()
|
|
mime_types = {".mp4": "video/mp4", ".mov": "video/quicktime", ".avi": "video/x-msvideo", ".webm": "video/webm"}
|
|
mime_type = mime_types.get(ext, "video/mp4")
|
|
return video_bytes, mime_type
|
|
|
|
if isinstance(video_input, bytes):
|
|
return video_input, "video/mp4"
|
|
|
|
raise ValueError(f"Unsupported video input type: {type(video_input)}")
|
|
|
|
|
|
class SoraTextToVideo:
|
|
"""
|
|
Submits a job to the Sora API and optionally polls until a video is ready.
|
|
The node returns diagnostic metadata along with the downloaded file path (if requested).
|
|
"""
|
|
|
|
CATEGORY = "Sora"
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("job_id", "status", "result")
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": "A cinematic drone shot of waves crashing."}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"multiline": False, "default": ""}),
|
|
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"duration_seconds": ("FLOAT", {"default": 5.0, "min": 1.0, "step": 0.5}),
|
|
"aspect_ratio": ("STRING", {"default": "16:9"}),
|
|
"seed": ("STRING", {"default": ""}),
|
|
"fps": ("INT", {"default": 24, "min": 1}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5, "min": 0.0, "step": 0.1}),
|
|
"base_url": ("STRING", {"default": DEFAULT_SORA_BASE_URL}),
|
|
"endpoint": ("STRING", {"default": "/videos"}),
|
|
"metadata": ("STRING", {"multiline": True, "default": ""}),
|
|
"webhook_url": ("STRING", {"default": ""}),
|
|
"wait_for_result": ("BOOLEAN", {"default": True}),
|
|
"poll_interval": ("FLOAT", {"default": 3.0, "min": 1.0}),
|
|
"max_wait_seconds": ("FLOAT", {"default": 120.0, "min": 5.0}),
|
|
"download_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
prompt: str,
|
|
api_key: str = "",
|
|
negative_prompt: str = "",
|
|
duration_seconds: float = 5.0,
|
|
aspect_ratio: str = "16:9",
|
|
seed: str = "",
|
|
fps: int = 24,
|
|
guidance_scale: float = 7.5,
|
|
base_url: str = DEFAULT_SORA_BASE_URL,
|
|
endpoint: str = "/videos",
|
|
metadata: str = "",
|
|
webhook_url: str = "",
|
|
wait_for_result: bool = True,
|
|
poll_interval: float = 3.0,
|
|
max_wait_seconds: float = 120.0,
|
|
download_path: str = "",
|
|
) -> tuple[str, str, str]:
|
|
client = SoraClient(
|
|
api_key=_resolve_api_key(api_key),
|
|
base_url=base_url,
|
|
endpoint=endpoint,
|
|
)
|
|
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt or None,
|
|
"duration_seconds": duration_seconds,
|
|
"aspect_ratio": aspect_ratio,
|
|
"seed": seed or None,
|
|
"fps": fps,
|
|
"guidance_scale": guidance_scale,
|
|
}
|
|
if metadata:
|
|
try:
|
|
payload["metadata"] = json.loads(metadata)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Metadata must be valid JSON: {error.msg}") from error
|
|
if webhook_url:
|
|
payload["webhook_url"] = webhook_url
|
|
|
|
response = client.create_job(payload)
|
|
job_id = str(response.get("id") or response.get("job_id") or "")
|
|
status = response.get("status", "submitted")
|
|
job_result = SoraJobResult(job_id=job_id, status=status, response_json=response, downloaded_file=None)
|
|
|
|
if wait_for_result and job_id:
|
|
job_result = self._poll_until_complete(
|
|
client=client,
|
|
job_id=job_id,
|
|
poll_interval=poll_interval,
|
|
max_wait_seconds=max_wait_seconds,
|
|
download_path=download_path,
|
|
)
|
|
elif wait_for_result and not job_id:
|
|
raise SoraAPIError("Sora API response did not include a job id.")
|
|
|
|
result_payload = {
|
|
"job_id": job_result.job_id,
|
|
"status": job_result.status,
|
|
"downloaded_file": job_result.downloaded_file,
|
|
"raw_response": job_result.response_json,
|
|
}
|
|
return (
|
|
job_result.job_id or "",
|
|
job_result.status,
|
|
json.dumps(result_payload, ensure_ascii=False),
|
|
)
|
|
|
|
def _poll_until_complete(
|
|
self,
|
|
client: SoraClient,
|
|
job_id: str,
|
|
poll_interval: float,
|
|
max_wait_seconds: float,
|
|
download_path: str,
|
|
) -> SoraJobResult:
|
|
deadline = time.time() + max_wait_seconds
|
|
last_response: Dict[str, Any] = {}
|
|
downloaded_file: Optional[str] = None
|
|
|
|
while time.time() < deadline:
|
|
last_response = client.fetch_job(job_id)
|
|
status = str(last_response.get("status", "unknown")).lower()
|
|
if status in {"succeeded", "completed", "ready"}:
|
|
downloaded_file = self._maybe_download(last_response, client, download_path, job_id)
|
|
return SoraJobResult(
|
|
job_id=job_id,
|
|
status=status,
|
|
response_json=last_response,
|
|
downloaded_file=downloaded_file,
|
|
)
|
|
if status in {"failed", "error"}:
|
|
message = last_response.get("error") or last_response
|
|
raise SoraAPIError(f"Sora generation failed: {message}")
|
|
time.sleep(poll_interval)
|
|
|
|
raise TimeoutError(f"Sora job {job_id} did not finish within {max_wait_seconds} seconds.")
|
|
|
|
def _maybe_download(
|
|
self,
|
|
job_response: Dict[str, Any],
|
|
client: SoraClient,
|
|
download_path: str,
|
|
job_id: str,
|
|
) -> Optional[str]:
|
|
if not download_path:
|
|
return None
|
|
|
|
video_url = self._extract_video_url(job_response)
|
|
if not video_url:
|
|
return None
|
|
|
|
target = download_path
|
|
if os.path.isdir(download_path):
|
|
safe_job_id = job_id or "sora"
|
|
filename = f"{safe_job_id}.mp4"
|
|
target = os.path.join(download_path, filename)
|
|
|
|
return client.download_asset(video_url, target)
|
|
|
|
@staticmethod
|
|
def _extract_video_url(job_response: Dict[str, Any]) -> Optional[str]:
|
|
video_url = job_response.get("result_url") or job_response.get("video_url")
|
|
if video_url:
|
|
return str(video_url)
|
|
|
|
outputs = job_response.get("outputs") or job_response.get("videos")
|
|
if isinstance(outputs, list) and outputs:
|
|
first = outputs[0]
|
|
if isinstance(first, dict):
|
|
url = first.get("url") or first.get("download_url")
|
|
if url:
|
|
return str(url)
|
|
return None
|
|
|
|
|
|
class SoraImageToVideo:
|
|
"""
|
|
Image-to-Video node: Generates video from an image input.
|
|
Supports ComfyUI image inputs (PIL Image, numpy arrays, file paths).
|
|
"""
|
|
|
|
CATEGORY = "Sora"
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("job_id", "status", "result")
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # ComfyUI image input
|
|
"prompt": ("STRING", {"multiline": True, "default": "A cinematic video based on this image."}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"multiline": False, "default": ""}),
|
|
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"duration_seconds": ("FLOAT", {"default": 5.0, "min": 1.0, "step": 0.5}),
|
|
"aspect_ratio": ("STRING", {"default": "16:9"}),
|
|
"seed": ("STRING", {"default": ""}),
|
|
"fps": ("INT", {"default": 24, "min": 1}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5, "min": 0.0, "step": 0.1}),
|
|
"base_url": ("STRING", {"default": DEFAULT_SORA_BASE_URL}),
|
|
"endpoint": ("STRING", {"default": "/videos"}),
|
|
"metadata": ("STRING", {"multiline": True, "default": ""}),
|
|
"webhook_url": ("STRING", {"default": ""}),
|
|
"wait_for_result": ("BOOLEAN", {"default": True}),
|
|
"poll_interval": ("FLOAT", {"default": 3.0, "min": 1.0}),
|
|
"max_wait_seconds": ("FLOAT", {"default": 120.0, "min": 5.0}),
|
|
"download_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
image: Any,
|
|
prompt: str,
|
|
api_key: str = "",
|
|
negative_prompt: str = "",
|
|
duration_seconds: float = 5.0,
|
|
aspect_ratio: str = "16:9",
|
|
seed: str = "",
|
|
fps: int = 24,
|
|
guidance_scale: float = 7.5,
|
|
base_url: str = DEFAULT_SORA_BASE_URL,
|
|
endpoint: str = "/videos",
|
|
metadata: str = "",
|
|
webhook_url: str = "",
|
|
wait_for_result: bool = True,
|
|
poll_interval: float = 3.0,
|
|
max_wait_seconds: float = 120.0,
|
|
download_path: str = "",
|
|
) -> tuple[str, str, str]:
|
|
client = SoraClient(
|
|
api_key=_resolve_api_key(api_key),
|
|
base_url=base_url,
|
|
endpoint=endpoint,
|
|
)
|
|
|
|
# 处理图像输入
|
|
try:
|
|
image_bytes, mime_type = _process_image_input(image)
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to process image input: {e}") from e
|
|
|
|
# 构建 payload
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt or None,
|
|
"duration_seconds": duration_seconds,
|
|
"aspect_ratio": aspect_ratio,
|
|
"seed": seed or None,
|
|
"fps": fps,
|
|
"guidance_scale": guidance_scale,
|
|
}
|
|
if metadata:
|
|
try:
|
|
payload["metadata"] = json.loads(metadata)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Metadata must be valid JSON: {error.msg}") from error
|
|
if webhook_url:
|
|
payload["webhook_url"] = webhook_url
|
|
|
|
# 上传图像和创建任务
|
|
files = {"image": (image_bytes, mime_type)}
|
|
response = client.create_job_with_files(payload, files)
|
|
|
|
job_id = str(response.get("id") or response.get("job_id") or "")
|
|
status = response.get("status", "submitted")
|
|
job_result = SoraJobResult(job_id=job_id, status=status, response_json=response, downloaded_file=None)
|
|
|
|
if wait_for_result and job_id:
|
|
job_result = self._poll_until_complete(
|
|
client=client,
|
|
job_id=job_id,
|
|
poll_interval=poll_interval,
|
|
max_wait_seconds=max_wait_seconds,
|
|
download_path=download_path,
|
|
)
|
|
elif wait_for_result and not job_id:
|
|
raise SoraAPIError("Sora API response did not include a job id.")
|
|
|
|
result_payload = {
|
|
"job_id": job_result.job_id,
|
|
"status": job_result.status,
|
|
"downloaded_file": job_result.downloaded_file,
|
|
"raw_response": job_result.response_json,
|
|
}
|
|
return (
|
|
job_result.job_id or "",
|
|
job_result.status,
|
|
json.dumps(result_payload, ensure_ascii=False),
|
|
)
|
|
|
|
def _poll_until_complete(
|
|
self,
|
|
client: SoraClient,
|
|
job_id: str,
|
|
poll_interval: float,
|
|
max_wait_seconds: float,
|
|
download_path: str,
|
|
) -> SoraJobResult:
|
|
deadline = time.time() + max_wait_seconds
|
|
last_response: Dict[str, Any] = {}
|
|
downloaded_file: Optional[str] = None
|
|
|
|
while time.time() < deadline:
|
|
last_response = client.fetch_job(job_id)
|
|
status = str(last_response.get("status", "unknown")).lower()
|
|
if status in {"succeeded", "completed", "ready"}:
|
|
downloaded_file = self._maybe_download(last_response, client, download_path, job_id)
|
|
return SoraJobResult(
|
|
job_id=job_id,
|
|
status=status,
|
|
response_json=last_response,
|
|
downloaded_file=downloaded_file,
|
|
)
|
|
if status in {"failed", "error"}:
|
|
message = last_response.get("error") or last_response
|
|
raise SoraAPIError(f"Sora generation failed: {message}")
|
|
time.sleep(poll_interval)
|
|
|
|
raise TimeoutError(f"Sora job {job_id} did not finish within {max_wait_seconds} seconds.")
|
|
|
|
def _maybe_download(
|
|
self,
|
|
job_response: Dict[str, Any],
|
|
client: SoraClient,
|
|
download_path: str,
|
|
job_id: str,
|
|
) -> Optional[str]:
|
|
if not download_path:
|
|
return None
|
|
|
|
video_url = self._extract_video_url(job_response)
|
|
if not video_url:
|
|
return None
|
|
|
|
target = download_path
|
|
if os.path.isdir(download_path):
|
|
safe_job_id = job_id or "sora"
|
|
filename = f"{safe_job_id}.mp4"
|
|
target = os.path.join(download_path, filename)
|
|
|
|
return client.download_asset(video_url, target)
|
|
|
|
@staticmethod
|
|
def _extract_video_url(job_response: Dict[str, Any]) -> Optional[str]:
|
|
return SoraTextToVideo._extract_video_url(job_response)
|
|
|
|
|
|
class SoraVideoToVideo:
|
|
"""
|
|
Video-to-Video node: Edits or extends an existing video.
|
|
Supports video extension, inpainting, and style transfer.
|
|
"""
|
|
|
|
CATEGORY = "Sora"
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("job_id", "status", "result")
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
|
|
return {
|
|
"required": {
|
|
"video": ("STRING", {"default": ""}), # Video file path or base64
|
|
"prompt": ("STRING", {"multiline": True, "default": "Extend this video with smooth continuation."}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"multiline": False, "default": ""}),
|
|
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"operation": (["extend", "edit", "inpaint", "style_transfer"], {"default": "extend"}),
|
|
"extension_seconds": ("FLOAT", {"default": 5.0, "min": 1.0, "step": 0.5}),
|
|
"duration_seconds": ("FLOAT", {"default": 5.0, "min": 1.0, "step": 0.5}),
|
|
"aspect_ratio": ("STRING", {"default": "16:9"}),
|
|
"seed": ("STRING", {"default": ""}),
|
|
"fps": ("INT", {"default": 24, "min": 1}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5, "min": 0.0, "step": 0.1}),
|
|
"base_url": ("STRING", {"default": DEFAULT_SORA_BASE_URL}),
|
|
"endpoint": ("STRING", {"default": "/videos"}),
|
|
"metadata": ("STRING", {"multiline": True, "default": ""}),
|
|
"webhook_url": ("STRING", {"default": ""}),
|
|
"wait_for_result": ("BOOLEAN", {"default": True}),
|
|
"poll_interval": ("FLOAT", {"default": 3.0, "min": 1.0}),
|
|
"max_wait_seconds": ("FLOAT", {"default": 120.0, "min": 5.0}),
|
|
"download_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
video: str,
|
|
prompt: str,
|
|
api_key: str = "",
|
|
negative_prompt: str = "",
|
|
operation: str = "extend",
|
|
extension_seconds: float = 5.0,
|
|
duration_seconds: float = 5.0,
|
|
aspect_ratio: str = "16:9",
|
|
seed: str = "",
|
|
fps: int = 24,
|
|
guidance_scale: float = 7.5,
|
|
base_url: str = DEFAULT_SORA_BASE_URL,
|
|
endpoint: str = "/videos",
|
|
metadata: str = "",
|
|
webhook_url: str = "",
|
|
wait_for_result: bool = True,
|
|
poll_interval: float = 3.0,
|
|
max_wait_seconds: float = 120.0,
|
|
download_path: str = "",
|
|
) -> tuple[str, str, str]:
|
|
client = SoraClient(
|
|
api_key=_resolve_api_key(api_key),
|
|
base_url=base_url,
|
|
endpoint=endpoint,
|
|
)
|
|
|
|
# 处理视频输入
|
|
try:
|
|
video_bytes, mime_type = _process_video_input(video)
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to process video input: {e}") from e
|
|
|
|
# 构建 payload
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt or None,
|
|
"operation": operation,
|
|
"duration_seconds": duration_seconds,
|
|
"aspect_ratio": aspect_ratio,
|
|
"seed": seed or None,
|
|
"fps": fps,
|
|
"guidance_scale": guidance_scale,
|
|
}
|
|
|
|
# 根据操作类型添加特定参数
|
|
if operation == "extend":
|
|
payload["extension_seconds"] = extension_seconds
|
|
|
|
if metadata:
|
|
try:
|
|
payload["metadata"] = json.loads(metadata)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Metadata must be valid JSON: {error.msg}") from error
|
|
if webhook_url:
|
|
payload["webhook_url"] = webhook_url
|
|
|
|
# 上传视频和创建任务
|
|
files = {"video": (video_bytes, mime_type)}
|
|
response = client.create_job_with_files(payload, files)
|
|
|
|
job_id = str(response.get("id") or response.get("job_id") or "")
|
|
status = response.get("status", "submitted")
|
|
job_result = SoraJobResult(job_id=job_id, status=status, response_json=response, downloaded_file=None)
|
|
|
|
if wait_for_result and job_id:
|
|
job_result = self._poll_until_complete(
|
|
client=client,
|
|
job_id=job_id,
|
|
poll_interval=poll_interval,
|
|
max_wait_seconds=max_wait_seconds,
|
|
download_path=download_path,
|
|
)
|
|
elif wait_for_result and not job_id:
|
|
raise SoraAPIError("Sora API response did not include a job id.")
|
|
|
|
result_payload = {
|
|
"job_id": job_result.job_id,
|
|
"status": job_result.status,
|
|
"downloaded_file": job_result.downloaded_file,
|
|
"raw_response": job_result.response_json,
|
|
}
|
|
return (
|
|
job_result.job_id or "",
|
|
job_result.status,
|
|
json.dumps(result_payload, ensure_ascii=False),
|
|
)
|
|
|
|
def _poll_until_complete(
|
|
self,
|
|
client: SoraClient,
|
|
job_id: str,
|
|
poll_interval: float,
|
|
max_wait_seconds: float,
|
|
download_path: str,
|
|
) -> SoraJobResult:
|
|
deadline = time.time() + max_wait_seconds
|
|
last_response: Dict[str, Any] = {}
|
|
downloaded_file: Optional[str] = None
|
|
|
|
while time.time() < deadline:
|
|
last_response = client.fetch_job(job_id)
|
|
status = str(last_response.get("status", "unknown")).lower()
|
|
if status in {"succeeded", "completed", "ready"}:
|
|
downloaded_file = self._maybe_download(last_response, client, download_path, job_id)
|
|
return SoraJobResult(
|
|
job_id=job_id,
|
|
status=status,
|
|
response_json=last_response,
|
|
downloaded_file=downloaded_file,
|
|
)
|
|
if status in {"failed", "error"}:
|
|
message = last_response.get("error") or last_response
|
|
raise SoraAPIError(f"Sora generation failed: {message}")
|
|
time.sleep(poll_interval)
|
|
|
|
raise TimeoutError(f"Sora job {job_id} did not finish within {max_wait_seconds} seconds.")
|
|
|
|
def _maybe_download(
|
|
self,
|
|
job_response: Dict[str, Any],
|
|
client: SoraClient,
|
|
download_path: str,
|
|
job_id: str,
|
|
) -> Optional[str]:
|
|
if not download_path:
|
|
return None
|
|
|
|
video_url = self._extract_video_url(job_response)
|
|
if not video_url:
|
|
return None
|
|
|
|
target = download_path
|
|
if os.path.isdir(download_path):
|
|
safe_job_id = job_id or "sora"
|
|
filename = f"{safe_job_id}.mp4"
|
|
target = os.path.join(download_path, filename)
|
|
|
|
return client.download_asset(video_url, target)
|
|
|
|
@staticmethod
|
|
def _extract_video_url(job_response: Dict[str, Any]) -> Optional[str]:
|
|
return SoraTextToVideo._extract_video_url(job_response)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SoraTextToVideo": SoraTextToVideo,
|
|
"SoraImageToVideo": SoraImageToVideo,
|
|
"SoraVideoToVideo": SoraVideoToVideo,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SoraTextToVideo": "Sora Text-To-Video",
|
|
"SoraImageToVideo": "Sora Image-To-Video",
|
|
"SoraVideoToVideo": "Sora Video-To-Video",
|
|
}
|