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.
397 lines
16 KiB
Python
397 lines
16 KiB
Python
"""
|
||
K3 动作控制节点
|
||
|
||
用参考视频驱动参考图中人物动作，生成视频。
|
||
|
||
支持的模型：
|
||
- v3：标准动作控制，支持 5~30s 时长
|
||
- v2-6：标准动作控制，支持 5~30s 时长
|
||
- v3-t / v2-6-t：腾讯 Kling 网关渠道（保留兼容）
|
||
|
||
接口端点：
|
||
- 动作控制：POST /kling/v1/videos/motion-control
|
||
- 腾讯渠道：POST /v1/videos
|
||
"""
|
||
|
||
import asyncio
|
||
import io as _stdio
|
||
import json
|
||
import os
|
||
import struct
|
||
import tempfile
|
||
|
||
import aiohttp
|
||
|
||
from comfy_api.latest import io
|
||
|
||
from ..utils.config import get_api_key_or_raise, get_async_api_base_url, get_base_url_by_route
|
||
from ..utils.r2_uploader import upload_video, upload_image
|
||
from ..utils.image_utils import tensor_to_pil
|
||
from ..utils.http_error import async_request_with_retry
|
||
from ..utils.video_task import (
|
||
PollDeadline,
|
||
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,
|
||
)
|
||
|
||
try:
|
||
from comfy_api.latest import InputImpl
|
||
import folder_paths
|
||
_FOLDER_PATHS_OK = True
|
||
except Exception:
|
||
_FOLDER_PATHS_OK = False
|
||
|
||
|
||
# ── 常量 ──────────────────────────────────────────────────────────────────────
|
||
|
||
# 官方标准模型名映射
|
||
_STANDARD_MODELS = {
|
||
"v3": "kling-v3",
|
||
"v2-6": "kling-v2-6",
|
||
}
|
||
|
||
# 官方标准端点
|
||
_ENDPOINT_CREATE = "/kling/v1/videos/motion-control"
|
||
_ENDPOINT_STATUS = "/kling/v1/videos/motion-control/{task_id}"
|
||
|
||
# 腾讯 Kling 网关渠道（-t）：保留兼容
|
||
_ENDPOINT_T_CREATE = "/v1/videos"
|
||
_ENDPOINT_T_STATUS = "/v1/videos/{task_id}"
|
||
|
||
# -t 渠道模型名映射（服务端已部署，需在「模型倍率」各配一行 =1）
|
||
_MODEL_T_MAP = {
|
||
"v3-t": "kling-v3-motion-t",
|
||
"v2-6-t": "kling-v2-6-motion-t",
|
||
}
|
||
|
||
_POLL_INIT = 5
|
||
_POLL_MAX = 15
|
||
|
||
|
||
# ── 工具函数 ───────────────────────────────────────────────────────────────────
|
||
|
||
|
||
|
||
# ── 视频时长检测（纯标准库，跨平台） ──────────────────────────────────────────
|
||
|
||
def _parse_video_duration(data: bytes) -> float | None:
|
||
"""从 MP4/MOV 原始字节解析时长（秒）。读取 mvhd box。"""
|
||
idx = data.find(b"mvhd")
|
||
if idx == -1:
|
||
return None
|
||
box = data[idx + 4:]
|
||
if len(box) < 32:
|
||
return None
|
||
version = box[0]
|
||
try:
|
||
if version == 0:
|
||
timescale = struct.unpack(">I", box[12:16])[0]
|
||
duration = struct.unpack(">I", box[16:20])[0]
|
||
else: # version == 1
|
||
timescale = struct.unpack(">I", box[20:24])[0]
|
||
duration = struct.unpack(">Q", box[24:32])[0]
|
||
except struct.error:
|
||
return None
|
||
return (duration / timescale) if timescale > 0 else None
|
||
|
||
|
||
def _get_video_duration(reference_video) -> float | None:
|
||
"""从 ComfyUI VIDEO 对象获取视频时长（秒），失败返回 None。"""
|
||
try:
|
||
source = reference_video.get_stream_source()
|
||
if isinstance(source, str) and os.path.isfile(source):
|
||
with open(source, "rb") as f:
|
||
data = f.read()
|
||
elif isinstance(source, _stdio.BytesIO):
|
||
source.seek(0)
|
||
data = source.read()
|
||
else:
|
||
return None
|
||
return _parse_video_duration(data)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _validate_video_duration(reference_video, character_orientation: str):
|
||
"""校验视频时长，超限时抛出 ValueError。解析失败时静默跳过。"""
|
||
duration = _get_video_duration(reference_video)
|
||
if duration is None:
|
||
print("[K3 动作控制] 无法解析视频时长，跳过校验。")
|
||
return
|
||
limit = 10 if character_orientation == "image" else 30
|
||
print(f"[K3 动作控制] 检测到视频时长: {duration:.2f}s（限制: 3~{limit}s）")
|
||
if not (3 <= duration <= limit):
|
||
orientation_label = "图片" if character_orientation == "image" else "视频"
|
||
raise ValueError(
|
||
f"参考视频时长 {duration:.1f}s 不符合要求。\n"
|
||
f"角色朝向为「{orientation_label}」时，时长须在 3~{limit}s 之间。"
|
||
)
|
||
|
||
|
||
# ── 模型 DynamicCombo 选项构建 ─────────────────────────────────────────────────
|
||
|
||
def _build_model_input():
|
||
"""构建「模型」DynamicCombo。
|
||
|
||
支持的模型：
|
||
- v3：标准动作控制
|
||
- v2-6：标准动作控制
|
||
- v3-t / v2-6-t：腾讯网关渠道（保留兼容）
|
||
"""
|
||
def _duration_input():
|
||
return io.Combo.Input(
|
||
"时长", options=[5, 10, 15, 20, 25, 30], default=5,
|
||
tooltip="输出视频时长（秒）。须 ≥ 参考视频时长。",
|
||
)
|
||
return io.DynamicCombo.Input(
|
||
"模型",
|
||
options=[
|
||
io.DynamicCombo.Option("v3", [_duration_input()]),
|
||
io.DynamicCombo.Option("v2-6", [_duration_input()]),
|
||
io.DynamicCombo.Option("v3-t", []), # 腾讯网关，无时长参数
|
||
io.DynamicCombo.Option("v2-6-t", []), # 腾讯网关，无时长参数
|
||
],
|
||
tooltip="v3/v2-6：官方标准模型；v3-t/v2-6-t：腾讯网关渠道（兼容）。",
|
||
)
|
||
|
||
|
||
# ── 节点 ──────────────────────────────────────────────────────────────────────
|
||
|
||
class K3MotionControl(io.ComfyNode):
|
||
"""K3 动作控制 自研 —— 用参考视频驱动参考图人物动作"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="K3MotionControl",
|
||
display_name="K 动作模仿",
|
||
category="comfyui_o1key/KVideo",
|
||
inputs=[
|
||
io.Image.Input("参考图片"),
|
||
io.Video.Input("参考视频"),
|
||
_build_model_input(),
|
||
io.String.Input("提示词", multiline=True, default=""),
|
||
io.Combo.Input("模式", options=["720p", "1080p"], default="1080p"),
|
||
io.Combo.Input("角色朝向", options=["图片", "视频"], default="图片"),
|
||
io.Combo.Input("保留原声", options=["打开", "关闭"], default="打开"),
|
||
io.Int.Input("seed", default=0, min=0, max=2147483647,
|
||
tooltip="seed 仅控制节点是否重新运行，结果本身不可复现。"),
|
||
],
|
||
outputs=[io.Video.Output(display_name="视频")],
|
||
accept_all_inputs=True,
|
||
)
|
||
|
||
@classmethod
|
||
async def execute(cls, 参考图片, 参考视频, 模型, 提示词, 模式, 角色朝向, 保留原声, seed, **_kwargs) -> io.NodeOutput:
|
||
api_key = get_api_key_or_raise()
|
||
|
||
# ── 渠道判定（模型为 DynamicCombo dict）─────────────────────────
|
||
模型代号 = 模型["模型"]
|
||
is_t_channel = 模型代号 in _MODEL_T_MAP
|
||
时长 = int(模型.get("时长", 5)) # -t 渠道无此子输入
|
||
mode_api = "std" if 模式 == "720p" else "pro"
|
||
character_orientation = "image" if 角色朝向 == "图片" else "video"
|
||
keep_sound = "yes" if 保留原声 == "打开" else "no"
|
||
prompt = 提示词.strip()
|
||
|
||
base_url = get_base_url_by_route()
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
if len(prompt) > 2500:
|
||
raise ValueError(f"提示词长度（{len(prompt)}）超过上限 2500 个字符，请缩短后重试。")
|
||
|
||
# ── 进度条 ────────────────────────────────────────────────────
|
||
try:
|
||
from comfy.utils import ProgressBar
|
||
pbar = ProgressBar(100)
|
||
except Exception:
|
||
pbar = None
|
||
|
||
def _stage(s: str):
|
||
if s == "uploading":
|
||
print("[K3 动作控制] 上传图片/视频到 OSS...")
|
||
if pbar: pbar.update_absolute(0, 100)
|
||
elif s == "submitting":
|
||
print("[K3 动作控制] 提交任务...")
|
||
if pbar: pbar.update_absolute(10, 100)
|
||
elif s.startswith("submitted:"):
|
||
print(f"[K3 动作控制] 任务已提交 → {s.split(':', 1)[1]}")
|
||
if pbar: pbar.update_absolute(15, 100)
|
||
elif s == "downloading":
|
||
print("[K3 动作控制] 下载视频...")
|
||
if pbar: pbar.update_absolute(99, 100)
|
||
elif s == "done":
|
||
print("[K3 动作控制] 完成")
|
||
if pbar: pbar.update_absolute(100, 100)
|
||
|
||
def _progress(pct: int):
|
||
if pbar: pbar.update_absolute(15 + int(pct * 0.84), 100)
|
||
|
||
# ── 视频时长校验 ──────────────────────────────────────────────
|
||
# 参考视频时长约束（image≤10s / video≤30s，下限 3s）两渠道通用。
|
||
_validate_video_duration(参考视频, character_orientation)
|
||
|
||
# 「参考视频不得超过所选时长」仅标准渠道有意义：-t 渠道无时长入参
|
||
if not is_t_channel:
|
||
_dur = _get_video_duration(参考视频)
|
||
if _dur is not None and _dur > 时长 + 0.5:
|
||
raise ValueError(
|
||
f"参考视频时长 {_dur:.1f}s 超过所选时长 {时长}s。\n"
|
||
f"请将时长调整为 ≥{_dur:.0f}s 的档位，或更换更短的参考视频。"
|
||
)
|
||
|
||
# ── 图片 & 视频上传 OSS → 获取公网 URL ────────────────────────
|
||
_stage("uploading")
|
||
check_interrupt()
|
||
pil_list = tensor_to_pil(参考图片)
|
||
img = pil_list[0]
|
||
# 转换为 RGBA 以支持透明通道，PNG 格式上传
|
||
if img.mode not in ("RGBA", "RGB"):
|
||
img = img.convert("RGBA" if "A" in img.mode or img.mode == "LA" else "RGB")
|
||
image_url = await upload_image(img, base_url=base_url)
|
||
check_interrupt()
|
||
video_url = await upload_video(参考视频, base_url=base_url)
|
||
|
||
# ── 构建请求体 ────────────────────────────────────────────────
|
||
if is_t_channel:
|
||
# 腾讯 Kling 网关渠道：动作控制专有字段进 metadata 透传（PascalCase），
|
||
# 顶层只放标准字段。无 duration / mode 由网关按模型处理。
|
||
body = {
|
||
"model": _MODEL_T_MAP[模型代号],
|
||
"prompt": prompt or "动作与参考视频保持一致", # 网关强制非空
|
||
"image": image_url,
|
||
"metadata": {
|
||
"Video": video_url,
|
||
"CharacterOrientation": character_orientation,
|
||
"KeepOriginalSound": keep_sound,
|
||
"Mode": mode_api, # 720p→std / 1080p→pro
|
||
},
|
||
}
|
||
create_path = _ENDPOINT_T_CREATE
|
||
status_path = _ENDPOINT_T_STATUS
|
||
else:
|
||
# 标准渠道：使用官方标准接口，参数扁平传递
|
||
# 获取实际的模型名（v3 → kling-v3）
|
||
actual_model_name = _STANDARD_MODELS.get(模型代号, f"kling-{模型代号}")
|
||
|
||
body = {
|
||
"model_name": actual_model_name,
|
||
"image_url": image_url,
|
||
"video_url": video_url,
|
||
"character_orientation": character_orientation,
|
||
"mode": mode_api,
|
||
"keep_original_sound": keep_sound,
|
||
"duration": str(时长),
|
||
}
|
||
if prompt:
|
||
body["prompt"] = prompt
|
||
create_path = _ENDPOINT_CREATE
|
||
status_path = _ENDPOINT_STATUS
|
||
|
||
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
|
||
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k3_motion_")
|
||
|
||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
|
||
# 1. 提交任务
|
||
check_interrupt()
|
||
_stage("submitting")
|
||
create_url = f"{base_url}{create_path}"
|
||
resp = await run_with_interrupt(async_request_with_retry(
|
||
session, "POST", create_url,
|
||
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
|
||
headers=headers, prefix="K3 动作控制提交: "
|
||
))
|
||
check_interrupt()
|
||
text = await resp.text()
|
||
create_resp = json.loads(text)
|
||
|
||
# task_id 兼容扁平结构和 data 嵌套结构
|
||
task_id = (
|
||
create_resp.get("task_id")
|
||
or create_resp.get("id")
|
||
or create_resp.get("data", {}).get("task_id")
|
||
)
|
||
if not task_id:
|
||
raise RuntimeError(f"API 未返回任务 ID，响应：{create_resp}")
|
||
_stage(f"submitted:{task_id}")
|
||
|
||
# 2. 轮询
|
||
status_url = f"{base_url}{status_path.format(task_id=task_id)}"
|
||
interval = _POLL_INIT
|
||
video_result_url = None
|
||
deadline = PollDeadline(label="K3 动作控制")
|
||
|
||
while True:
|
||
deadline.check()
|
||
await interruptible_sleep(interval)
|
||
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("message") or text
|
||
except Exception:
|
||
msg = text
|
||
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
|
||
sr = json.loads(text)
|
||
|
||
# 兼容扁平结构和 data 嵌套结构
|
||
data = sr.get("data", sr)
|
||
status = extract_status(sr)
|
||
|
||
pct = extract_progress(sr)
|
||
print(f"[K3 动作控制] 生成中 {pct}%")
|
||
_progress(pct)
|
||
|
||
if is_success_status(status):
|
||
video_result_url = extract_video_url(sr)
|
||
break
|
||
elif is_failure_status(status, sr):
|
||
err_msg = extract_error_message(sr)
|
||
raise RuntimeError(f"K3 动作控制生成失败：{err_msg}")
|
||
|
||
interval = min(interval * 1.3, _POLL_MAX)
|
||
|
||
if not video_result_url:
|
||
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
|
||
|
||
# 3. 下载视频（抗超时 / 断点续传 / 无限重试 / 可取消）
|
||
check_interrupt()
|
||
_stage("downloading")
|
||
os.close(tmp_fd)
|
||
await download_video_to_file(
|
||
session, video_result_url, save_path, label="K3 动作控制",
|
||
)
|
||
|
||
_stage("done")
|
||
|
||
if _FOLDER_PATHS_OK:
|
||
return io.NodeOutput(InputImpl.VideoFromFile(save_path))
|
||
return io.NodeOutput(save_path)
|
||
|
||
|
||
# ── 节点注册 ──────────────────────────────────────────────────────────────────
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"K3MotionControl": K3MotionControl,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"K3MotionControl": "K 动作模仿",
|
||
}
|