commit b06aa2bbfc03cddfc6bd990a2a1018c8539f828a Author: maxczc Date: Sat Dec 13 12:21:10 2025 +0800 Initial commit diff --git a/FEATURES.md b/FEATURES.md new file mode 100644 index 0000000..4bc3056 --- /dev/null +++ b/FEATURES.md @@ -0,0 +1,132 @@ +# Sora 自定义节点功能说明 + +## 已实现的功能 + +### ✅ 1. Text-to-Video（文本生成视频） +**节点名称**: `SoraTextToVideo` + +从文本提示生成视频，支持完整的 Sora 参数配置。 + +**主要参数**: +- `prompt`: 视频描述文本 +- `negative_prompt`: 负面提示 +- `duration_seconds`: 视频时长 +- `aspect_ratio`: 宽高比 +- `fps`: 帧率 +- `guidance_scale`: 引导强度 +- `seed`: 随机种子 + +### ✅ 2. Image-to-Video（图像生成视频） +**节点名称**: `SoraImageToVideo` + +从图像生成视频，完全支持 ComfyUI 的图像输入格式。 + +**支持的图像输入格式**: +- ✅ PIL Image 对象 +- ✅ Numpy 数组（自动处理各种形状） +- ✅ 文件路径（PNG, JPG, WebP） +- ✅ Base64 编码字符串 + +**主要参数**: +- `image`: ComfyUI 图像输入（必需） +- `prompt`: 视频描述（必需） +- 其他参数与 Text-to-Video 相同 + +### ✅ 3. Video-to-Video（视频编辑/扩展） +**节点名称**: `SoraVideoToVideo` + +对现有视频进行编辑、扩展或风格转换。 + +**支持的操作类型**: +- `extend`: 扩展视频时长 +- `edit`: 根据提示编辑视频内容 +- `inpaint`: 视频修复/填充 +- `style_transfer`: 风格迁移 + +**主要参数**: +- `video`: 视频文件路径或 base64 字符串（必需） +- `prompt`: 编辑描述（必需） +- `operation`: 操作类型（extend/edit/inpaint/style_transfer） +- `extension_seconds`: 扩展时长（用于 extend 操作） +- 其他参数与其他节点相同 + +## 技术特性 + +### 🔧 文件上传支持 +- 扩展了 `SoraClient` 以支持 `multipart/form-data` 上传 +- 自动处理图像和视频文件的编码和上传 +- 支持多种 MIME 类型 + +### 🖼️ 图像处理 +- 自动格式转换（PIL ↔ Numpy ↔ 文件） +- 自动处理不同的数组形状（batch, channels-first, channels-last） +- 自动归一化（0-1 范围转换为 0-255） +- 自动颜色空间转换（RGBA → RGB） + +### 📹 视频处理 +- 支持多种视频格式（MP4, MOV, AVI, WebM） +- 支持文件路径和 base64 编码 +- 自动 MIME 类型检测 + +### 🔄 轮询机制 +- 所有节点都支持异步任务轮询 +- 可配置的轮询间隔和超时时间 +- 自动状态检测和错误处理 + +### 📥 下载支持 +- 可选的视频下载功能 +- 支持目录或文件路径 +- 自动文件命名（使用 job_id） + +## 使用示例 + +### Text-to-Video +```python +# 在 ComfyUI 中直接使用节点 +# 输入: prompt = "A beautiful sunset over the ocean" +# 输出: job_id, status, result_json +``` + +### Image-to-Video +```python +# 连接 ComfyUI 的图像节点到 SoraImageToVideo +# 输入: image (来自 LoadImage 节点), prompt = "Animate this image" +# 输出: job_id, status, result_json +``` + +### Video-to-Video +```python +# 扩展视频 +# 输入: video = "/path/to/video.mp4", operation = "extend", extension_seconds = 5.0 +# 输出: job_id, status, result_json +``` + +## 兼容性 + +- ✅ 仅使用 Python 标准库（核心功能） +- ✅ 可选依赖：Pillow（PIL）和 numpy（用于更好的图像处理） +- ✅ 完全兼容 ComfyUI 的节点系统 +- ✅ 支持所有 ComfyUI 图像格式 + +## API 兼容性 + +节点设计为与 Sora 兼容的 REST API 配合使用，支持： +- OpenAI Sora API（当可用时） +- 自定义 Sora 兼容后端 +- 可配置的 base_url 和 endpoint + +## 错误处理 + +- ✅ 清晰的错误消息 +- ✅ API 错误自动捕获和报告 +- ✅ 超时保护 +- ✅ 输入验证 + +## 未来可能的扩展 + +- [ ] 批量处理支持 +- [ ] 视频预览功能 +- [ ] 更多视频操作类型 +- [ ] 实时进度回调 +- [ ] 视频质量/分辨率选项 + diff --git a/README.md b/README.md new file mode 100644 index 0000000..2f7825a --- /dev/null +++ b/README.md @@ -0,0 +1,174 @@ +# Sora Custom Node for ComfyUI + +This repository contains a comprehensive set of ComfyUI custom nodes for interacting +with a Sora-compatible REST API. The nodes support the full range of Sora capabilities: + +- **Text-to-Video**: Generate videos from text prompts +- **Image-to-Video**: Generate videos from images +- **Video-to-Video**: Edit, extend, or transform existing videos + +All nodes use only Python's standard library (with optional PIL/numpy support for image processing). + +## Installation + +1. Clone or copy this folder into the `ComfyUI/custom_nodes/` directory: + ```bash + cd ComfyUI/custom_nodes + git clone sora + ``` +2. Restart ComfyUI so it can load the new nodes. +3. **Optional dependencies** (for better image handling): + + - `Pillow` (PIL): For processing PIL Image objects + - `numpy`: For processing numpy array inputs + + These are optional - the nodes will work with file paths and base64 strings without them. + +## Available Nodes + +### 1. Sora Text-To-Video + +Generates videos from text prompts. + +**Inputs:** + +- `prompt` (required): Text description of the video +- `negative_prompt`: What to avoid in the video +- `duration_seconds`: Video length (default: 5.0) +- `aspect_ratio`: Video aspect ratio (default: "16:9") +- `fps`: Frames per second (default: 24) +- `guidance_scale`: Generation guidance (default: 7.5) +- `seed`: Random seed for reproducibility +- And more... + +### 2. Sora Image-To-Video + +Generates videos from images. Accepts ComfyUI image inputs (PIL Image, numpy arrays, or file paths). + +**Inputs:** + +- `image` (required): Image input (from ComfyUI image nodes) +- `prompt` (required): Text description for video generation +- All the same optional parameters as Text-To-Video + +**Image Input Support:** + +- PIL Image objects +- Numpy arrays (various shapes supported) +- File paths (PNG, JPG, WebP) +- Base64 encoded strings + +### 3. Sora Video-To-Video + +Edits or extends existing videos. Supports multiple operations: + +- **extend**: Extend video duration +- **edit**: Edit video content based on prompt +- **inpaint**: Fill masked regions +- **style_transfer**: Apply style to video + +**Inputs:** + +- `video` (required): Video file path or base64 string +- `prompt` (required): Description of desired changes +- `operation`: Type of video operation (default: "extend") +- `extension_seconds`: For extend operation (default: 5.0) +- All the same optional parameters as other nodes + +## Configuration + +All nodes look for an API key using the following priority: + +1. The `api_key` input field on the node. +2. The `SORA_API_KEY` environment variable. +3. The `OPENAI_API_KEY` environment variable. + +You can also configure: + +- `base_url`: defaults to `https://api.openai.com/v1/sora`. +- `endpoint`: defaults to `/videos`, change it if your backend uses another path. +- `download_path`: optional local directory or file path for saving the video. +- `wait_for_result`: Whether to poll until completion (default: True) +- `poll_interval`: Seconds between status checks (default: 3.0) +- `max_wait_seconds`: Maximum wait time (default: 120.0) + +## Outputs + +All nodes return three strings: + +1. **Job ID**: Unique identifier for the generation job +2. **Status**: Current status (`submitted`, `processing`, `succeeded`, `failed`, etc.) +3. **Result JSON**: Contains the raw API response and downloaded file path (if any) + +## Local Testing + +You can test the custom node locally without running ComfyUI: + +### 1. Basic Functionality Test + +Run the test script to verify node functionality: + +```bash +# Run basic tests (no API calls) +python test_sora_node.py + +# Run full unittest suite +python test_sora_node.py --unittest + +# Test with real API (requires API key) +python test_sora_node.py --test-client --api-key YOUR_KEY +``` + +### 2. Test in ComfyUI + +1. Copy this folder to `ComfyUI/custom_nodes/sora/` +2. Start ComfyUI +3. The node should appear in the node menu under "Sora" category +4. You can test it directly in the ComfyUI interface + +### 3. Mock Testing + +The test script uses mocks to test API interactions without making real network calls. This is useful for: + +- Testing error handling +- Testing response parsing +- Validating node logic + +## Features + +### Full Sora Capability Support + +- ✅ Text-to-Video generation +- ✅ Image-to-Video generation +- ✅ Video-to-Video editing and extension +- ✅ Multiple video operations (extend, edit, inpaint, style transfer) + +### ComfyUI Integration + +- Seamless integration with ComfyUI image nodes +- Supports all ComfyUI image formats (PIL, numpy, file paths) +- Automatic format conversion and handling + +### Robust Error Handling + +- Clear error messages for API failures +- Timeout protection +- Automatic retry logic for polling + +### Flexible Configuration + +- Environment variable support for API keys +- Customizable endpoints and base URLs +- Optional result downloading +- Webhook support for async operations + +## Notes + +- Nodes poll the job until completion when `wait_for_result` is enabled. +- If `download_path` is a folder, videos are saved as `.mp4`. +- The nodes automatically search for video URLs in various response formats: + - `result_url` or `video_url` (direct) + - `outputs[0].url` or `outputs[0].download_url` (nested) + - `videos` array +- Image inputs are automatically converted to the format required by the API. +- Video inputs support file paths and base64 encoded strings. diff --git a/__pycache__/sora_node.cpython-312.pyc b/__pycache__/sora_node.cpython-312.pyc new file mode 100644 index 0000000..feffac8 Binary files /dev/null and b/__pycache__/sora_node.cpython-312.pyc differ diff --git a/__pycache__/sora_node.cpython-313.pyc b/__pycache__/sora_node.cpython-313.pyc new file mode 100644 index 0000000..a7de8c2 Binary files /dev/null and b/__pycache__/sora_node.cpython-313.pyc differ diff --git a/quick_test.py b/quick_test.py new file mode 100644 index 0000000..bab6bc6 --- /dev/null +++ b/quick_test.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +""" +快速测试脚本 - 验证节点是否可以正常导入和使用 + +这个脚本可以快速检查： +1. 节点代码是否可以正常导入 +2. 节点类是否可以实例化 +3. 节点映射是否正确 +""" + +def quick_test(): + """快速测试节点基本功能""" + print("🔍 开始快速测试...\n") + + try: + # 1. 测试导入 + print("1. 测试模块导入...") + from sora_node import ( + SoraClient, + SoraTextToVideo, + SoraImageToVideo, + SoraVideoToVideo, + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS + ) + print(" ✓ 模块导入成功\n") + + # 2. 测试节点类 + print("2. 测试节点类...") + # Text-to-Video + node_t2v = SoraTextToVideo() + input_types = node_t2v.INPUT_TYPES() + assert "required" in input_types + assert "prompt" in input_types["required"] + print(" ✓ SoraTextToVideo 节点正常") + + # Image-to-Video + node_i2v = SoraImageToVideo() + input_types_i2v = node_i2v.INPUT_TYPES() + assert "image" in input_types_i2v["required"] + assert "prompt" in input_types_i2v["required"] + print(" ✓ SoraImageToVideo 节点正常") + + # Video-to-Video + node_v2v = SoraVideoToVideo() + input_types_v2v = node_v2v.INPUT_TYPES() + assert "video" in input_types_v2v["required"] + assert "prompt" in input_types_v2v["required"] + print(" ✓ SoraVideoToVideo 节点正常") + print(" ✓ 所有节点类实例化成功\n") + + # 3. 测试节点映射 + print("3. 测试节点映射...") + assert "SoraTextToVideo" in NODE_CLASS_MAPPINGS + assert "SoraImageToVideo" in NODE_CLASS_MAPPINGS + assert "SoraVideoToVideo" in NODE_CLASS_MAPPINGS + assert NODE_CLASS_MAPPINGS["SoraTextToVideo"] == SoraTextToVideo + assert NODE_CLASS_MAPPINGS["SoraImageToVideo"] == SoraImageToVideo + assert NODE_CLASS_MAPPINGS["SoraVideoToVideo"] == SoraVideoToVideo + assert "SoraTextToVideo" in NODE_DISPLAY_NAME_MAPPINGS + assert "SoraImageToVideo" in NODE_DISPLAY_NAME_MAPPINGS + assert "SoraVideoToVideo" in NODE_DISPLAY_NAME_MAPPINGS + print(" ✓ 所有节点映射正确\n") + + # 4. 测试客户端初始化（不需要真实 API） + print("4. 测试客户端类...") + try: + client = SoraClient(api_key="test-key-123") + assert client.api_key == "test-key-123" + print(" ✓ 客户端类正常\n") + except Exception as e: + print(f" ✗ 客户端测试失败: {e}\n") + + print("=" * 50) + print("✅ 所有快速测试通过！") + print("=" * 50) + print("\n节点已准备好，可以：") + print(" 1. 复制到 ComfyUI/custom_nodes/ 目录") + print(" 2. 重启 ComfyUI") + print(" 3. 在节点菜单中找到以下节点：") + print(" - Sora Text-To-Video") + print(" - Sora Image-To-Video") + print(" - Sora Video-To-Video") + return True + + except ImportError as e: + print(f" ✗ 导入失败: {e}") + print("\n请确保 sora_node.py 在同一目录下") + return False + except Exception as e: + print(f" ✗ 测试失败: {e}") + import traceback + traceback.print_exc() + return False + + +if __name__ == "__main__": + quick_test() + diff --git a/sora_node.py b/sora_node.py new file mode 100644 index 0000000..48d9523 --- /dev/null +++ b/sora_node.py @@ -0,0 +1,847 @@ +""" +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", +} diff --git a/test_sora_node.py b/test_sora_node.py new file mode 100644 index 0000000..e70ee63 --- /dev/null +++ b/test_sora_node.py @@ -0,0 +1,248 @@ +""" +本地测试脚本 - 用于测试 Sora 自定义节点 + +使用方法： +1. 直接运行测试（需要 mock API）： + python test_sora_node.py + +2. 测试节点类的基本功能： + python test_sora_node.py --test-class + +3. 测试 API 客户端（需要真实的 API key）： + python test_sora_node.py --test-client --api-key YOUR_KEY +""" + +import json +import os +import sys +import unittest +from unittest.mock import Mock, patch, MagicMock +from urllib.error import HTTPError + +# 导入节点代码 +from sora_node import SoraClient, SoraTextToVideo, SoraAPIError, DEFAULT_SORA_BASE_URL + + +class TestSoraClient(unittest.TestCase): + """测试 SoraClient 类""" + + def test_client_init_with_api_key(self): + """测试客户端初始化""" + client = SoraClient(api_key="test-key") + self.assertEqual(client.api_key, "test-key") + self.assertEqual(client.base_url, DEFAULT_SORA_BASE_URL) + + def test_client_init_without_api_key(self): + """测试缺少 API key 时抛出错误""" + with self.assertRaises(ValueError): + SoraClient(api_key="") + + @patch('sora_node.urllib.request.urlopen') + def test_create_job_success(self, mock_urlopen): + """测试创建任务成功""" + # Mock HTTP 响应 + mock_response = MagicMock() + mock_response.read.return_value = b'{"id": "job-123", "status": "submitted"}' + mock_urlopen.return_value.__enter__.return_value = mock_response + + client = SoraClient(api_key="test-key") + result = client.create_job({"prompt": "test prompt"}) + + self.assertEqual(result["id"], "job-123") + self.assertEqual(result["status"], "submitted") + + @patch('sora_node.urllib.request.urlopen') + def test_create_job_error(self, mock_urlopen): + """测试创建任务失败""" + # Mock HTTP 错误 + mock_error = HTTPError("http://test.com", 400, "Bad Request", {}, None) + mock_error.read.return_value = b'{"error": "Invalid prompt"}' + mock_urlopen.side_effect = mock_error + + client = SoraClient(api_key="test-key") + with self.assertRaises(SoraAPIError): + client.create_job({"prompt": "test"}) + + +class TestSoraTextToVideo(unittest.TestCase): + """测试 SoraTextToVideo 节点类""" + + def test_input_types(self): + """测试节点输入类型定义""" + node = SoraTextToVideo() + input_types = node.INPUT_TYPES() + + self.assertIn("required", input_types) + self.assertIn("optional", input_types) + self.assertIn("prompt", input_types["required"]) + + def test_resolve_api_key_from_env(self): + """测试从环境变量解析 API key""" + # 测试 SORA_API_KEY + os.environ["SORA_API_KEY"] = "env-key-123" + from sora_node import _resolve_api_key + key = _resolve_api_key(None) + self.assertEqual(key, "env-key-123") + + # 测试 OPENAI_API_KEY（fallback） + del os.environ["SORA_API_KEY"] + os.environ["OPENAI_API_KEY"] = "openai-key-456" + key = _resolve_api_key(None) + self.assertEqual(key, "openai-key-456") + + # 清理 + if "OPENAI_API_KEY" in os.environ: + del os.environ["OPENAI_API_KEY"] + + @patch('sora_node.SoraClient') + def test_generate_basic(self, mock_client_class): + """测试基本的生成功能""" + # Mock 客户端 + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + # Mock API 响应 + mock_client.create_job.return_value = { + "id": "test-job-123", + "status": "submitted" + } + + node = SoraTextToVideo() + job_id, status, result = node.generate( + prompt="test prompt", + api_key="test-key", + wait_for_result=False + ) + + self.assertEqual(job_id, "test-job-123") + self.assertEqual(status, "submitted") + result_data = json.loads(result) + self.assertEqual(result_data["job_id"], "test-job-123") + + @patch('sora_node.SoraClient') + def test_generate_with_polling(self, mock_client_class): + """测试带轮询的生成""" + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + # Mock 创建任务 + mock_client.create_job.return_value = { + "id": "test-job-456", + "status": "processing" + } + + # Mock 轮询响应 + mock_client.fetch_job.return_value = { + "id": "test-job-456", + "status": "succeeded", + "result_url": "https://example.com/video.mp4" + } + + node = SoraTextToVideo() + job_id, status, result = node.generate( + prompt="test prompt", + api_key="test-key", + wait_for_result=True, + max_wait_seconds=10.0, + poll_interval=1.0 + ) + + self.assertEqual(job_id, "test-job-456") + self.assertEqual(status, "succeeded") + # 验证轮询被调用 + mock_client.fetch_job.assert_called() + + def test_extract_video_url(self): + """测试视频 URL 提取""" + node = SoraTextToVideo() + + # 测试 result_url + response1 = {"result_url": "https://example.com/video1.mp4"} + url1 = node._extract_video_url(response1) + self.assertEqual(url1, "https://example.com/video1.mp4") + + # 测试 outputs 数组 + response2 = { + "outputs": [{"url": "https://example.com/video2.mp4"}] + } + url2 = node._extract_video_url(response2) + self.assertEqual(url2, "https://example.com/video2.mp4") + + # 测试空响应 + response3 = {} + url3 = node._extract_video_url(response3) + self.assertIsNone(url3) + + +def test_node_mappings(): + """测试节点映射是否正确""" + from sora_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + + assert "SoraTextToVideo" in NODE_CLASS_MAPPINGS + assert NODE_CLASS_MAPPINGS["SoraTextToVideo"] == SoraTextToVideo + assert "SoraTextToVideo" in NODE_DISPLAY_NAME_MAPPINGS + print("✓ 节点映射测试通过") + + +def run_basic_tests(): + """运行基本功能测试（不需要真实 API）""" + print("=" * 60) + print("运行基本功能测试...") + print("=" * 60) + + # 测试节点映射 + test_node_mappings() + + # 测试输入类型 + node = SoraTextToVideo() + input_types = node.INPUT_TYPES() + assert "required" in input_types + assert "prompt" in input_types["required"] + print("✓ 节点输入类型测试通过") + + # 测试 API key 解析 + os.environ["SORA_API_KEY"] = "test-key-123" + from sora_node import _resolve_api_key + key = _resolve_api_key(None) + assert key == "test-key-123" + del os.environ["SORA_API_KEY"] + print("✓ API key 解析测试通过") + + print("\n所有基本测试通过！") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="测试 Sora 自定义节点") + parser.add_argument("--test-class", action="store_true", help="运行类测试") + parser.add_argument("--test-client", action="store_true", help="测试客户端（需要真实 API key）") + parser.add_argument("--api-key", type=str, help="API key（用于客户端测试）") + parser.add_argument("--unittest", action="store_true", help="运行完整的 unittest 测试套件") + + args = parser.parse_args() + + if args.unittest: + # 运行完整的 unittest 测试 + unittest.main(argv=[''], exit=False, verbosity=2) + elif args.test_class: + # 运行基本类测试 + run_basic_tests() + elif args.test_client and args.api_key: + # 测试真实 API 客户端（需要网络连接和有效 API key） + print("测试真实 API 客户端...") + print("注意：这需要有效的 API key 和网络连接") + try: + client = SoraClient(api_key=args.api_key) + # 这里可以添加真实的 API 调用测试 + print("客户端初始化成功") + except Exception as e: + print(f"客户端测试失败: {e}") + else: + # 默认运行基本测试 + run_basic_tests() + print("\n提示：") + print(" - 运行完整测试: python test_sora_node.py --unittest") + print(" - 运行基本测试: python test_sora_node.py --test-class") + print(" - 测试真实 API: python test_sora_node.py --test-client --api-key YOUR_KEY") +