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.
362 lines
16 KiB
Python
362 lines
16 KiB
Python
"""Client for the complete O1Key Grok Imagine Video API."""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from typing import Any, Callable, Dict, List, Optional
|
||
from urllib.parse import quote
|
||
|
||
import aiohttp
|
||
|
||
from .base_client import BaseAPIClient
|
||
from ..utils.config import get_api_base_url, get_api_key_or_raise
|
||
from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message
|
||
from ..utils.video_task import (
|
||
POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
|
||
check_interrupt,
|
||
download_video_to_file,
|
||
extract_error_message,
|
||
interruptible_sleep,
|
||
run_with_interrupt,
|
||
)
|
||
|
||
|
||
class GrokVideoClient(BaseAPIClient):
|
||
"""Submit, poll, and download Grok video generation, edit, or extension tasks."""
|
||
|
||
ENDPOINTS = {
|
||
"generate": "/grok/v1/videos/generations",
|
||
"edit": "/grok/v1/videos/edits",
|
||
"extend": "/grok/v1/videos/extensions",
|
||
}
|
||
STATUS_ENDPOINT = "/grok/v1/videos/{request_id}"
|
||
|
||
BASE_MODEL = "grok-imagine-video"
|
||
LATEST_MODEL = "grok-imagine-video-1.5"
|
||
DEFAULT_MODEL = LATEST_MODEL
|
||
# Kept as a compatibility alias for callers that imported the old constant.
|
||
TEXT_TO_VIDEO_MODEL = BASE_MODEL
|
||
IMAGE_TO_VIDEO_MODELS = (LATEST_MODEL,)
|
||
MODEL_OPTIONS = (BASE_MODEL, LATEST_MODEL)
|
||
ASPECT_RATIO_OPTIONS = ("16:9", "9:16", "1:1", "4:3", "3:4", "3:2", "2:3")
|
||
RESOLUTION_OPTIONS = ("480p", "720p", "1080p")
|
||
SUCCESS_STATUSES = {"done"}
|
||
FAILURE_STATUSES = {"failed", "expired"}
|
||
POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS
|
||
|
||
def __init__(self, base_url: Optional[str] = None):
|
||
super().__init__(
|
||
base_url=(base_url or get_api_base_url()).rstrip("/"),
|
||
api_key=get_api_key_or_raise("O1KEY_API_KEY"),
|
||
)
|
||
|
||
def get_endpoint(self, operation: str = "generate", **kwargs) -> str:
|
||
try:
|
||
return self.ENDPOINTS[operation]
|
||
except KeyError:
|
||
raise ValueError(f"不支持的 Grok 操作：{operation}。") from None
|
||
|
||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||
return self.build_video_body(**kwargs)
|
||
|
||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||
return response
|
||
|
||
@staticmethod
|
||
def _locator(
|
||
value: Optional[Dict[str, str]],
|
||
label: str,
|
||
allowed_keys: tuple[str, ...],
|
||
) -> Dict[str, str]:
|
||
if not isinstance(value, dict):
|
||
raise ValueError(f"{label}必须提供媒体定位对象。")
|
||
known_keys = ("url", "image_url", "file_id", "voice_id")
|
||
provided_keys = {
|
||
key
|
||
for key in known_keys
|
||
if value.get(key) is not None and str(value[key]).strip()
|
||
}
|
||
locator = {
|
||
key: str(value[key]).strip()
|
||
for key in allowed_keys
|
||
if value.get(key) is not None and str(value[key]).strip()
|
||
}
|
||
if len(locator) != 1 or provided_keys != set(locator):
|
||
supported = "、".join(allowed_keys)
|
||
raise ValueError(f"{label}必须且只能提供 {supported} 中的一项。")
|
||
return locator
|
||
|
||
@classmethod
|
||
def _image_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
|
||
return cls._locator(value, label, ("url", "image_url"))
|
||
|
||
@classmethod
|
||
def _audio_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
|
||
return cls._locator(value, label, ("url", "voice_id"))
|
||
|
||
@classmethod
|
||
def _video_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
|
||
return cls._locator(value, label, ("url", "file_id"))
|
||
|
||
@classmethod
|
||
def _validate_common_generation(
|
||
cls, model: str, duration: int, aspect_ratio: str, resolution: str
|
||
) -> int:
|
||
if model not in cls.MODEL_OPTIONS:
|
||
raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
|
||
try:
|
||
duration = int(duration)
|
||
except (TypeError, ValueError):
|
||
raise ValueError("时长必须是整数。") from None
|
||
if not 1 <= duration <= 15:
|
||
raise ValueError("生成时长仅支持 1 到 15 秒。")
|
||
if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
|
||
raise ValueError(f"宽高比仅支持：{', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
|
||
if resolution not in cls.RESOLUTION_OPTIONS:
|
||
raise ValueError(f"分辨率仅支持：{', '.join(cls.RESOLUTION_OPTIONS)}。")
|
||
return duration
|
||
|
||
@classmethod
|
||
def build_video_body(
|
||
cls,
|
||
*,
|
||
operation: str,
|
||
prompt: str,
|
||
model: str,
|
||
duration: Optional[int] = None,
|
||
aspect_ratio: str = "16:9",
|
||
resolution: str = "480p",
|
||
image: Optional[Dict[str, str]] = None,
|
||
reference_images: Optional[List[Dict[str, str]]] = None,
|
||
reference_audios: Optional[List[Dict[str, str]]] = None,
|
||
video: Optional[Dict[str, str]] = None,
|
||
) -> Dict[str, Any]:
|
||
if operation not in cls.ENDPOINTS:
|
||
raise ValueError(f"不支持的 Grok 操作：{operation}。")
|
||
prompt = (prompt or "").strip()
|
||
if reference_images is not None and not isinstance(reference_images, (list, tuple)):
|
||
raise ValueError("reference_images 必须是数组。")
|
||
if reference_audios is not None and not isinstance(reference_audios, (list, tuple)):
|
||
raise ValueError("reference_audios 必须是数组。")
|
||
references = list(reference_images or [])
|
||
audios = list(reference_audios or [])
|
||
|
||
if operation == "generate":
|
||
duration = cls._validate_common_generation(model, duration, aspect_ratio, resolution)
|
||
normal_image = cls._image_locator(image, "图生视频参考图") if image else None
|
||
normal_references = [cls._image_locator(item, "参考图") for item in references]
|
||
normal_audios = [cls._audio_locator(item, "参考音频") for item in audios]
|
||
if normal_image and normal_references:
|
||
raise ValueError("image 和 reference_images 不能同时使用。")
|
||
if len(normal_references) > 7:
|
||
raise ValueError("参考生视频最多支持 7 张参考图。")
|
||
if len(normal_audios) > 3:
|
||
raise ValueError("参考生视频最多支持 3 个参考音频。")
|
||
|
||
has_reference_assets = bool(normal_references or normal_audios)
|
||
if has_reference_assets:
|
||
if not prompt:
|
||
raise ValueError("参考图/音频生视频必须填写提示词。")
|
||
if resolution == "1080p":
|
||
raise ValueError("参考图/音频生视频不支持 1080p。")
|
||
elif not normal_image and not prompt:
|
||
raise ValueError("文生视频必须填写提示词。")
|
||
|
||
if resolution == "1080p" and model != cls.LATEST_MODEL:
|
||
raise ValueError("1080p 仅支持 grok-imagine-video-1.5 的文生或图生视频。")
|
||
|
||
body: Dict[str, Any] = {
|
||
"model": model,
|
||
"duration": duration,
|
||
"aspect_ratio": aspect_ratio,
|
||
"resolution": resolution,
|
||
}
|
||
if prompt:
|
||
body["prompt"] = prompt
|
||
if normal_image:
|
||
body["image"] = normal_image
|
||
if normal_references:
|
||
body["reference_images"] = normal_references
|
||
if normal_audios:
|
||
body["reference_audios"] = normal_audios
|
||
return body
|
||
|
||
if model not in cls.MODEL_OPTIONS:
|
||
raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
|
||
if not prompt:
|
||
raise ValueError(f"{operation} 必须填写提示词。")
|
||
normal_video = cls._video_locator(video, "输入视频")
|
||
if operation == "edit":
|
||
return {"model": model, "prompt": prompt, "video": normal_video}
|
||
|
||
try:
|
||
duration = int(duration)
|
||
except (TypeError, ValueError):
|
||
raise ValueError("续写时长必须是整数。") from None
|
||
if not 2 <= duration <= 10:
|
||
raise ValueError("视频续写时长仅支持 2 到 10 秒。")
|
||
return {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"video": normal_video,
|
||
"duration": duration,
|
||
}
|
||
|
||
@staticmethod
|
||
def _extract_request_id(payload: Dict[str, Any]) -> Optional[str]:
|
||
for source in (payload, payload.get("data")):
|
||
if isinstance(source, dict) and source.get("request_id"):
|
||
return str(source["request_id"])
|
||
return None
|
||
|
||
@staticmethod
|
||
def _safe_filename(request_id: str) -> str:
|
||
return re.sub(r"[^A-Za-z0-9_.-]+", "_", request_id).strip("._") or "grok_video"
|
||
|
||
@staticmethod
|
||
def _safe_error_message(value: object) -> str:
|
||
message = str(value or "").strip()
|
||
message = re.sub(r"data:[^\s,;]+;base64,[A-Za-z0-9+/=_-]+", "<base64 omitted>", message)
|
||
message = re.sub(r"https?://[^\s\"'<>]+", "<temporary URL omitted>", message)
|
||
return message[:500]
|
||
|
||
async def _request_json(
|
||
self,
|
||
method: str,
|
||
endpoint: str,
|
||
session: aiohttp.ClientSession,
|
||
*,
|
||
json_body: Optional[Dict[str, Any]] = None,
|
||
timeout_seconds: int = 120,
|
||
request_id: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
url = f"{self.base_url}{endpoint}"
|
||
timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
|
||
last_status, last_text = 0, ""
|
||
for attempt in range(4):
|
||
check_interrupt()
|
||
response = None
|
||
try:
|
||
response = await run_with_interrupt(
|
||
session.request(
|
||
method, url, json=json_body,
|
||
headers=self.get_headers(use_bearer_token=True), timeout=timeout,
|
||
)
|
||
)
|
||
text = await run_with_interrupt(response.text())
|
||
last_status, last_text = response.status, text
|
||
if 200 <= response.status < 300:
|
||
try:
|
||
return json.loads(text) if text.strip() else {}
|
||
except json.JSONDecodeError:
|
||
raise RuntimeError("Grok Video 响应不是有效 JSON。") from None
|
||
if response.status not in RETRYABLE_STATUS_CODES or attempt == 3:
|
||
break
|
||
delay = min(2 ** attempt, 8)
|
||
print(f"Grok Video：HTTP {response.status}，{delay}s 后重试…")
|
||
await interruptible_sleep(delay)
|
||
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||
if attempt == 3:
|
||
raise RuntimeError(
|
||
f"Grok Video 网络错误：{type(exc).__name__}"
|
||
) from None
|
||
delay = min(2 ** attempt, 8)
|
||
print(f"Grok Video：网络错误，{delay}s 后重试…")
|
||
await interruptible_sleep(delay)
|
||
finally:
|
||
if response is not None:
|
||
response.release()
|
||
|
||
message = self._safe_error_message(
|
||
get_friendly_message(last_status, last_text) or "请求失败"
|
||
)
|
||
detail = f"Grok Video 请求失败：HTTP {last_status}，{message}"
|
||
if request_id:
|
||
detail += f"（request_id: {request_id}）"
|
||
raise RuntimeError(detail)
|
||
|
||
async def _poll(
|
||
self, request_id: str, session: aiohttp.ClientSession, *, poll_interval: int,
|
||
timeout: int, progress_callback: Optional[Callable[[int, str, float], None]],
|
||
) -> Dict[str, Any]:
|
||
endpoint = self.STATUS_ENDPOINT.format(request_id=quote(request_id, safe=""))
|
||
started_at = time.monotonic()
|
||
while True:
|
||
await interruptible_sleep(poll_interval)
|
||
response = await self._request_json(
|
||
"GET", endpoint, session, timeout_seconds=60, request_id=request_id
|
||
)
|
||
status = str(response.get("status", "")).strip().lower()
|
||
try:
|
||
progress = max(0, min(100, int(float(response.get("progress") or 0))))
|
||
except (TypeError, ValueError):
|
||
progress = 0
|
||
elapsed = time.monotonic() - started_at
|
||
if progress_callback:
|
||
progress_callback(progress, status, elapsed)
|
||
if status in self.SUCCESS_STATUSES:
|
||
return response
|
||
if status in self.FAILURE_STATUSES:
|
||
message = self._safe_error_message(
|
||
extract_error_message(response, default="未知错误")
|
||
)
|
||
raise RuntimeError(
|
||
f"Grok Video 任务{status}（request_id: {request_id}）："
|
||
f"{message}"
|
||
)
|
||
if elapsed >= timeout:
|
||
raise TimeoutError(
|
||
f"Grok Video 轮询超时（request_id: {request_id}，状态：{status or 'unknown'}）。"
|
||
)
|
||
|
||
def run_video_sync(
|
||
self, *, operation: str, prompt: str, model: str, duration: Optional[int] = None,
|
||
aspect_ratio: str = "16:9", resolution: str = "480p",
|
||
image: Optional[Dict[str, str]] = None,
|
||
reference_images: Optional[List[Dict[str, str]]] = None,
|
||
reference_audios: Optional[List[Dict[str, str]]] = None,
|
||
video: Optional[Dict[str, str]] = None, output_dir: Optional[str] = None,
|
||
poll_interval: int = 5, timeout: int = VIDEO_POLL_DEADLINE_SECONDS,
|
||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||
) -> Dict[str, Any]:
|
||
async def run_request() -> Dict[str, Any]:
|
||
async with self._make_session() as session:
|
||
endpoint = self.get_endpoint(operation)
|
||
body = self.build_video_body(
|
||
operation=operation, prompt=prompt, model=model, duration=duration,
|
||
aspect_ratio=aspect_ratio, resolution=resolution, image=image,
|
||
reference_images=reference_images, reference_audios=reference_audios,
|
||
video=video,
|
||
)
|
||
print(f"Grok Video：正在提交{operation}任务…")
|
||
created = await self._request_json(
|
||
"POST", endpoint, session, json_body=body, timeout_seconds=180
|
||
)
|
||
request_id = self._extract_request_id(created)
|
||
if not request_id:
|
||
raise RuntimeError("Grok Video 创建响应中没有 request_id。")
|
||
print(f"Grok Video：任务已提交，request_id：{request_id}")
|
||
completed = await self._poll(
|
||
request_id, session, poll_interval=max(1, int(poll_interval)),
|
||
timeout=timeout, progress_callback=progress_callback,
|
||
)
|
||
video_data = completed.get("video")
|
||
video_url = video_data.get("url") if isinstance(video_data, dict) else None
|
||
if not video_url:
|
||
raise RuntimeError(f"Grok Video 完成响应中没有 video.url（request_id: {request_id}）。")
|
||
directory = output_dir or os.getcwd()
|
||
os.makedirs(directory, exist_ok=True)
|
||
save_path = os.path.join(directory, f"{self._safe_filename(request_id)}.mp4")
|
||
print("Grok Video：视频生成完成，正在下载…")
|
||
video_path = await download_video_to_file(session, video_url, save_path, label="Grok Video")
|
||
return {
|
||
"request_id": request_id,
|
||
"video_path": video_path,
|
||
"duration": video_data.get("duration"),
|
||
"raw_json": {"create": created, "status": completed},
|
||
}
|
||
|
||
return self.run_async_in_thread(run_request())
|