Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
359 lines
16 KiB
Python
359 lines
16 KiB
Python
"""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}")
|