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.
1063 lines
41 KiB
Python
1063 lines
41 KiB
Python
"""Parallel background jobs for the panel-style o1key video generator."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from collections import deque
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Callable
|
||
|
||
import aiohttp
|
||
from PIL import Image
|
||
|
||
from ..clients.seedance_client import SeedanceClient
|
||
from ..clients.seedance_element_client import SeedanceElementClient
|
||
from .config import get_base_url_by_route
|
||
from .http_error import format_o1key_video_error
|
||
from .o1key_image_save import normalize_save_location
|
||
from .o1key_video_catalog import (
|
||
SEEDANCE_REFERENCE_AUDIO_MAX_BYTES,
|
||
SEEDANCE_REFERENCE_IMAGE_MAX_BYTES,
|
||
SEEDANCE_REFERENCE_VIDEO_MAX_BYTES,
|
||
normalize_seedance_parameters,
|
||
public_video_capabilities,
|
||
validate_seedance_media_counts,
|
||
validate_seedance_reference_dimensions,
|
||
)
|
||
from .r2_uploader import upload_audio, upload_image, upload_video
|
||
from .seedance_assets import (
|
||
SeedanceAssetService,
|
||
seedance_asset_fingerprint,
|
||
seedance_asset_request_type,
|
||
)
|
||
|
||
|
||
_BATCH_ID_RE = re.compile(r"^[A-Za-z0-9_-]{8,96}$")
|
||
_ASSET_ID_RE = re.compile(r"^[A-Za-z0-9._:-]{1,256}$")
|
||
_UNSAFE_FILENAME = re.compile(r'[<>:"/\\|?*\x00-\x1f]+')
|
||
_MEDIA_RULES = {
|
||
"first_frame": ({".png", ".jpg", ".jpeg", ".webp", ".bmp"}, SEEDANCE_REFERENCE_IMAGE_MAX_BYTES),
|
||
"last_frame": ({".png", ".jpg", ".jpeg", ".webp", ".bmp"}, SEEDANCE_REFERENCE_IMAGE_MAX_BYTES),
|
||
"reference_images": ({".png", ".jpg", ".jpeg", ".webp", ".bmp"}, SEEDANCE_REFERENCE_IMAGE_MAX_BYTES),
|
||
"reference_videos": ({".mp4", ".mov"}, SEEDANCE_REFERENCE_VIDEO_MAX_BYTES),
|
||
"reference_audios": ({".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg"}, SEEDANCE_REFERENCE_AUDIO_MAX_BYTES),
|
||
}
|
||
_LIST_MEDIA_ROLES = {"reference_images", "reference_videos", "reference_audios"}
|
||
_IMAGE_MEDIA_ROLES = {"first_frame", "last_frame", "reference_images"}
|
||
_VIDEO_MEDIA_ROLES = {"reference_videos"}
|
||
_IMAGE_ROLE_LABELS = {
|
||
"first_frame": "首帧图片",
|
||
"last_frame": "尾帧图片",
|
||
"reference_images": "参考图片",
|
||
}
|
||
_VIDEO_ROLE_LABELS = {
|
||
"reference_videos": "参考视频",
|
||
}
|
||
_SAVE_LOCK = threading.Lock()
|
||
_MAX_HISTORY_ITEMS = 128
|
||
|
||
|
||
def _canonical_batch_id(value: Any) -> str:
|
||
batch_id = str(value or "").strip()
|
||
if not _BATCH_ID_RE.fullmatch(batch_id):
|
||
raise ValueError("视频任务 ID 无效")
|
||
return batch_id
|
||
|
||
|
||
def _safe_node_id(value: Any, label: str) -> int:
|
||
try:
|
||
node_id = int(value)
|
||
except (TypeError, ValueError):
|
||
raise ValueError(f"{label}无效") from None
|
||
if node_id < 0:
|
||
raise ValueError(f"{label}无效")
|
||
return node_id
|
||
|
||
|
||
def _parse_json_object(value: Any, label: str) -> dict[str, Any]:
|
||
if isinstance(value, str):
|
||
try:
|
||
value = json.loads(value or "{}")
|
||
except json.JSONDecodeError:
|
||
raise ValueError(f"{label}不是有效 JSON") from None
|
||
if value is None:
|
||
return {}
|
||
if not isinstance(value, dict):
|
||
raise ValueError(f"{label}必须是对象")
|
||
return value
|
||
|
||
|
||
def _clean_input_descriptor(value: Any, label: str) -> dict[str, str]:
|
||
if not isinstance(value, dict):
|
||
raise ValueError(f"{label}描述符无效")
|
||
name = os.path.basename(str(value.get("name") or "").strip())
|
||
subfolder = str(value.get("subfolder") or "").strip().replace("\\", "/")
|
||
if not name or value.get("type", "input") != "input":
|
||
raise ValueError(f"{label}必须来自 ComfyUI input 目录")
|
||
if subfolder.startswith("/") or any(part == ".." for part in subfolder.split("/")):
|
||
raise ValueError(f"{label}子目录无效")
|
||
return {"name": name, "subfolder": subfolder, "type": "input"}
|
||
|
||
|
||
def _clean_asset_ids(value: Any, label: str) -> list[str]:
|
||
if value is None:
|
||
return []
|
||
if isinstance(value, str):
|
||
value = re.split(r"[,，\n]", value)
|
||
if not isinstance(value, list):
|
||
raise ValueError(f"{label}必须是数组或逗号分隔文本")
|
||
result = []
|
||
for item in value:
|
||
asset_id = str(item or "").strip().removeprefix("asset://")
|
||
if not asset_id:
|
||
continue
|
||
if not _ASSET_ID_RE.fullmatch(asset_id):
|
||
raise ValueError(f"{label}包含无效素材 ID")
|
||
result.append(asset_id)
|
||
return result
|
||
|
||
|
||
def normalize_video_job_payload(payload: Any) -> dict[str, Any]:
|
||
if not isinstance(payload, dict):
|
||
raise ValueError("请求体必须是对象")
|
||
normalized = normalize_seedance_parameters(payload)
|
||
media_raw = _parse_json_object(
|
||
payload.get("media") if "media" in payload else payload.get("media_manifest"),
|
||
"媒体清单",
|
||
)
|
||
assets_raw = _parse_json_object(
|
||
payload.get("assets") if "assets" in payload else payload.get("asset_manifest"),
|
||
"素材 ID 清单",
|
||
)
|
||
|
||
media: dict[str, Any] = {}
|
||
for role in _MEDIA_RULES:
|
||
raw_value = media_raw.get(role)
|
||
if role in _LIST_MEDIA_ROLES:
|
||
if raw_value is None:
|
||
raw_value = []
|
||
if not isinstance(raw_value, list):
|
||
raise ValueError(f"{role} 必须是数组")
|
||
media[role] = [
|
||
_clean_input_descriptor(item, role) for item in raw_value
|
||
]
|
||
else:
|
||
media[role] = (
|
||
_clean_input_descriptor(raw_value, role) if raw_value else None
|
||
)
|
||
|
||
raw_image_assets = (
|
||
assets_raw.get("images")
|
||
if "images" in assets_raw
|
||
else assets_raw.get("persons")
|
||
)
|
||
assets = {
|
||
"images": _clean_asset_ids(raw_image_assets, "图片素材 ID"),
|
||
"videos": _clean_asset_ids(assets_raw.get("videos"), "视频素材 ID"),
|
||
"audios": _clean_asset_ids(assets_raw.get("audios"), "音频素材 ID"),
|
||
}
|
||
|
||
mode = normalized["generation_mode"]
|
||
first_frame = media["first_frame"]
|
||
last_frame = media["last_frame"]
|
||
reference_count = len(media["reference_images"])
|
||
video_count = len(media["reference_videos"])
|
||
audio_count = len(media["reference_audios"])
|
||
asset_count = sum(len(items) for items in assets.values())
|
||
# Released requests used an assets object without a mode. Preserve those as
|
||
# manual-ID requests while new panel submissions always send the mode.
|
||
legacy_asset_request = "asset_creation_mode" not in payload and bool(asset_count)
|
||
if legacy_asset_request:
|
||
normalized["asset_creation_mode"] = "manual"
|
||
manual_assets = normalized["asset_creation_mode"] == "manual"
|
||
|
||
if mode == "text":
|
||
if manual_assets:
|
||
raise ValueError("文生视频模式不需要手动素材创建")
|
||
if first_frame or last_frame or reference_count or video_count or audio_count or asset_count:
|
||
raise ValueError("文生视频模式不能携带参考素材")
|
||
elif manual_assets:
|
||
if (
|
||
(first_frame or last_frame or reference_count or video_count or audio_count)
|
||
and not legacy_asset_request
|
||
):
|
||
raise ValueError("手动素材创建模式只能填写素材 ID，不能同时上传参考素材")
|
||
if mode == "first_frame":
|
||
if len(assets["images"]) != 1 or assets["videos"] or assets["audios"]:
|
||
raise ValueError("首帧图生视频的手动模式必须且只能填写一个图片素材 ID")
|
||
elif mode == "first_last_frame":
|
||
if len(assets["images"]) != 2 or assets["videos"] or assets["audios"]:
|
||
raise ValueError("首尾帧生视频的手动模式必须且只能填写两个图片素材 ID")
|
||
elif not asset_count:
|
||
raise ValueError("多模态手动模式至少需要填写一个素材 ID")
|
||
else:
|
||
if asset_count:
|
||
raise ValueError("自动素材创建模式不能填写手动素材 ID")
|
||
if mode == "first_frame":
|
||
if not first_frame or last_frame or reference_count or video_count or audio_count:
|
||
raise ValueError("首帧图生视频必须且只能提供一张首帧图片")
|
||
elif mode == "first_last_frame":
|
||
if not first_frame or not last_frame or reference_count or video_count or audio_count:
|
||
raise ValueError("首尾帧生视频必须且只能提供首帧和尾帧图片")
|
||
elif not (
|
||
normalized["prompt"]
|
||
or reference_count
|
||
or video_count
|
||
or audio_count
|
||
):
|
||
raise ValueError("多模态模式下提示词和参考素材不能同时为空")
|
||
|
||
normalized.pop("capabilities")
|
||
image_total = reference_count + len(assets["images"])
|
||
if mode == "first_frame":
|
||
image_total = 1
|
||
elif mode == "first_last_frame":
|
||
image_total = 2
|
||
video_total = video_count + len(assets["videos"])
|
||
audio_total = audio_count + len(assets["audios"])
|
||
validate_seedance_media_counts(
|
||
normalized["model"],
|
||
image_total,
|
||
video_total,
|
||
audio_total,
|
||
)
|
||
if (
|
||
mode == "multimodal"
|
||
and normalized["model"] != "seedance-2.5"
|
||
and audio_total
|
||
and not image_total
|
||
and not video_total
|
||
):
|
||
raise ValueError("Seedance 2.0 系列不可单独使用参考音频")
|
||
|
||
normalized.update({
|
||
"batch_id": _canonical_batch_id(payload.get("batch_id")),
|
||
"generator_node_id": _safe_node_id(payload.get("generator_node_id"), "生成节点 ID"),
|
||
"result_node_id": _safe_node_id(payload.get("result_node_id"), "结果节点 ID"),
|
||
"media": media,
|
||
"assets": assets,
|
||
"submitted_at": time.time(),
|
||
})
|
||
return normalized
|
||
|
||
|
||
def _resolve_input_path(input_directory: str, descriptor: dict[str, str], label: str) -> str:
|
||
root = os.path.realpath(os.path.abspath(input_directory))
|
||
path = os.path.realpath(
|
||
os.path.abspath(os.path.join(root, descriptor["subfolder"], descriptor["name"]))
|
||
)
|
||
try:
|
||
inside = os.path.commonpath((root, path)) == root
|
||
except ValueError:
|
||
inside = False
|
||
if not inside or not os.path.isfile(path):
|
||
raise ValueError(f"{label}文件不存在或超出 ComfyUI input 目录")
|
||
return path
|
||
|
||
|
||
def _validate_media_file(path: str, role: str) -> None:
|
||
extensions, maximum_size = _MEDIA_RULES[role]
|
||
extension = os.path.splitext(path)[1].lower()
|
||
if extension not in extensions:
|
||
raise ValueError(f"{role} 不支持文件格式 {extension or '无扩展名'}")
|
||
size = os.path.getsize(path)
|
||
if size <= 0 or size > maximum_size:
|
||
raise ValueError(
|
||
f"{role} 文件大小必须在 1 字节到 {maximum_size // 1024 // 1024}MB 之间"
|
||
)
|
||
if role in _IMAGE_MEDIA_ROLES:
|
||
try:
|
||
with Image.open(path) as image:
|
||
width, height = image.size
|
||
image.verify()
|
||
except Exception:
|
||
raise ValueError(f"{role} 不是可读取的图片") from None
|
||
validate_seedance_reference_dimensions(width, height, _IMAGE_ROLE_LABELS[role])
|
||
elif role in _VIDEO_MEDIA_ROLES:
|
||
width, height = _probe_video_dimensions(path)
|
||
validate_seedance_reference_dimensions(
|
||
width,
|
||
height,
|
||
_VIDEO_ROLE_LABELS[role],
|
||
require_video_pixel_range=True,
|
||
)
|
||
|
||
|
||
_validate_reference_dimensions = validate_seedance_reference_dimensions
|
||
|
||
|
||
def _probe_video_dimensions(path: str) -> tuple[int, int]:
|
||
"""Read video dimensions without decoding or rewriting the uploaded media."""
|
||
|
||
try:
|
||
import av
|
||
except ImportError:
|
||
raise ValueError("当前 ComfyUI 环境缺少 PyAV，无法校验参考视频分辨率") from None
|
||
|
||
try:
|
||
container = av.open(path)
|
||
try:
|
||
streams = list(container.streams.video)
|
||
if not streams:
|
||
raise ValueError("参考视频不包含视频轨道")
|
||
stream = streams[0]
|
||
width = int(stream.codec_context.width or stream.width or 0)
|
||
height = int(stream.codec_context.height or stream.height or 0)
|
||
finally:
|
||
container.close()
|
||
except ValueError:
|
||
raise
|
||
except Exception:
|
||
raise ValueError("参考视频无法读取分辨率，请确认文件可正常播放") from None
|
||
|
||
if width <= 0 or height <= 0:
|
||
raise ValueError("参考视频未包含有效的宽高信息")
|
||
return width, height
|
||
|
||
|
||
def snapshot_video_job_media(job: dict[str, Any], input_directory: str, temp_directory: str) -> None:
|
||
batch_root = os.path.join(temp_directory, "o1key_video_jobs", job["batch_id"])
|
||
input_root = os.path.join(batch_root, "inputs")
|
||
os.makedirs(input_root, exist_ok=False)
|
||
snapshots: dict[str, Any] = {}
|
||
for role, raw_value in job["media"].items():
|
||
values = raw_value if isinstance(raw_value, list) else ([raw_value] if raw_value else [])
|
||
copied = []
|
||
for index, descriptor in enumerate(values, start=1):
|
||
source = _resolve_input_path(input_directory, descriptor, role)
|
||
_validate_media_file(source, role)
|
||
extension = os.path.splitext(source)[1].lower()
|
||
target = os.path.join(input_root, f"{role}_{index:02d}{extension}")
|
||
shutil.copy2(source, target)
|
||
copied.append(target)
|
||
snapshots[role] = copied if role in _LIST_MEDIA_ROLES else (copied[0] if copied else None)
|
||
job["snapshot_root"] = batch_root
|
||
job["snapshot_media"] = snapshots
|
||
|
||
|
||
def cleanup_video_job_snapshot(job: dict[str, Any], temp_directory: str) -> None:
|
||
expected_root = os.path.realpath(os.path.abspath(os.path.join(temp_directory, "o1key_video_jobs")))
|
||
candidate = os.path.realpath(os.path.abspath(str(job.get("snapshot_root") or "")))
|
||
try:
|
||
inside = candidate != expected_root and os.path.commonpath((expected_root, candidate)) == expected_root
|
||
except ValueError:
|
||
inside = False
|
||
if inside and os.path.isdir(candidate):
|
||
shutil.rmtree(candidate, ignore_errors=True)
|
||
|
||
|
||
def _asset_url(value: str) -> str:
|
||
return value if value.startswith("asset://") else f"asset://{value}"
|
||
|
||
|
||
async def _prepare_seedance_media(job: dict[str, Any], update) -> dict[str, Any]:
|
||
base_url = get_base_url_by_route()
|
||
snapshots = job["snapshot_media"]
|
||
prepared = {
|
||
"first_frame": None,
|
||
"last_frame": None,
|
||
"reference_images": [],
|
||
"reference_videos": [],
|
||
"reference_audios": [],
|
||
}
|
||
|
||
manual_assets = job.get("asset_creation_mode") == "manual"
|
||
if manual_assets:
|
||
job["resolved_assets"] = {
|
||
kind: list(job.get("assets", {}).get(kind) or [])
|
||
for kind in ("images", "videos", "audios")
|
||
}
|
||
return prepared
|
||
|
||
items = []
|
||
if snapshots.get("first_frame"):
|
||
items.append(("first_frame", "image", snapshots["first_frame"]))
|
||
if snapshots.get("last_frame"):
|
||
items.append(("last_frame", "image", snapshots["last_frame"]))
|
||
items.extend(
|
||
("reference_images", "image", path)
|
||
for path in snapshots.get("reference_images") or []
|
||
)
|
||
items.extend(
|
||
("reference_videos", "video", path)
|
||
for path in snapshots.get("reference_videos") or []
|
||
)
|
||
items.extend(
|
||
("reference_audios", "audio", path)
|
||
for path in snapshots.get("reference_audios") or []
|
||
)
|
||
if not items:
|
||
job["resolved_assets"] = {"images": [], "videos": [], "audios": []}
|
||
return prepared
|
||
|
||
request_type = seedance_asset_request_type(job["route"])
|
||
request_label = "HC" if request_type == "hc" else "Doubao"
|
||
service = SeedanceAssetService(
|
||
client=SeedanceElementClient(base_url=base_url),
|
||
cache_path=job.get("asset_cache_path"),
|
||
)
|
||
semaphore = asyncio.Semaphore(3)
|
||
progress_lock = asyncio.Lock()
|
||
uploaded_count = 0
|
||
active_count = 0
|
||
total = len(items)
|
||
|
||
async def prepare(item_number: int, path: str, kind: str) -> tuple[str, str]:
|
||
nonlocal uploaded_count, active_count
|
||
async with semaphore:
|
||
fingerprint = await asyncio.to_thread(
|
||
seedance_asset_fingerprint,
|
||
path,
|
||
request_type,
|
||
kind,
|
||
)
|
||
|
||
async def upload() -> str:
|
||
nonlocal uploaded_count
|
||
update(stage=f"上传素材 {item_number}/{total}")
|
||
if kind == "image":
|
||
with Image.open(path) as opened:
|
||
opened.load()
|
||
value = await upload_image(opened.convert("RGB"), base_url=base_url)
|
||
elif kind == "video":
|
||
value = await upload_video(path, base_url=base_url)
|
||
else:
|
||
value = await upload_audio(path, base_url=base_url)
|
||
async with progress_lock:
|
||
uploaded_count += 1
|
||
update(
|
||
stage=f"已上传素材 {uploaded_count}/{total}",
|
||
progress=0.02 + (uploaded_count / total) * 0.06,
|
||
)
|
||
update(stage=f"创建 {request_label} 素材 {item_number}/{total}")
|
||
return value
|
||
|
||
update(stage=f"检查 {request_label} 素材 {item_number}/{total}")
|
||
created = await service.create_from_url(
|
||
name=f"o1key-{kind}-{item_number}",
|
||
asset_url_factory=upload,
|
||
asset_type=kind,
|
||
request_type=request_type,
|
||
fingerprint=fingerprint,
|
||
)
|
||
asset_id = str(created.get("Id") or "").strip()
|
||
if not asset_id:
|
||
raise RuntimeError(f"Seedance {kind} 素材创建成功但未返回 ID")
|
||
async with progress_lock:
|
||
active_count += 1
|
||
action = "复用" if created.get("_reused") else "激活"
|
||
update(
|
||
stage=f"素材已{action} {active_count}/{total}",
|
||
progress=0.08 + (active_count / total) * 0.12,
|
||
)
|
||
return _asset_url(asset_id), asset_id
|
||
|
||
results = await asyncio.gather(*(
|
||
prepare(item_number, path, kind)
|
||
for item_number, (_role, kind, path) in enumerate(items, start=1)
|
||
))
|
||
resolved = {"images": [], "videos": [], "audios": []}
|
||
for (role, kind, _path), (asset_url, asset_id) in zip(items, results):
|
||
if role in {"first_frame", "last_frame"}:
|
||
prepared[role] = asset_url
|
||
else:
|
||
prepared[role].append(asset_url)
|
||
resolved[f"{kind}s"].append(asset_id)
|
||
job["resolved_assets"] = resolved
|
||
return prepared
|
||
|
||
|
||
def build_seedance_video_body(job: dict[str, Any], prepared: dict[str, Any]) -> dict[str, Any]:
|
||
content = []
|
||
if job["prompt"]:
|
||
content.append({"type": "text", "text": job["prompt"]})
|
||
mode = job["generation_mode"]
|
||
image_asset_ids = job["assets"].get("images", job["assets"].get("persons", []))
|
||
manual_assets = job.get("asset_creation_mode") == "manual"
|
||
if mode == "first_frame":
|
||
image_url = _asset_url(image_asset_ids[0]) if manual_assets else prepared["first_frame"]
|
||
content.append({
|
||
"type": "image_url",
|
||
"image_url": {"url": image_url},
|
||
})
|
||
elif mode == "first_last_frame":
|
||
image_urls = (
|
||
[_asset_url(value) for value in image_asset_ids]
|
||
if manual_assets
|
||
else [prepared["first_frame"], prepared["last_frame"]]
|
||
)
|
||
content.extend([
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": image_urls[0]},
|
||
"role": "first_frame",
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": image_urls[1]},
|
||
"role": "last_frame",
|
||
},
|
||
])
|
||
elif mode == "multimodal":
|
||
for index, asset_id in enumerate(image_asset_ids):
|
||
content.append({
|
||
"type": "image_url",
|
||
"image_url": {"url": _asset_url(asset_id)},
|
||
"role": "subject" if index == 0 else "reference_image",
|
||
})
|
||
content.extend(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": url},
|
||
"role": "reference_image",
|
||
}
|
||
for url in prepared["reference_images"]
|
||
)
|
||
content.extend(
|
||
{
|
||
"type": "video_url",
|
||
"video_url": {"url": _asset_url(asset_id)},
|
||
"role": "reference_video",
|
||
}
|
||
for asset_id in job["assets"]["videos"]
|
||
)
|
||
content.extend(
|
||
{
|
||
"type": "video_url",
|
||
"video_url": {"url": url},
|
||
"role": "reference_video",
|
||
}
|
||
for url in prepared["reference_videos"]
|
||
)
|
||
content.extend(
|
||
{
|
||
"type": "audio_url",
|
||
"audio_url": {"url": _asset_url(asset_id)},
|
||
"role": "reference_audio",
|
||
}
|
||
for asset_id in job["assets"]["audios"]
|
||
)
|
||
content.extend(
|
||
{
|
||
"type": "audio_url",
|
||
"audio_url": {"url": url},
|
||
"role": "reference_audio",
|
||
}
|
||
for url in prepared["reference_audios"]
|
||
)
|
||
|
||
body = {
|
||
"model": job["actual_model"],
|
||
"content": content,
|
||
"duration": job["duration"],
|
||
"resolution": job["resolution"],
|
||
"ratio": "16:9" if job["aspect_ratio"] == "auto" else job["aspect_ratio"],
|
||
"generate_audio": job["generate_audio"],
|
||
"watermark": False,
|
||
"return_last_frame": job["return_last_frame"],
|
||
}
|
||
if job["seed"]:
|
||
body["seed"] = job["seed"]
|
||
return body
|
||
|
||
|
||
def _safe_prefix(value: str) -> str:
|
||
prefix = _UNSAFE_FILENAME.sub("_", str(value or "o1key_video")).strip(" ._")
|
||
return prefix[:180] or "o1key_video"
|
||
|
||
|
||
def _prepare_save_directory(job: dict[str, Any], output_directory: str) -> None:
|
||
"""Fail on an unusable destination before uploads or provider submission."""
|
||
|
||
location = normalize_save_location(job.get("save_location"))
|
||
directory = location if location and os.path.isabs(location) else os.path.join(output_directory, location)
|
||
os.makedirs(directory, exist_ok=True)
|
||
if not os.path.isdir(directory) or not os.access(directory, os.W_OK):
|
||
raise ValueError("视频保存位置不可写")
|
||
|
||
|
||
def _allocate_output_path(job: dict[str, Any], output_directory: str) -> tuple[str, str, str, bool]:
|
||
location = normalize_save_location(job.get("save_location"))
|
||
external = bool(location and os.path.isabs(location))
|
||
relative = "" if external else location
|
||
directory = location if external else os.path.join(output_directory, relative)
|
||
os.makedirs(directory, exist_ok=True)
|
||
prefix = _safe_prefix(job["filename_prefix"])
|
||
with _SAVE_LOCK:
|
||
for counter in range(1, 1_000_000):
|
||
filename = f"{prefix}_{counter:05d}_.mp4"
|
||
path = os.path.join(directory, filename)
|
||
try:
|
||
with open(path, "xb"):
|
||
pass
|
||
except FileExistsError:
|
||
continue
|
||
return path, filename, relative.replace("\\", "/"), external
|
||
raise RuntimeError("无法分配视频输出文件名")
|
||
|
||
|
||
def _atomic_promote(source: str, destination: str) -> None:
|
||
temporary = f"{destination}.{uuid.uuid4().hex}.tmp"
|
||
try:
|
||
shutil.copy2(source, temporary)
|
||
os.replace(temporary, destination)
|
||
except Exception:
|
||
# _allocate_output_path reserves a zero-byte destination so concurrent
|
||
# jobs cannot select the same counter. Do not leave that reservation
|
||
# behind when promotion fails.
|
||
try:
|
||
if os.path.isfile(destination) and os.path.getsize(destination) == 0:
|
||
os.remove(destination)
|
||
except OSError:
|
||
pass
|
||
raise
|
||
finally:
|
||
if os.path.exists(temporary):
|
||
try:
|
||
os.remove(temporary)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _preview_descriptor_for_external(
|
||
source: str,
|
||
filename: str,
|
||
batch_id: str,
|
||
temp_directory: str,
|
||
) -> dict[str, str]:
|
||
subfolder = os.path.join("o1key_video_preview", batch_id)
|
||
directory = os.path.join(temp_directory, subfolder)
|
||
os.makedirs(directory, exist_ok=True)
|
||
target = os.path.join(directory, filename)
|
||
_atomic_promote(source, target)
|
||
return {"filename": filename, "subfolder": subfolder.replace("\\", "/"), "type": "temp"}
|
||
|
||
|
||
async def _download_last_frame(
|
||
url: str,
|
||
job: dict[str, Any],
|
||
output_directory: str,
|
||
temp_directory: str,
|
||
video_path: str,
|
||
video_filename: str,
|
||
relative: str,
|
||
external: bool,
|
||
) -> dict[str, str] | None:
|
||
if not url or not str(url).startswith(("http://", "https://")):
|
||
return None
|
||
timeout = aiohttp.ClientTimeout(total=120)
|
||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||
async with session.get(url, allow_redirects=True) as response:
|
||
if response.status != 200:
|
||
return None
|
||
data = await response.read()
|
||
if not data or len(data) > 32 * 1024 * 1024:
|
||
return None
|
||
from io import BytesIO
|
||
|
||
with Image.open(BytesIO(data)) as opened:
|
||
opened.load()
|
||
frame = opened.convert("RGB")
|
||
frame_filename = os.path.splitext(video_filename)[0] + "_last_frame.png"
|
||
destination = os.path.join(os.path.dirname(video_path), frame_filename)
|
||
temporary = f"{destination}.{uuid.uuid4().hex}.tmp"
|
||
try:
|
||
frame.save(temporary, format="PNG")
|
||
os.replace(temporary, destination)
|
||
finally:
|
||
if os.path.exists(temporary):
|
||
try:
|
||
os.remove(temporary)
|
||
except OSError:
|
||
pass
|
||
if external:
|
||
return _preview_descriptor_for_external(
|
||
destination, frame_filename, job["batch_id"], temp_directory
|
||
)
|
||
return {"filename": frame_filename, "subfolder": relative, "type": "output"}
|
||
|
||
|
||
async def execute_video_job(
|
||
job: dict[str, Any],
|
||
output_directory: str,
|
||
temp_directory: str,
|
||
update: Callable[..., None],
|
||
) -> dict[str, Any]:
|
||
update(stage="preparing", progress=0.02)
|
||
await asyncio.to_thread(_prepare_save_directory, job, output_directory)
|
||
prepared = await _prepare_seedance_media(job, update)
|
||
body = build_seedance_video_body(job, prepared)
|
||
|
||
generated_path = os.path.join(job["snapshot_root"], "provider_result.mp4")
|
||
client = SeedanceClient()
|
||
client.base_url = get_base_url_by_route()
|
||
|
||
def on_stage(stage: str) -> None:
|
||
if stage.startswith("submitted:"):
|
||
update(stage="polling", progress=0.22, provider_task_id=stage.split(":", 1)[1])
|
||
elif stage == "submitting":
|
||
update(stage="submitting", progress=0.21)
|
||
elif stage == "downloading":
|
||
update(stage="downloading", progress=0.97)
|
||
|
||
def on_progress(value: int) -> None:
|
||
update(stage="polling", progress=0.22 + max(0, min(100, int(value))) * 0.0074)
|
||
|
||
try:
|
||
generated_path, last_frame_url = await client.generate_async(
|
||
body=body,
|
||
save_path=generated_path,
|
||
on_stage=on_stage,
|
||
on_progress=on_progress,
|
||
use_new_format=True,
|
||
)
|
||
except Exception as exc:
|
||
message = format_o1key_video_error(exc)
|
||
if message == str(exc):
|
||
raise
|
||
raise RuntimeError(message) from None
|
||
|
||
if not os.path.isfile(generated_path) or os.path.getsize(generated_path) <= 0:
|
||
raise RuntimeError("Seedance 返回的视频文件为空")
|
||
destination, filename, relative, external = _allocate_output_path(job, output_directory)
|
||
update(stage="saving", progress=0.99)
|
||
_atomic_promote(generated_path, destination)
|
||
if external:
|
||
video_descriptor = _preview_descriptor_for_external(
|
||
destination, filename, job["batch_id"], temp_directory
|
||
)
|
||
else:
|
||
video_descriptor = {"filename": filename, "subfolder": relative, "type": "output"}
|
||
last_frame_descriptor = await _download_last_frame(
|
||
last_frame_url or "",
|
||
job,
|
||
output_directory,
|
||
temp_directory,
|
||
destination,
|
||
filename,
|
||
relative,
|
||
external,
|
||
) if job["return_last_frame"] else None
|
||
return {
|
||
"video": video_descriptor,
|
||
"last_frame": last_frame_descriptor,
|
||
"resolved_assets": job.get("resolved_assets"),
|
||
}
|
||
|
||
|
||
class VideoJobHistory:
|
||
def __init__(self, path: str):
|
||
self.path = path
|
||
self._lock = threading.Lock()
|
||
|
||
def _read(self) -> list[dict[str, Any]]:
|
||
if not os.path.isfile(self.path):
|
||
return []
|
||
try:
|
||
with open(self.path, "r", encoding="utf-8") as handle:
|
||
value = json.load(handle)
|
||
return value if isinstance(value, list) else []
|
||
except (OSError, json.JSONDecodeError):
|
||
return []
|
||
|
||
def upsert(self, item: dict[str, Any]) -> None:
|
||
with self._lock:
|
||
items = [value for value in self._read() if value.get("batch_id") != item.get("batch_id")]
|
||
items.insert(0, item)
|
||
items = items[:_MAX_HISTORY_ITEMS]
|
||
os.makedirs(os.path.dirname(self.path), exist_ok=True)
|
||
temporary = f"{self.path}.{uuid.uuid4().hex}.tmp"
|
||
with open(temporary, "w", encoding="utf-8") as handle:
|
||
json.dump(items, handle, ensure_ascii=False, indent=2)
|
||
os.replace(temporary, self.path)
|
||
|
||
def list(self, limit: int = 64) -> list[dict[str, Any]]:
|
||
with self._lock:
|
||
return self._read()[: max(0, min(_MAX_HISTORY_ITEMS, int(limit)))]
|
||
|
||
def get(self, batch_id: str) -> dict[str, Any] | None:
|
||
with self._lock:
|
||
return next((item for item in self._read() if item.get("batch_id") == batch_id), None)
|
||
|
||
def delete(self, batch_id: str) -> None:
|
||
with self._lock:
|
||
items = [value for value in self._read() if value.get("batch_id") != batch_id]
|
||
os.makedirs(os.path.dirname(self.path), exist_ok=True)
|
||
temporary = f"{self.path}.{uuid.uuid4().hex}.tmp"
|
||
with open(temporary, "w", encoding="utf-8") as handle:
|
||
json.dump(items, handle, ensure_ascii=False, indent=2)
|
||
os.replace(temporary, self.path)
|
||
|
||
|
||
@dataclass
|
||
class VideoJobRecord:
|
||
job: dict[str, Any]
|
||
state: str = "queued"
|
||
stage: str = "queued"
|
||
progress: float = 0.0
|
||
provider_task_id: str = ""
|
||
video: dict[str, Any] | None = None
|
||
last_frame: dict[str, Any] | None = None
|
||
resolved_assets: dict[str, list[str]] | None = None
|
||
error: str = ""
|
||
started_at: float = 0.0
|
||
ended_at: float = 0.0
|
||
task: asyncio.Task | None = field(default=None, repr=False)
|
||
|
||
|
||
class ParallelVideoJobManager:
|
||
def __init__(self, executor, event_sender, *, history_store=None):
|
||
self.executor = executor
|
||
self.event_sender = event_sender
|
||
self.history_store = history_store
|
||
self.jobs: dict[str, VideoJobRecord] = {}
|
||
self.completed_order: deque[str] = deque()
|
||
self._loop: asyncio.AbstractEventLoop | None = None
|
||
self._lock: asyncio.Lock | None = None
|
||
|
||
def _ensure_loop(self) -> None:
|
||
loop = asyncio.get_running_loop()
|
||
if self._loop is loop:
|
||
return
|
||
self._loop = loop
|
||
self._lock = asyncio.Lock()
|
||
|
||
def _counts(self) -> tuple[int, int]:
|
||
return (
|
||
sum(record.state == "running" for record in self.jobs.values()),
|
||
sum(record.state == "queued" for record in self.jobs.values()),
|
||
)
|
||
|
||
def status(self, batch_id: str) -> dict[str, Any] | None:
|
||
record = self.jobs.get(batch_id)
|
||
if record is None:
|
||
return None
|
||
active, queued = self._counts()
|
||
return {
|
||
"batch_id": batch_id,
|
||
"generator_node_id": record.job["generator_node_id"],
|
||
"result_node_id": record.job["result_node_id"],
|
||
"state": record.state,
|
||
"stage": record.stage,
|
||
"progress": record.progress,
|
||
"progress_percent": round(record.progress * 100),
|
||
"provider": record.job["provider"],
|
||
"model": record.job["model"],
|
||
"provider_task_id": record.provider_task_id,
|
||
"video": dict(record.video) if record.video else None,
|
||
"last_frame": dict(record.last_frame) if record.last_frame else None,
|
||
"resolved_assets": {
|
||
kind: list(values)
|
||
for kind, values in (
|
||
record.resolved_assets or record.job.get("resolved_assets") or {}
|
||
).items()
|
||
if kind in {"images", "videos", "audios"} and isinstance(values, list)
|
||
} or None,
|
||
"error": record.error,
|
||
"active_jobs": active,
|
||
"queued_jobs": queued,
|
||
"create_time": round(record.job["submitted_at"] * 1000),
|
||
"execution_start_time": round(record.started_at * 1000) if record.started_at else 0,
|
||
"execution_end_time": round(record.ended_at * 1000) if record.ended_at else 0,
|
||
}
|
||
|
||
async def _emit(self, record: VideoJobRecord) -> None:
|
||
try:
|
||
await self.event_sender("o1key.video_job", self.status(record.job["batch_id"]))
|
||
except Exception:
|
||
pass
|
||
|
||
async def submit(self, job: dict[str, Any]) -> dict[str, Any]:
|
||
self._ensure_loop()
|
||
async with self._lock:
|
||
if job["batch_id"] in self.jobs:
|
||
raise ValueError("视频任务 ID 已存在")
|
||
record = VideoJobRecord(
|
||
job=dict(job),
|
||
state="running",
|
||
stage="starting",
|
||
started_at=time.time(),
|
||
)
|
||
self.jobs[job["batch_id"]] = record
|
||
record.task = asyncio.create_task(self._run(record))
|
||
return self.status(job["batch_id"])
|
||
|
||
async def _run(self, record: VideoJobRecord) -> None:
|
||
try:
|
||
await self._emit(record)
|
||
record.stage = "preparing"
|
||
await self._emit(record)
|
||
|
||
def update(*, stage=None, progress=None, provider_task_id=None):
|
||
if stage:
|
||
record.stage = str(stage)
|
||
if progress is not None:
|
||
record.progress = max(record.progress, min(1.0, float(progress)))
|
||
if provider_task_id:
|
||
record.provider_task_id = str(provider_task_id)
|
||
asyncio.create_task(self._emit(record))
|
||
|
||
result = await self.executor(record.job, update)
|
||
record.video = result.get("video")
|
||
record.last_frame = result.get("last_frame")
|
||
record.resolved_assets = result.get("resolved_assets")
|
||
record.progress = 1.0
|
||
record.stage = "completed"
|
||
record.state = "completed"
|
||
record.ended_at = time.time()
|
||
await self._emit(record)
|
||
except asyncio.CancelledError:
|
||
record.state = "cancelled"
|
||
record.stage = "cancelled"
|
||
record.error = (
|
||
"已停止本地等待；远端任务可能仍在继续"
|
||
if record.provider_task_id
|
||
else "任务已取消"
|
||
)
|
||
record.ended_at = time.time()
|
||
await self._emit(record)
|
||
except Exception as exc:
|
||
record.state = "failed"
|
||
record.stage = "failed"
|
||
record.error = " ".join(format_o1key_video_error(exc).split())[:600]
|
||
record.ended_at = time.time()
|
||
await self._emit(record)
|
||
finally:
|
||
if self.history_store is not None:
|
||
try:
|
||
await asyncio.to_thread(self.history_store.upsert, self.status(record.job["batch_id"]))
|
||
except OSError:
|
||
pass
|
||
await asyncio.to_thread(cleanup_video_job_snapshot, record.job, record.job["temp_directory"])
|
||
self.completed_order.append(record.job["batch_id"])
|
||
while len(self.completed_order) > _MAX_HISTORY_ITEMS:
|
||
self.jobs.pop(self.completed_order.popleft(), None)
|
||
|
||
async def cancel(self, batch_id: str) -> dict[str, Any] | None:
|
||
self._ensure_loop()
|
||
async with self._lock:
|
||
record = self.jobs.get(batch_id)
|
||
if record is None:
|
||
return None
|
||
if record.task and not record.task.done() and record.state in {"queued", "running"}:
|
||
record.task.cancel()
|
||
if record.task and not record.task.done():
|
||
try:
|
||
await record.task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
return self.status(batch_id)
|
||
|
||
|
||
def register_o1key_video_job_routes(PromptServer, web, folder_paths):
|
||
async def send_event(event: str, payload: dict[str, Any]) -> None:
|
||
await PromptServer.instance.send_json(event, payload)
|
||
|
||
async def executor(job: dict[str, Any], update):
|
||
return await execute_video_job(
|
||
job,
|
||
folder_paths.get_output_directory(),
|
||
folder_paths.get_temp_directory(),
|
||
update,
|
||
)
|
||
|
||
get_user_directory = getattr(folder_paths, "get_user_directory", None)
|
||
history_root = get_user_directory() if callable(get_user_directory) else folder_paths.get_output_directory()
|
||
history = VideoJobHistory(os.path.join(history_root, "o1key", "video_jobs.json"))
|
||
asset_cache_path = os.path.join(history_root, "o1key", "seedance_asset_cache.json")
|
||
manager = ParallelVideoJobManager(executor, send_event, history_store=history)
|
||
|
||
@PromptServer.instance.routes.get("/o1key/video/capabilities")
|
||
async def get_video_capabilities(_request):
|
||
return web.json_response(public_video_capabilities(), headers={"Cache-Control": "no-store"})
|
||
|
||
@PromptServer.instance.routes.post("/o1key/video/jobs")
|
||
async def submit_video_job(request):
|
||
job = None
|
||
try:
|
||
job = normalize_video_job_payload(await request.json())
|
||
job["temp_directory"] = folder_paths.get_temp_directory()
|
||
job["asset_cache_path"] = asset_cache_path
|
||
await asyncio.to_thread(
|
||
snapshot_video_job_media,
|
||
job,
|
||
folder_paths.get_input_directory(),
|
||
folder_paths.get_temp_directory(),
|
||
)
|
||
return web.json_response(
|
||
await manager.submit(job),
|
||
status=202,
|
||
headers={"Cache-Control": "no-store"},
|
||
)
|
||
except ValueError as exc:
|
||
if job is not None:
|
||
await asyncio.to_thread(cleanup_video_job_snapshot, job, folder_paths.get_temp_directory())
|
||
return web.json_response({"error": str(exc)}, status=400)
|
||
except Exception as exc:
|
||
if job is not None:
|
||
await asyncio.to_thread(cleanup_video_job_snapshot, job, folder_paths.get_temp_directory())
|
||
return web.json_response({"error": str(exc)}, status=500)
|
||
|
||
@PromptServer.instance.routes.get("/o1key/video/jobs/history")
|
||
async def get_video_history(request):
|
||
try:
|
||
limit = int(request.query.get("limit", 64))
|
||
except (TypeError, ValueError):
|
||
limit = 64
|
||
return web.json_response(
|
||
{"items": await asyncio.to_thread(history.list, limit)},
|
||
headers={"Cache-Control": "no-store"},
|
||
)
|
||
|
||
@PromptServer.instance.routes.post("/o1key/video/jobs/history")
|
||
async def delete_video_history(request):
|
||
try:
|
||
payload = await request.json()
|
||
batch_id = _canonical_batch_id(payload.get("batch_id"))
|
||
await asyncio.to_thread(history.delete, batch_id)
|
||
return web.json_response({"ok": True}, headers={"Cache-Control": "no-store"})
|
||
except ValueError as exc:
|
||
return web.json_response({"error": str(exc)}, status=400)
|
||
|
||
@PromptServer.instance.routes.get("/o1key/video/jobs/{batch_id}")
|
||
async def get_video_job(request):
|
||
try:
|
||
batch_id = _canonical_batch_id(request.match_info.get("batch_id"))
|
||
except ValueError as exc:
|
||
return web.json_response({"error": str(exc)}, status=400)
|
||
result = manager.status(batch_id) or await asyncio.to_thread(history.get, batch_id)
|
||
if result is None:
|
||
return web.json_response({"error": "视频任务不存在"}, status=404)
|
||
return web.json_response(result, headers={"Cache-Control": "no-store"})
|
||
|
||
@PromptServer.instance.routes.post("/o1key/video/jobs/{batch_id}/cancel")
|
||
async def cancel_video_job(request):
|
||
try:
|
||
batch_id = _canonical_batch_id(request.match_info.get("batch_id"))
|
||
except ValueError as exc:
|
||
return web.json_response({"error": str(exc)}, status=400)
|
||
result = await manager.cancel(batch_id)
|
||
if result is None:
|
||
return web.json_response({"error": "视频任务不存在"}, status=404)
|
||
return web.json_response(result, headers={"Cache-Control": "no-store"})
|
||
|
||
return manager
|
||
|
||
|
||
__all__ = [
|
||
"ParallelVideoJobManager",
|
||
"VideoJobHistory",
|
||
"build_seedance_video_body",
|
||
"execute_video_job",
|
||
"normalize_video_job_payload",
|
||
"register_o1key_video_job_routes",
|
||
"snapshot_video_job_media",
|
||
]
|