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.
This commit is contained in:
+262
-384
@@ -1,19 +1,12 @@
|
||||
"""
|
||||
Grok Video API client.
|
||||
|
||||
Flow:
|
||||
1. POST /v1/videos
|
||||
2. GET /v1/videos/{task_id}
|
||||
3. GET /v1/videos/{task_id}/content, or download a URL from the status body
|
||||
"""
|
||||
"""Client for the complete O1Key Grok Imagine Video API."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
import aiohttp
|
||||
|
||||
@@ -21,44 +14,49 @@ 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,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
|
||||
class GrokVideoClient(BaseAPIClient):
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{task_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
|
||||
"""Submit, poll, and download Grok video generation, edit, or extension tasks."""
|
||||
|
||||
MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
|
||||
ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
|
||||
QUALITY_OPTIONS = ["720p"]
|
||||
MODEL_SECONDS_OPTIONS = {
|
||||
"grok-imagine-1.0-video": [6, 10, 12, 16, 20],
|
||||
}
|
||||
QUALITY_API_MAP = {
|
||||
"720p": "high",
|
||||
"high": "high",
|
||||
ENDPOINTS = {
|
||||
"generate": "/grok/v1/videos/generations",
|
||||
"edit": "/grok/v1/videos/edits",
|
||||
"extend": "/grok/v1/videos/extensions",
|
||||
}
|
||||
STATUS_ENDPOINT = "/grok/v1/videos/{request_id}"
|
||||
|
||||
SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"}
|
||||
FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"}
|
||||
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):
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
resolved_base_url = (base_url or "").strip() or get_api_base_url()
|
||||
super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
|
||||
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, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
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)
|
||||
@@ -66,418 +64,298 @@ class GrokVideoClient(BaseAPIClient):
|
||||
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,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str = "720p",
|
||||
images: Optional[List[str]] = None,
|
||||
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 not prompt:
|
||||
raise ValueError("提示词不能为空。")
|
||||
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 aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
|
||||
raise ValueError(f"宽高比仅支持: {', '.join(cls.ASPECT_RATIO_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:
|
||||
seconds_value = int(seconds)
|
||||
duration = int(duration)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("秒数必须是整数。") from None
|
||||
|
||||
allowed_seconds = cls.MODEL_SECONDS_OPTIONS.get(model)
|
||||
if allowed_seconds is not None:
|
||||
if seconds_value not in allowed_seconds:
|
||||
raise ValueError(
|
||||
f"模型 {model} 仅支持秒数: "
|
||||
f"{', '.join(str(s) for s in allowed_seconds)}。"
|
||||
"请修改为正确的秒数后再发起请求。"
|
||||
)
|
||||
elif seconds_value < 5 or seconds_value > 15:
|
||||
raise ValueError("秒数仅支持 5 到 15。")
|
||||
|
||||
api_quality = cls.QUALITY_API_MAP.get(str(quality), str(quality))
|
||||
if api_quality != "high":
|
||||
raise ValueError("画质仅支持 720p。")
|
||||
|
||||
body: Dict[str, Any] = {
|
||||
raise ValueError("续写时长必须是整数。") from None
|
||||
if not 2 <= duration <= 10:
|
||||
raise ValueError("视频续写时长仅支持 2 到 10 秒。")
|
||||
return {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"seconds": str(seconds_value),
|
||||
"quality": api_quality,
|
||||
"video": normal_video,
|
||||
"duration": duration,
|
||||
}
|
||||
|
||||
image_list = [img for img in (images or []) if img]
|
||||
if image_list:
|
||||
body["images"] = image_list[:3]
|
||||
|
||||
return body
|
||||
|
||||
@staticmethod
|
||||
def _safe_task_filename(task_id: str) -> str:
|
||||
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
|
||||
return safe or "grok_video"
|
||||
|
||||
@staticmethod
|
||||
def _mask_body_for_log(body: Dict[str, Any]) -> Dict[str, Any]:
|
||||
log_body = dict(body)
|
||||
images = log_body.get("images")
|
||||
if isinstance(images, list):
|
||||
log_body["images"] = [f"<data-url chars={len(item)}>" for item in images]
|
||||
return log_body
|
||||
|
||||
@staticmethod
|
||||
def _extract_task_id(payload: Dict[str, Any]) -> Optional[str]:
|
||||
sources = [payload]
|
||||
data = payload.get("data")
|
||||
if isinstance(data, dict):
|
||||
sources.append(data)
|
||||
|
||||
for source in sources:
|
||||
for key in ("id", "task_id", "video_id"):
|
||||
value = source.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
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 _format_http_error(endpoint: str, status: int, error_text: str, task_id: Optional[str] = None) -> str:
|
||||
message = get_friendly_message(status, error_text)
|
||||
parts = [
|
||||
"Grok Video 请求失败。",
|
||||
f"endpoint: {endpoint}",
|
||||
f"http_status: {status}",
|
||||
]
|
||||
if task_id:
|
||||
parts.append(f"task_id: {task_id}")
|
||||
if message:
|
||||
parts.append(f"message: {message}")
|
||||
return "\n".join(parts)
|
||||
def _safe_filename(request_id: str) -> str:
|
||||
return re.sub(r"[^A-Za-z0-9_.-]+", "_", request_id).strip("._") or "grok_video"
|
||||
|
||||
@classmethod
|
||||
def _format_task_failure(cls, task_id: str, payload: Dict[str, Any]) -> str:
|
||||
return "\n".join(
|
||||
[
|
||||
"Grok Video 任务失败。",
|
||||
f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
|
||||
f"task_id: {task_id}",
|
||||
f"message: {extract_error_message(payload)}",
|
||||
]
|
||||
)
|
||||
@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_with_retry(
|
||||
async def _request_json(
|
||||
self,
|
||||
method: str,
|
||||
endpoint: str,
|
||||
session: aiohttp.ClientSession,
|
||||
task_id: Optional[str] = None,
|
||||
*,
|
||||
json_body: Optional[Dict[str, Any]] = None,
|
||||
max_retries: int = 3,
|
||||
timeout_seconds: int = 120,
|
||||
request_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
|
||||
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
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=headers, timeout=timeout)
|
||||
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 = response.status
|
||||
last_text = text
|
||||
|
||||
last_status, last_text = response.status, text
|
||||
if 200 <= response.status < 300:
|
||||
if not text.strip():
|
||||
return {}
|
||||
try:
|
||||
return json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"Grok Video 响应 JSON 解析失败,原始内容:{text[:500]}") from None
|
||||
|
||||
if response.status in RETRYABLE_STATUS_CODES and attempt < max_retries:
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(
|
||||
f"Grok Video:{get_friendly_message(response.status)} "
|
||||
f"{delay}s 后重试 ({attempt + 1}/{max_retries})..."
|
||||
)
|
||||
await interruptible_sleep(delay)
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
|
||||
if attempt < max_retries:
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:网络错误,{delay}s 后重试 ({attempt + 1}/{max_retries})...")
|
||||
await interruptible_sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(f"Grok Video 网络错误: {e}") from None
|
||||
|
||||
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()
|
||||
|
||||
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
body: Dict[str, Any],
|
||||
session: aiohttp.ClientSession,
|
||||
) -> Dict[str, Any]:
|
||||
print("Grok Video:正在提交任务...")
|
||||
return await self._request_json_with_retry(
|
||||
"POST",
|
||||
self.CREATE_ENDPOINT,
|
||||
session=session,
|
||||
json_body=body,
|
||||
timeout_seconds=180,
|
||||
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_video_status_async(
|
||||
self,
|
||||
task_id: str,
|
||||
session: aiohttp.ClientSession,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
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(task_id=task_id)
|
||||
start = time.time()
|
||||
interval = max(1, int(poll_interval))
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
|
||||
endpoint = self.STATUS_ENDPOINT.format(request_id=quote(request_id, safe=""))
|
||||
started_at = time.monotonic()
|
||||
while True:
|
||||
data = await self._request_json_with_retry(
|
||||
"GET",
|
||||
endpoint,
|
||||
session=session,
|
||||
task_id=task_id,
|
||||
timeout_seconds=60,
|
||||
await interruptible_sleep(poll_interval)
|
||||
response = await self._request_json(
|
||||
"GET", endpoint, session, timeout_seconds=60, request_id=request_id
|
||||
)
|
||||
|
||||
status = extract_status(data)
|
||||
progress = extract_progress(data)
|
||||
elapsed = time.time() - start
|
||||
|
||||
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 or is_success_status(status):
|
||||
return data
|
||||
|
||||
if status in self.FAILURE_STATUSES or is_failure_status(status, data):
|
||||
raise RuntimeError(self._format_task_failure(task_id, data))
|
||||
|
||||
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(
|
||||
"Grok Video 任务轮询超时;任务未被标记为失败,可用 task_id 继续查询。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}\n"
|
||||
f"status: {status or 'unknown'}\n"
|
||||
f"timeout: {timeout}s"
|
||||
f"Grok Video 轮询超时(request_id: {request_id},状态:{status or 'unknown'})。"
|
||||
)
|
||||
|
||||
await interruptible_sleep(min(interval, max(0.0, timeout - elapsed)))
|
||||
|
||||
async def _download_url_to_file(
|
||||
self,
|
||||
url: str,
|
||||
save_path: str,
|
||||
session: aiohttp.ClientSession,
|
||||
max_retries: int = 3,
|
||||
) -> str:
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
headers = None
|
||||
resolved_url = url
|
||||
|
||||
if url.startswith("data:"):
|
||||
if "," not in url:
|
||||
raise RuntimeError("Grok Video 下载失败:data URL 格式无效。")
|
||||
_, b64_data = url.split(",", 1)
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
with open(save_path, "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError("Grok Video 下载失败:保存后的文件为空。")
|
||||
return save_path
|
||||
|
||||
if url.startswith("/"):
|
||||
resolved_url = f"{self.base_url}{url}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
check_interrupt()
|
||||
async with session.get(
|
||||
resolved_url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
allow_redirects=True,
|
||||
) as response:
|
||||
if 200 <= response.status < 300:
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
with open(save_path, "wb") as f:
|
||||
async for chunk in response.content.iter_chunked(1024 * 1024):
|
||||
check_interrupt()
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError("Grok Video 下载失败:保存后的文件为空。")
|
||||
return save_path
|
||||
|
||||
last_status = response.status
|
||||
last_text = await response.text()
|
||||
if response.status not in RETRYABLE_STATUS_CODES or attempt >= max_retries:
|
||||
break
|
||||
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:下载重试 {attempt + 1}/{max_retries},{delay}s 后继续...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
raise RuntimeError(self._format_http_error("download_url", last_status, last_text))
|
||||
|
||||
async def download_video_async(
|
||||
self,
|
||||
task_id: str,
|
||||
save_path: str,
|
||||
session: aiohttp.ClientSession,
|
||||
) -> str:
|
||||
endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
|
||||
for attempt in range(4):
|
||||
check_interrupt()
|
||||
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
|
||||
if 200 <= response.status < 300:
|
||||
content_type = response.headers.get("Content-Type", "").lower()
|
||||
if "application/json" in content_type:
|
||||
data = await response.json(content_type=None)
|
||||
download_url = extract_video_url(data)
|
||||
if not download_url:
|
||||
raise RuntimeError(
|
||||
"Grok Video 下载失败:content 响应为 JSON,但未包含视频 URL。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
return await self._download_url_to_file(download_url, save_path, session)
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
with open(save_path, "wb") as f:
|
||||
async for chunk in response.content.iter_chunked(1024 * 1024):
|
||||
check_interrupt()
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError(
|
||||
"Grok Video 下载失败:保存后的文件为空。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
return save_path
|
||||
|
||||
last_status = response.status
|
||||
last_text = await response.text()
|
||||
if response.status not in RETRYABLE_STATUS_CODES or attempt >= 3:
|
||||
break
|
||||
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:content 下载重试 {attempt + 1}/3,{delay}s 后继续...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
|
||||
|
||||
def generate_video_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str,
|
||||
images: Optional[List[str]],
|
||||
output_dir: Optional[str] = None,
|
||||
save_path: Optional[str] = None,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
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():
|
||||
async def run_request() -> Dict[str, Any]:
|
||||
async with self._make_session() as session:
|
||||
endpoint = self.get_endpoint(operation)
|
||||
body = self.build_video_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
images=images,
|
||||
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,
|
||||
)
|
||||
|
||||
create_response = await self.create_video_async(body, session)
|
||||
task_id = self._extract_task_id(create_response) or ""
|
||||
if not task_id:
|
||||
raise RuntimeError(
|
||||
"Grok Video 未返回任务 ID。\n"
|
||||
f"endpoint: {self.CREATE_ENDPOINT}\n"
|
||||
f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
|
||||
)
|
||||
|
||||
print(f"Grok Video:任务已提交,任务ID:{task_id}")
|
||||
print("Grok Video:视频生成中...")
|
||||
status_response = await self.poll_video_status_async(
|
||||
task_id=task_id,
|
||||
session=session,
|
||||
poll_interval=poll_interval,
|
||||
timeout=timeout,
|
||||
progress_callback=progress_callback,
|
||||
print(f"Grok Video:正在提交{operation}任务…")
|
||||
created = await self._request_json(
|
||||
"POST", endpoint, session, json_body=body, timeout_seconds=180
|
||||
)
|
||||
|
||||
video_url = extract_video_url(status_response)
|
||||
print("Grok Video:视频生成完成,正在下载...")
|
||||
if save_path is None:
|
||||
resolved_output_dir = output_dir or os.getcwd()
|
||||
os.makedirs(resolved_output_dir, exist_ok=True)
|
||||
target_path = os.path.join(
|
||||
resolved_output_dir,
|
||||
f"{self._safe_task_filename(task_id)}.mp4",
|
||||
)
|
||||
else:
|
||||
target_path = save_path
|
||||
|
||||
if video_url:
|
||||
video_path = await self._download_url_to_file(video_url, target_path, session)
|
||||
else:
|
||||
video_path = await self.download_video_async(task_id, target_path, session)
|
||||
|
||||
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 {
|
||||
"task_id": task_id,
|
||||
"status": extract_status(status_response),
|
||||
"request_id": request_id,
|
||||
"video_path": video_path,
|
||||
"raw_json": {
|
||||
"create": create_response,
|
||||
"status": status_response,
|
||||
},
|
||||
"duration": video_data.get("duration"),
|
||||
"raw_json": {"create": created, "status": completed},
|
||||
}
|
||||
|
||||
return self.run_async_in_thread(_run())
|
||||
return self.run_async_in_thread(run_request())
|
||||
|
||||
Reference in New Issue
Block a user