286 lines
11 KiB
Python
286 lines
11 KiB
Python
"""
|
|
Kling 视频生成 API 客户端
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
import aiohttp
|
|
|
|
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
|
from ..utils.http_error import async_request_with_retry
|
|
from ..utils.video_task import (
|
|
check_interrupt,
|
|
extract_error_message,
|
|
extract_progress,
|
|
extract_status,
|
|
interruptible_sleep,
|
|
is_failure_status,
|
|
is_success_status,
|
|
run_with_interrupt,
|
|
)
|
|
|
|
|
|
class KlingClient:
|
|
"""Kling 视频生成客户端"""
|
|
|
|
ENDPOINTS = {
|
|
"image2video": "/kling/v1/videos/image2video",
|
|
"text2video": "/kling/v1/videos/text2video",
|
|
"motion_control": "/kling/v1/videos/motion-control",
|
|
}
|
|
|
|
# new API 三段式端点（动作控制走这里）
|
|
NEW_API_CREATE = "/v1/videos"
|
|
NEW_API_STATUS = "/v1/videos/{video_id}"
|
|
NEW_API_CONTENT = "/v1/videos/{video_id}/content"
|
|
|
|
POLL_INITIAL_INTERVAL = 3
|
|
POLL_MAX_INTERVAL = 15
|
|
|
|
def __init__(self):
|
|
self.api_key = get_api_key_or_raise()
|
|
self.base_url = get_api_base_url()
|
|
|
|
def _headers(self) -> Dict[str, str]:
|
|
return {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
# ── 提交任务 ──────────────────────────────────────────────────────
|
|
|
|
async def create_video_async(
|
|
self,
|
|
endpoint_type: str,
|
|
body: Dict[str, Any],
|
|
session: aiohttp.ClientSession,
|
|
) -> Dict[str, Any]:
|
|
url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}"
|
|
|
|
check_interrupt()
|
|
resp = await run_with_interrupt(async_request_with_retry(
|
|
session, "POST", url, json=body, headers=self._headers(), prefix="Kling 提交: "
|
|
))
|
|
check_interrupt()
|
|
text = await resp.text()
|
|
return json.loads(text)
|
|
|
|
# ── 轮询状态 ──────────────────────────────────────────────────────
|
|
|
|
async def poll_status_async(
|
|
self,
|
|
task_id: str,
|
|
endpoint_type: str,
|
|
session: aiohttp.ClientSession,
|
|
on_progress: Optional[Callable[[int], None]] = None,
|
|
) -> Dict[str, Any]:
|
|
url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}/{task_id}"
|
|
interval = self.POLL_INITIAL_INTERVAL
|
|
|
|
while True:
|
|
check_interrupt()
|
|
async with session.get(url, headers=self._headers()) as resp:
|
|
text = await resp.text()
|
|
if resp.status != 200:
|
|
raise RuntimeError(f"状态查询失败 ({resp.status}): {text}")
|
|
result = json.loads(text)
|
|
|
|
data = result.get("data", {})
|
|
inner_data = data.get("data", {}) if isinstance(data, dict) else {}
|
|
status = extract_status(result)
|
|
|
|
progress_pct = extract_progress(result)
|
|
|
|
print(f"[视频生成] 生成中 {progress_pct}%")
|
|
|
|
if on_progress:
|
|
on_progress(progress_pct)
|
|
|
|
if is_success_status(status):
|
|
return result
|
|
elif is_failure_status(status, result):
|
|
error_msg = extract_error_message(result)
|
|
raise RuntimeError(f"生成失败：{error_msg}")
|
|
|
|
await interruptible_sleep(interval)
|
|
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
|
|
|
# ── 下载视频 ──────────────────────────────────────────────────────
|
|
|
|
async def download_video_async(
|
|
self,
|
|
video_url: str,
|
|
save_path: str,
|
|
session: aiohttp.ClientSession,
|
|
) -> str:
|
|
print("[视频生成] 下载视频...")
|
|
check_interrupt()
|
|
async with session.get(video_url, allow_redirects=True) as resp:
|
|
if resp.status != 200:
|
|
raise RuntimeError(f"视频下载失败 ({resp.status})")
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
with open(save_path, "wb") as f:
|
|
async for chunk in resp.content.iter_chunked(8192):
|
|
check_interrupt()
|
|
f.write(chunk)
|
|
return save_path
|
|
|
|
# ── 异步入口（供节点调用）────────────────────────────────────────
|
|
|
|
async def generate_async(
|
|
self,
|
|
endpoint_type: str,
|
|
body: Dict[str, Any],
|
|
save_path: str,
|
|
on_stage: Optional[Callable[[str], None]] = None,
|
|
on_progress: Optional[Callable[[int], None]] = None,
|
|
) -> str:
|
|
"""提交 → 轮询 → 下载，返回本地文件路径"""
|
|
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
if on_stage:
|
|
on_stage("submitting")
|
|
|
|
result = await self.create_video_async(endpoint_type, body, session)
|
|
# 提交响应结构：result.data.task_id
|
|
task_id = result.get("task_id") or result.get("data", {}).get("task_id")
|
|
if not task_id:
|
|
raise RuntimeError(f"API 未返回任务 ID，响应：{result}")
|
|
if on_stage:
|
|
on_stage(f"submitted:{task_id}")
|
|
|
|
if on_stage:
|
|
on_stage("polling")
|
|
final = await self.poll_status_async(
|
|
task_id, endpoint_type, session, on_progress=on_progress
|
|
)
|
|
|
|
# 兼容多种URL路径
|
|
# 响应结构：result.data.result_url 或 result.data.data.task_result.videos[0].url
|
|
data = final.get("data", {})
|
|
inner_data = data.get("data", {}) if isinstance(data, dict) else {}
|
|
video_url = (
|
|
data.get("result_url") or
|
|
final.get("url") or
|
|
final.get("video_url") or
|
|
(inner_data.get("task_result", {}).get("videos", [{}])[0].get("url")
|
|
if inner_data.get("task_result", {}).get("videos") else None)
|
|
)
|
|
if not video_url:
|
|
raise RuntimeError(f"API 未返回视频 URL，响应：{final}")
|
|
|
|
if on_stage:
|
|
on_stage("downloading")
|
|
path = await self.download_video_async(video_url, save_path, session)
|
|
|
|
if on_stage:
|
|
on_stage("done")
|
|
return path
|
|
|
|
# ── 动作控制：走 new API 三段式流程 ──────────────────────────────
|
|
|
|
async def motion_control_async(
|
|
self,
|
|
body: Dict[str, Any],
|
|
save_path: str,
|
|
on_stage: Optional[Callable[[str], None]] = None,
|
|
on_progress: Optional[Callable[[int], None]] = None,
|
|
) -> str:
|
|
"""
|
|
动作控制专用入口：
|
|
POST /v1/videos → GET /v1/videos/{id} → GET /v1/videos/{id}/content
|
|
body 字段与 Kling 官方动作控制接口一致（image_url/video_url/prompt/...）。
|
|
"""
|
|
headers = {"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json"}
|
|
interval = self.POLL_INITIAL_INTERVAL
|
|
|
|
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
|
|
# 1. 提交
|
|
check_interrupt()
|
|
if on_stage:
|
|
on_stage("submitting")
|
|
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
|
|
resp = await run_with_interrupt(async_request_with_retry(
|
|
session, "POST", create_url, json=body, headers=headers, prefix="Kling 动作控制提交: "
|
|
))
|
|
check_interrupt()
|
|
text = await resp.text()
|
|
create_resp = json.loads(text)
|
|
|
|
video_id = create_resp.get("id")
|
|
if not video_id:
|
|
raise RuntimeError(f"API 未返回视频 ID，响应：{create_resp}")
|
|
if on_stage:
|
|
on_stage(f"submitted:{video_id}")
|
|
|
|
# 2. 轮询
|
|
status_url = f"{self.base_url}{self.NEW_API_STATUS.format(video_id=video_id)}"
|
|
while True:
|
|
check_interrupt()
|
|
async with session.get(status_url, headers=headers) as resp:
|
|
text = await resp.text()
|
|
if resp.status != 200:
|
|
try:
|
|
err = json.loads(text)
|
|
msg = err.get("error", {}).get("message") or err.get("message") or text
|
|
except Exception:
|
|
msg = text
|
|
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
|
|
status_resp = json.loads(text)
|
|
|
|
status = extract_status(status_resp)
|
|
progress_pct = extract_progress(status_resp)
|
|
|
|
print(f"[动作控制] 生成中 {progress_pct}%")
|
|
if on_progress:
|
|
on_progress(progress_pct)
|
|
|
|
if is_success_status(status):
|
|
break
|
|
if is_failure_status(status, status_resp):
|
|
error_msg = extract_error_message(status_resp)
|
|
raise RuntimeError(f"动作控制生成失败：{error_msg}")
|
|
|
|
await interruptible_sleep(interval)
|
|
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
|
|
|
# 3. 下载
|
|
check_interrupt()
|
|
if on_stage:
|
|
on_stage("downloading")
|
|
content_url = f"{self.base_url}{self.NEW_API_CONTENT.format(video_id=video_id)}"
|
|
async with session.get(content_url, headers=headers,
|
|
allow_redirects=True) as resp:
|
|
if resp.status != 200:
|
|
raise RuntimeError(f"视频下载失败 ({resp.status})")
|
|
content_type = resp.headers.get("Content-Type", "")
|
|
if "application/json" in content_type:
|
|
data = await resp.json()
|
|
download_url = data.get("url") or data.get("download_url")
|
|
if not download_url:
|
|
raise RuntimeError("视频下载失败：响应中未找到下载链接")
|
|
async with session.get(download_url) as dl_resp:
|
|
if dl_resp.status != 200:
|
|
raise RuntimeError(f"从下载链接获取视频失败 ({dl_resp.status})")
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
with open(save_path, "wb") as f:
|
|
async for chunk in dl_resp.content.iter_chunked(8192):
|
|
check_interrupt()
|
|
f.write(chunk)
|
|
else:
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
with open(save_path, "wb") as f:
|
|
async for chunk in resp.content.iter_chunked(8192):
|
|
check_interrupt()
|
|
f.write(chunk)
|
|
|
|
if on_stage:
|
|
on_stage("done")
|
|
return save_path
|