Publish current ComfyUI O1Key code baseline

Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
2026-09-24 19:56:48 +08:00
parent 3e337722ab
commit ba920f2b66
183 changed files with 49496 additions and 9934 deletions
@@ -0,0 +1,358 @@
+"""O1Key Omni Flash JSON video tasks."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+import time
+from typing import Any, Callable
+
+import aiohttp
+
+from ..utils.config import get_api_key_or_raise, get_base_url_by_route
+from ..utils.video_task import (
+    InterruptProcessingException, check_interrupt, download_video_to_file,
+    extract_error_message, extract_progress, extract_video_url,
+    interruptible_sleep, is_failure_status, is_success_status,
+)
+
+
+MODELS = {"omni_flash_8s", "omni_flash_10s", "omni_flash_abra_edit"}
+RESOLUTIONS = {"720p", "1080p"}
+RATIOS = {"16:9", "9:16"}
+POLL_SECONDS = 7
+POLL_DEADLINE_SECONDS = 2000
+ERROR_HINTS = {
+    "invalid_request": "请求参数有误，请检查模型、分辨率、宽高比和素材",
+    "model_not_available": "当前模型不可用，请重新选择模型",
+    "image_url_required_for_i2v": "参考图地址缺失或无效，请连接图片后重试",
+    "invalid_api_key": "O1Key 令牌无效或已停用，请在令牌管理中更新",
+    "insufficient_balance": "O1Key 余额不足",
+    "task_not_found": "视频任务不存在或已失效",
+    "rate_limit_exceeded": "请求过于频繁，请稍后重试",
+}
+HTTP_HINTS = {
+    400: "请求参数错误", 401: "令牌验证失败", 402: "余额不足",
+    404: "任务不存在", 429: "请求过于频繁",
+}
+SENSITIVE_RESPONSE_KEYS = {
+    "authorization", "api_key", "apikey", "access_token", "refresh_token",
+    "token", "secret", "password", "b64_json", "base64", "image_base64",
+    "video_base64",
+}
+MAX_LOG_BODY = 16000
+
+
+def build_video_body(
+    *, model: str, prompt: str, resolution: str, aspect_ratio: str,
+    mode: str, references: list[str] | None = None, source_video_url: str = "",
+) -> dict[str, Any]:
+    """Validate all scalar inputs before a paid request."""
+    if model not in MODELS:
+        raise ValueError("Omni Flash 模型无效")
+    prompt = str(prompt or "").strip()
+    if not prompt:
+        raise ValueError("提示词不能为空")
+    if len(prompt) > 20000:
+        raise ValueError("提示词过长")
+    if resolution not in RESOLUTIONS or aspect_ratio not in RATIOS:
+        raise ValueError("分辨率或宽高比无效")
+    references = list(references or [])
+    if any(not isinstance(value, str) or len(value) > 4096 or not value.startswith(("https://", "http://")) for value in references):
+        raise ValueError("参考图必须是 HTTP(S) 直链")
+    body: dict[str, Any] = {
+        "model": model, "prompt": prompt,
+        "resolution": resolution, "aspect_ratio": aspect_ratio,
+    }
+    if mode == "edit":
+        if model != "omni_flash_abra_edit" or len(source_video_url) > 4096 or not source_video_url.startswith(("https://", "http://")):
+            raise ValueError("视频编辑需要编辑模型和源视频直链")
+        if len(references) > 5:
+            raise ValueError("视频编辑最多支持 5 张参考图")
+        body["source_video_url"] = source_video_url
+    elif mode in {"text", "reference", "first_last_frame"}:
+        if model == "omni_flash_abra_edit" or source_video_url:
+            raise ValueError("生成模式不能使用编辑模型或源视频")
+        if mode == "text" and references:
+            raise ValueError("文生视频不能提供参考图")
+        if mode == "reference" and not references:
+            raise ValueError("参考图模式至少需要 1 张图片")
+        if mode == "first_last_frame":
+            if not 1 <= len(references) <= 2:
+                raise ValueError("首尾帧模式需要首帧图片，尾帧图片可选")
+            # The provider's frame-pair flag is for a transition between two
+            # frames. A lone first frame uses the documented single-image i2v
+            # request, avoiding a pair request with a missing end frame.
+            if len(references) == 2:
+                body["first_last_frame"] = True
+    else:
+        raise ValueError("Omni Flash 生成模式无效")
+    if references:
+        body["input_reference"] = references[0] if len(references) == 1 else references
+    return body
+
+
+def _submission_payload(body: dict[str, Any]) -> dict[str, Any]:
+    """Use a scalar JSON reference, or repeat the field in multipart for several."""
+    references = body.get("input_reference")
+    if not isinstance(references, list):
+        return {"json": body}
+    form = aiohttp.FormData()
+    for name, value in body.items():
+        values = value if name == "input_reference" else [value]
+        for item in values:
+            text = "true" if item is True else "false" if item is False else str(item)
+            form.add_field(name, text, content_type="text/plain")
+    return {"data": form}
+
+
+def _redact_log_string(text: str) -> str:
+    text = re.sub(r"https?://[^\s\"'<>]+", "<URL 已隐藏>", text)
+    text = re.sub(r"(?i)bearer\s+[^\s\"']+", "Bearer <已隐藏>", text)
+    text = re.sub(r"(?i)(?:api[_-]?key|token|authorization)[\"']?\s*[:=]\s*[\"']?[^\s,;\"']+", "<凭据已隐藏>", text)
+    text = re.sub(r"(?i)data:[^,\s]+;base64,[A-Za-z0-9+/=]+", "<Base64 已隐藏>", text)
+    return re.sub(r"[A-Za-z0-9+/]{256,}={0,2}", "<长数据已隐藏>", text)
+
+
+def _safe_error(value: Any) -> str:
+    text = _redact_log_string(str(value or "请求失败"))
+    return text[:400]
+
+
+def _log_value(value: Any, key: str = "", depth: int = 0) -> Any:
+    if key.lower() in SENSITIVE_RESPONSE_KEYS:
+        return "<已隐藏>"
+    if depth >= 12:
+        return "<嵌套内容已省略>"
+    if isinstance(value, dict):
+        return {
+            _redact_log_string(str(name)[:200]): _log_value(item, str(name), depth + 1)
+            for name, item in value.items()
+        }
+    if isinstance(value, list):
+        return [_log_value(item, key, depth + 1) for item in value[:50]] + (
+            [f"<其余 {len(value) - 50} 项已省略>"] if len(value) > 50 else []
+        )
+    if isinstance(value, str):
+        if len(value) > 1200:
+            return f"<长文本 {len(value)} 字符已省略>"
+        return _redact_log_string(value)
+    return value
+
+
+def _log_response_body(stage: str, status: int, raw_body: str) -> None:
+    try:
+        payload = json.loads(raw_body)
+    except (ValueError, TypeError):
+        safe_body = _safe_error(raw_body) if raw_body else "<空响应体>"
+    else:
+        safe_body = json.dumps(_log_value(payload), ensure_ascii=False, separators=(",", ":"))
+    if len(safe_body) > MAX_LOG_BODY:
+        safe_body = f"{safe_body[:MAX_LOG_BODY]}...<后续内容已省略>"
+    print(f"[Omni Flash] {stage} HTTP {status} 原始响应体（敏感值已隐藏）：{safe_body}")
+
+
+async def _response_text(response: aiohttp.ClientResponse, stage: str) -> str:
+    raw_body = await response.text()
+    _log_response_body(stage, response.status, raw_body)
+    return raw_body
+
+
+def _error_code(payload: Any) -> str:
+    if not isinstance(payload, dict):
+        return ""
+    data = payload.get("data")
+    inner = data.get("data") if isinstance(data, dict) else None
+    for source in (inner, data, payload):
+        if not isinstance(source, dict):
+            continue
+        error = source.get("error")
+        for value in (error.get("code") if isinstance(error, dict) else None, source.get("code")):
+            if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{1,80}", value):
+                return value.lower()
+    return ""
+
+
+def _response_error(payload: Any, status: int) -> str:
+    detail = extract_error_message(payload, default="") if isinstance(payload, dict) else payload
+    code = _error_code(payload)
+    hint = ERROR_HINTS.get(code) or HTTP_HINTS.get(status, "视频接口请求失败")
+    detail = _safe_error(detail) if detail else ""
+    if detail == code or detail == hint:
+        detail = ""
+    suffix = f"：{detail}" if detail else ""
+    code_note = f"，{code}" if code else ""
+    return f"{hint}（HTTP {status}{code_note}）{suffix}"
+
+
+def _task_error(payload: dict[str, Any]) -> str:
+    code = _error_code(payload)
+    hint = ERROR_HINTS.get(code, "视频任务生成失败")
+    detail = extract_error_message(payload, default="")
+    detail = _safe_error(detail) if detail else ""
+    if detail == code or detail == hint:
+        detail = ""
+    code_note = f"（{code}）" if code else ""
+    return f"{hint}{code_note}{f'：{detail}' if detail else ''}"
+
+
+def _task_id(payload: Any) -> str | None:
+    if not isinstance(payload, dict):
+        return None
+    data = payload.get("data")
+    inner = data.get("data") if isinstance(data, dict) else None
+    for source in (inner, data, payload):
+        if isinstance(source, dict):
+            for name in ("id", "task_id", "video_id"):
+                value = source.get(name)
+                if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{8,128}", value):
+                    return value
+    return None
+
+
+def _task_status(payload: dict[str, Any]) -> str:
+    data = payload.get("data")
+    inner = data.get("data") if isinstance(data, dict) else None
+    for source in (inner, data, payload):
+        if isinstance(source, dict):
+            for name in ("task_status", "task_state", "status", "state"):
+                value = source.get(name)
+                if value is not None and str(value).strip():
+                    return str(value).strip().lower()
+    return ""
+
+
+def _video_url(payload: Any) -> str | None:
+    value = extract_video_url(payload) if isinstance(payload, dict) else None
+    if not value and isinstance(payload, dict):
+        data = payload.get("data")
+        for source in (data, payload):
+            if isinstance(source, dict):
+                output = source.get("output")
+                if isinstance(output, dict):
+                    value = output.get("video_url") or output.get("url")
+                    if value:
+                        break
+    return value if isinstance(value, str) and len(value) <= 8192 and value.startswith(("https://", "http://")) else None
+
+
+class OmniFlashClient:
+    def __init__(self, *, base_url: str | None = None, api_key: str | None = None):
+        self.base_url = (base_url or get_base_url_by_route()).rstrip("/")
+        self.api_key = api_key or get_api_key_or_raise("O1KEY_API_KEY")
+
+    async def _download_completed(
+        self, session: aiohttp.ClientSession, task_id: str, save_path: str,
+        headers: dict[str, str], status_payload: dict[str, Any],
+    ) -> None:
+        content_url = f"{self.base_url}/v1/videos/{task_id}/content"
+        result_url = _video_url(status_payload)
+        download_url, download_headers = content_url, headers
+        try:
+            async with session.get(content_url, headers=headers, allow_redirects=True) as response:
+                if response.status >= 300:
+                    raw_body = await _response_text(response, "下载")
+                    if not result_url:
+                        try:
+                            error = json.loads(raw_body)
+                        except ValueError:
+                            error = raw_body
+                        raise RuntimeError(_response_error(error, response.status))
+                    download_url, download_headers = result_url, {}
+                elif "json" in response.headers.get("Content-Type", "").lower():
+                    raw_body = await _response_text(response, "下载")
+                    try:
+                        content_payload = json.loads(raw_body)
+                    except ValueError:
+                        raise RuntimeError("视频下载接口返回了无效 JSON") from None
+                    if _error_code(content_payload) in ERROR_HINTS:
+                        raise RuntimeError(_task_error(content_payload))
+                    download_url = _video_url(content_payload) or result_url
+                    if not download_url:
+                        raise RuntimeError("任务已完成，但下载响应未提供视频地址")
+                    download_headers = {}
+                else:
+                    print(f"[Omni Flash] 下载 HTTP {response.status} 响应体：<视频二进制，未打印>")
+        except (aiohttp.ClientError, asyncio.TimeoutError):
+            # The streaming downloader handles transient connection failures and resumes.
+            if result_url:
+                download_url, download_headers = result_url, {}
+        await download_video_to_file(
+            session, download_url, save_path, headers=download_headers or None,
+            label="Omni Flash 视频",
+        )
+
+    async def generate(
+        self, body: dict[str, Any], save_path: str,
+        progress: Callable[[str, int, str], None] | None = None,
+    ) -> str:
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        submit_headers = dict(headers)
+        if body.get("model") == "omni_flash_abra_edit":
+            submit_headers["X-No-Watermark"] = "video"
+        timeout = aiohttp.ClientTimeout(total=120, connect=30)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(
+                f"{self.base_url}/v1/videos", headers=submit_headers,
+                **_submission_payload(body),
+            ) as response:
+                raw_body = await _response_text(response, "提交")
+                try:
+                    payload = json.loads(raw_body)
+                except ValueError:
+                    if response.status >= 300:
+                        raise RuntimeError(_response_error(raw_body, response.status)) from None
+                    raise RuntimeError("提交接口未返回有效 JSON") from None
+                if response.status >= 300:
+                    raise RuntimeError(_response_error(payload, response.status))
+                if _error_code(payload) in ERROR_HINTS:
+                    raise RuntimeError(_task_error(payload))
+            task_id = _task_id(payload)
+            if not task_id:
+                raise RuntimeError("接口未返回有效任务 ID")
+            if progress:
+                progress("polling", 0, task_id)
+            deadline = time.monotonic() + POLL_DEADLINE_SECONDS
+            last_status = ""
+            while time.monotonic() < deadline:
+                await interruptible_sleep(POLL_SECONDS)
+                check_interrupt()
+                try:
+                    async with session.get(f"{self.base_url}/v1/videos/{task_id}", headers=headers) as response:
+                        raw_body = await _response_text(response, "查询")
+                        if response.status in {408, 500, 502, 503, 504}:
+                            continue
+                        try:
+                            status_payload = json.loads(raw_body)
+                        except ValueError:
+                            if response.status >= 300:
+                                raise RuntimeError(_response_error(raw_body, response.status)) from None
+                            raise
+                        if response.status >= 300:
+                            raise RuntimeError(_response_error(status_payload, response.status))
+                except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, ValueError):
+                    continue
+                if not isinstance(status_payload, dict):
+                    continue
+                status = _task_status(status_payload)
+                last_status = status or last_status
+                code = _error_code(status_payload)
+                if is_failure_status(status) or code in ERROR_HINTS or (
+                    not status and code and code not in {"ok", "success", "0"}
+                ):
+                    raise RuntimeError(_task_error(status_payload))
+                if is_success_status(status) or _video_url(status_payload):
+                    if progress:
+                        progress("downloading", 100, task_id)
+                    try:
+                        await self._download_completed(session, task_id, save_path, headers, status_payload)
+                    except InterruptProcessingException:
+                        raise
+                    except Exception as exc:
+                        raise RuntimeError(_safe_error(exc)) from None
+                    return task_id
+                if progress:
+                    progress("polling", extract_progress(status_payload), task_id)
+            status_note = f"，最后状态：{_safe_error(last_status)}" if last_status else ""
+            raise TimeoutError(f"Omni Flash 任务 {task_id} 等待超时{status_note}")