Files
comfyui_o1key/nodes/K3_motion_control.py
T
Jony ba920f2b66 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.
2026-09-24 19:56:48 +08:00

397 lines
16 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
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 动作模仿",
}