- 新增 K3VideoFirstLast 节点(首尾帧 K3 自研):基于 K3Video 去掉分镜,新增可选尾帧输入,尾帧通过 metadata.image_tail 传递 - nano_banana_pro / batch_nano_banana_pro:将参数「代理加速」重命名为「代理端口(如7897)」 - gemini_client:新增 build_proxy_url、_estimate_body_size、_scale_images_to_fit 工具方法 - base_client:小幅优化
308 lines
12 KiB
Python
308 lines
12 KiB
Python
"""
|
||
首尾帧 K3 自研节点
|
||
基于 K3 图生视频 自研,去掉分镜功能,新增尾帧可选输入。
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
|
||
import aiohttp
|
||
|
||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||
|
||
try:
|
||
from comfy_api.latest import InputImpl
|
||
import folder_paths
|
||
_FOLDER_PATHS_OK = True
|
||
except Exception:
|
||
_FOLDER_PATHS_OK = False
|
||
|
||
|
||
# ── 常量 ──────────────────────────────────────────────────────────────────────
|
||
|
||
_MODEL_BASE = "kling-v3"
|
||
_MODES = ["标准", "专家", "4K"]
|
||
_MODE_MAP = {"标准": "std", "专家": "pro", "4K": "4k"}
|
||
|
||
_ENDPOINT_CREATE = "/v1/video/generations"
|
||
_ENDPOINT_STATUS = "/v1/video/generations/{task_id}"
|
||
|
||
_POLL_INIT = 3
|
||
_POLL_MAX = 15
|
||
|
||
|
||
# ── 工具函数 ───────────────────────────────────────────────────────────────────
|
||
|
||
def _get_video_dir() -> str:
|
||
if _FOLDER_PATHS_OK:
|
||
base = folder_paths.get_output_directory()
|
||
else:
|
||
plugin = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
base = os.path.join(os.path.dirname(os.path.dirname(plugin)), "output")
|
||
d = os.path.join(base, "video")
|
||
os.makedirs(d, exist_ok=True)
|
||
return d
|
||
|
||
|
||
def _next_counter(directory: str, prefix: str) -> int:
|
||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||
max_n = 0
|
||
if os.path.exists(directory):
|
||
for f in os.listdir(directory):
|
||
m = pattern.match(f)
|
||
if m:
|
||
max_n = max(max_n, int(m.group(1)))
|
||
return max_n + 1
|
||
|
||
|
||
def _prepare_image_base64(tensor) -> str:
|
||
"""转换并校验图片,不符合约束时自动等比缩放后返回 base64。"""
|
||
import io
|
||
import base64
|
||
|
||
pil_list = tensor_to_pil(tensor)
|
||
img = pil_list[0].convert("RGB")
|
||
w, h = img.size
|
||
|
||
# 1. 宽高比校验
|
||
ratio = w / h
|
||
if ratio < 1 / 2.5 or ratio > 2.5:
|
||
raise RuntimeError(
|
||
f"图片宽高比 {w}:{h}({ratio:.2f})超出允许范围 1:2.5 ~ 2.5:1,请裁剪后重试。"
|
||
)
|
||
|
||
# 2. 最小尺寸:任意边 < 300px 时等比放大
|
||
if w < 300 or h < 300:
|
||
scale = max(300 / w, 300 / h)
|
||
img = img.resize((int(w * scale), int(h * scale)), resample=1)
|
||
|
||
# 3. 文件大小:循环等比缩小直到 ≤ 10MB
|
||
MAX_BYTES = 10 * 1024 * 1024
|
||
for _ in range(20):
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
if buf.tell() <= MAX_BYTES:
|
||
break
|
||
scale = (MAX_BYTES / buf.tell()) ** 0.5 * 0.95
|
||
new_w = int(img.width * scale)
|
||
new_h = int(img.height * scale)
|
||
if new_w < 300 or new_h < 300:
|
||
raise RuntimeError(
|
||
f"图片压缩至 10MB 以内后尺寸({new_w}x{new_h})低于最小限制 300px,无法同时满足两项约束。"
|
||
)
|
||
img = img.resize((new_w, new_h), resample=1)
|
||
else:
|
||
raise RuntimeError("图片经过 20 次缩放仍超过 10MB,请检查原始图片。")
|
||
|
||
buf.seek(0)
|
||
return base64.b64encode(buf.read()).decode("utf-8")
|
||
|
||
|
||
# ── 节点 ──────────────────────────────────────────────────────────────────────
|
||
|
||
class K3VideoFirstLast:
|
||
"""首尾帧 K3 自研"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"起始帧": ("IMAGE",),
|
||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||
"负向提示词": ("STRING", {"multiline": True, "default": ""}),
|
||
"时长": ([5, 10, 15], {"default": 5}),
|
||
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||
"模式": (_MODES, {"default": "标准"}),
|
||
"seed": ("INT", {
|
||
"default": 0, "min": 0, "max": 2147483647,
|
||
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
|
||
}),
|
||
},
|
||
"optional": {
|
||
"尾帧": ("IMAGE", {"tooltip": "可选。传入后将作为视频尾帧参考。"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("VIDEO",)
|
||
RETURN_NAMES = ("视频",)
|
||
FUNCTION = "generate"
|
||
CATEGORY = "comfyui_o1key/KVideo"
|
||
|
||
async def generate(self, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, seed, 尾帧=None):
|
||
api_key = get_api_key_or_raise()
|
||
base_url = get_api_base_url()
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
voice = "voice" if 生成音频 == "打开" else "novoice"
|
||
mode_api = _MODE_MAP[模式]
|
||
if mode_api == "4k":
|
||
model_name = f"{_MODEL_BASE}-4k-{时长}s"
|
||
else:
|
||
model_name = f"{_MODEL_BASE}-{mode_api}-{时长}s-{voice}"
|
||
|
||
if not 提示词.strip():
|
||
raise RuntimeError("提示词不能为空。")
|
||
|
||
# ── 构建请求体 ────────────────────────────────────────────────
|
||
body: dict = {
|
||
"model": model_name,
|
||
"prompt": 提示词.strip(),
|
||
"mode": mode_api,
|
||
"duration": 时长,
|
||
"image": _prepare_image_base64(起始帧),
|
||
}
|
||
|
||
if 负向提示词.strip():
|
||
body["negative_prompt"] = 负向提示词.strip()
|
||
|
||
# metadata:尾帧 + 音频
|
||
metadata: dict = {}
|
||
if 尾帧 is not None:
|
||
metadata["image_tail"] = _prepare_image_base64(尾帧)
|
||
if 生成音频 == "打开":
|
||
metadata["sound"] = "on"
|
||
if metadata:
|
||
body["metadata"] = metadata
|
||
|
||
# generate_audio 字段(非 metadata 路径)
|
||
if 生成音频 == "打开" and not metadata.get("sound"):
|
||
body["generate_audio"] = True
|
||
|
||
# ── 进度条 ────────────────────────────────────────────────────
|
||
try:
|
||
from comfy.utils import ProgressBar
|
||
pbar = ProgressBar(100)
|
||
except Exception:
|
||
pbar = None
|
||
|
||
def _stage(s: str):
|
||
if s == "submitting":
|
||
print("[K3 首尾帧] 提交中...")
|
||
if pbar: pbar.update_absolute(0, 100)
|
||
elif s.startswith("submitted:"):
|
||
print(f"[K3 首尾帧] 任务已提交 → {s.split(':', 1)[1]}")
|
||
if pbar: pbar.update_absolute(5, 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(5 + int(pct * 0.94), 100)
|
||
|
||
# ── 保存路径 ──────────────────────────────────────────────────
|
||
video_dir = _get_video_dir()
|
||
counter = _next_counter(video_dir, "k3fl")
|
||
save_path = os.path.join(video_dir, f"k3fl_{counter:05d}.mp4")
|
||
|
||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
|
||
# 1. 提交
|
||
_stage("submitting")
|
||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||
async with session.post(create_url, json=body, 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"K3 首尾帧提交失败 ({resp.status}): {msg}")
|
||
create_resp = json.loads(text)
|
||
|
||
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}{_ENDPOINT_STATUS.format(task_id=task_id)}"
|
||
interval = _POLL_INIT
|
||
video_url = None
|
||
|
||
while True:
|
||
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}")
|
||
sr = json.loads(text)
|
||
|
||
data = sr.get("data", sr)
|
||
status = (data.get("status") or sr.get("status") or "").lower()
|
||
|
||
pct_raw = data.get("progress", 0)
|
||
try:
|
||
pct = int(str(pct_raw).rstrip("%").strip())
|
||
except (ValueError, AttributeError):
|
||
pct = 0
|
||
print(f"[K3 首尾帧] 生成中 {pct}%")
|
||
_progress(pct)
|
||
|
||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||
video_url = (
|
||
data.get("video_url")
|
||
or data.get("result_url")
|
||
or data.get("url")
|
||
or (data.get("result", {}) or {}).get("url")
|
||
or sr.get("video_url")
|
||
or sr.get("url")
|
||
)
|
||
break
|
||
if status in ("failed", "fail"):
|
||
err_info = data.get("error") or sr.get("error") or {}
|
||
err_msg = (err_info.get("message", "未知错误")
|
||
if isinstance(err_info, dict) else str(err_info))
|
||
raise RuntimeError(f"K3 首尾帧生成失败:{err_msg}")
|
||
|
||
await asyncio.sleep(interval)
|
||
interval = min(interval * 1.5, _POLL_MAX)
|
||
|
||
if not video_url:
|
||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||
|
||
# 3. 下载
|
||
_stage("downloading")
|
||
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):
|
||
f.write(chunk)
|
||
|
||
_stage("done")
|
||
|
||
if _FOLDER_PATHS_OK:
|
||
return (InputImpl.VideoFromFile(save_path),)
|
||
return (save_path,)
|
||
|
||
|
||
# ── 节点注册 ──────────────────────────────────────────────────────────────────
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"K3VideoFirstLast": K3VideoFirstLast,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"K3VideoFirstLast": "首尾帧 K3 自研",
|
||
}
|