669 lines
22 KiB
Python
669 lines
22 KiB
Python
import io
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import tempfile
|
||
import time
|
||
import uuid
|
||
from time import sleep
|
||
from typing import Any, Dict
|
||
|
||
import cv2
|
||
import folder_paths
|
||
import numpy as np
|
||
import requests
|
||
import torch
|
||
from loguru import logger
|
||
from PIL import Image
|
||
from torchvision.transforms import transforms
|
||
from tqdm import tqdm
|
||
|
||
from ..utils.config_utils import config
|
||
from ..utils.object_storage import UploadResult, get_provider
|
||
|
||
|
||
class JMUtils:
|
||
"""
|
||
即梦AI工具类
|
||
|
||
提供即梦AI视频生成服务的完整功能，包括：
|
||
- 图像上传到云存储
|
||
- 任务提交和状态查询
|
||
- 视频下载和处理
|
||
- 张量和图像格式转换
|
||
|
||
使用统一的存储抽象层，支持多种云存储服务。
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""
|
||
初始化即梦工具实例
|
||
|
||
从配置中读取API密钥和存储配置信息
|
||
"""
|
||
try:
|
||
# 获取即梦API配置
|
||
self.api_key = config.get_config("jm_api_key")
|
||
if not self.api_key:
|
||
raise ValueError("即梦API密钥未配置")
|
||
|
||
# 获取COS存储配置（用于素材上传）
|
||
cos_config = config.get_cos_config()
|
||
self.cos_bucket_name = cos_config.get("bucket_name") or config.get_config(
|
||
"cos_sucai_bucket_name"
|
||
)
|
||
|
||
if not self.cos_bucket_name:
|
||
raise ValueError("COS素材存储桶未配置")
|
||
|
||
# 获取存储提供者
|
||
self.storage_provider = get_provider("cos")
|
||
|
||
logger.info(f"即梦工具初始化成功，使用存储桶: {self.cos_bucket_name}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"即梦工具初始化失败: {e}")
|
||
raise
|
||
|
||
def submit_task(
|
||
self, prompt: str, img_url: str, duration: str = "10", resolution: str = "720p"
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
提交即梦AI视频生成任务
|
||
|
||
Args:
|
||
prompt: 生成提示词
|
||
img_url: 输入图像URL
|
||
duration: 视频时长（秒）
|
||
resolution: 视频分辨率
|
||
|
||
Returns:
|
||
Dict: 任务提交结果
|
||
- status: 是否成功
|
||
- data: 任务ID或原图URL
|
||
- msg: 消息
|
||
"""
|
||
try:
|
||
# 验证输入参数
|
||
if not prompt or not prompt.strip():
|
||
return {"status": False, "data": None, "msg": "提示词不能为空"}
|
||
if not img_url or not img_url.strip():
|
||
return {"status": False, "data": None, "msg": "图像URL不能为空"}
|
||
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
}
|
||
|
||
json_data = {
|
||
"model": "doubao-seedance-1-0-pro-250528",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"{prompt.strip()} --resolution {resolution} --dur {duration} --camerafixed false",
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": img_url.strip(),
|
||
},
|
||
},
|
||
],
|
||
}
|
||
|
||
logger.info(
|
||
f"即梦任务提交中: prompt='{prompt[:50]}...', resolution={resolution}, duration={duration}"
|
||
)
|
||
|
||
response = requests.post(
|
||
"https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks",
|
||
headers=headers,
|
||
json=json_data,
|
||
timeout=30,
|
||
)
|
||
response.raise_for_status()
|
||
|
||
resp_json = response.json()
|
||
logger.info(
|
||
f"即梦任务提交响应: {json.dumps(resp_json, ensure_ascii=False)}"
|
||
)
|
||
|
||
if "id" not in resp_json:
|
||
error_msg = "未知错误"
|
||
if "error" in resp_json and "message" in resp_json["error"]:
|
||
error_msg = resp_json["error"]["message"]
|
||
return {
|
||
"status": False,
|
||
"data": img_url,
|
||
"msg": f"任务提交失败: {error_msg}",
|
||
}
|
||
else:
|
||
job_id = resp_json["id"]
|
||
logger.info(f"即梦任务提交成功，任务ID: {job_id}")
|
||
return {"data": job_id, "status": True, "msg": "任务提交成功"}
|
||
|
||
except requests.RequestException as e:
|
||
logger.error(f"即梦API请求失败: {e}")
|
||
return {"data": None, "status": False, "msg": f"网络请求失败: {str(e)}"}
|
||
except Exception as e:
|
||
logger.error(f"即梦任务提交异常: {e}")
|
||
return {"data": None, "status": False, "msg": str(e)}
|
||
|
||
def query_status(self, job_id: str) -> Dict[str, Any]:
|
||
"""
|
||
查询即梦AI任务状态
|
||
|
||
Args:
|
||
job_id: 任务ID
|
||
|
||
Returns:
|
||
Dict: 任务状态查询结果
|
||
- status: 任务是否完成成功
|
||
- data: 视频URL（如果完成）
|
||
- msg: 状态消息
|
||
"""
|
||
resp_dict = {"status": False, "data": None, "msg": ""}
|
||
|
||
try:
|
||
if not job_id or not job_id.strip():
|
||
resp_dict["msg"] = "任务ID不能为空"
|
||
return resp_dict
|
||
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
}
|
||
|
||
response = requests.get(
|
||
f"https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks/{job_id.strip()}",
|
||
headers=headers,
|
||
timeout=15,
|
||
)
|
||
response.raise_for_status()
|
||
|
||
resp_json = response.json()
|
||
task_status = resp_json.get("status", "unknown")
|
||
|
||
# 任务完成成功
|
||
if task_status == "succeeded":
|
||
resp_dict["status"] = True
|
||
resp_dict["msg"] = "任务完成"
|
||
if "content" in resp_json and "video_url" in resp_json["content"]:
|
||
resp_dict["data"] = resp_json["content"]["video_url"]
|
||
else:
|
||
resp_dict["status"] = False
|
||
resp_dict["msg"] = "任务完成但未找到视频URL"
|
||
|
||
# 任务失败
|
||
elif task_status in ["failed", "error"]:
|
||
resp_dict["status"] = False
|
||
error_msg = "任务失败"
|
||
if "error" in resp_json:
|
||
error_msg += f": {resp_json['error'].get('message', '未知错误')}"
|
||
resp_dict["msg"] = error_msg
|
||
|
||
# 任务进行中
|
||
elif task_status in ["pending", "running", "processing"]:
|
||
resp_dict["status"] = False
|
||
resp_dict["msg"] = f"任务进行中: {task_status}"
|
||
|
||
# 其他状态
|
||
else:
|
||
resp_dict["status"] = False
|
||
resp_dict["msg"] = f"未知任务状态: {task_status}"
|
||
|
||
except requests.RequestException as e:
|
||
logger.error(f"即梦状态查询网络错误: {e}")
|
||
resp_dict["msg"] = f"网络请求失败: {str(e)}"
|
||
except Exception as e:
|
||
logger.error(f"即梦状态查询异常: {e}")
|
||
resp_dict["msg"] = str(e)
|
||
|
||
return resp_dict
|
||
|
||
def upload_io_to_cos(
|
||
self, file: io.IOBase, mime_type: str = "image/png"
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
上传IO对象到COS存储
|
||
|
||
Args:
|
||
file: 文件IO对象
|
||
mime_type: MIME类型
|
||
|
||
Returns:
|
||
Dict: 包含上传结果的字典
|
||
- status: 是否成功
|
||
- data: 上传后的URL
|
||
- msg: 消息
|
||
"""
|
||
resp_data = {"status": True, "data": "", "msg": ""}
|
||
|
||
try:
|
||
# 解析MIME类型
|
||
parts = mime_type.split("/")
|
||
category = parts[0] if len(parts) > 0 else "file"
|
||
suffix = parts[1] if len(parts) > 1 else "bin"
|
||
|
||
# 生成存储键名
|
||
object_key = f"tk/{category}/{uuid.uuid4()}.{suffix}"
|
||
|
||
logger.info(f"开始上传文件到COS: {object_key}")
|
||
|
||
# 读取文件内容
|
||
file_content = file.read()
|
||
file.seek(0) # 重置文件指针
|
||
|
||
# 使用统一存储接口上传
|
||
result: UploadResult = self.storage_provider.upload_bytes(
|
||
file_content, object_key, bucket_name=self.cos_bucket_name
|
||
)
|
||
|
||
if result.success:
|
||
# 构造COS URL（如果result中没有提供）
|
||
if result.url:
|
||
resp_data["data"] = result.url
|
||
else:
|
||
# 构造默认的COS URL
|
||
cos_config = config.get_cos_config()
|
||
region = cos_config.get("region", "ap-beijing")
|
||
resp_data["data"] = (
|
||
f"https://{self.cos_bucket_name}.cos.{region}.myqcloud.com/{object_key}"
|
||
)
|
||
|
||
resp_data["msg"] = "上传成功"
|
||
logger.info(f"文件上传成功: {resp_data['data']}")
|
||
else:
|
||
resp_data["status"] = False
|
||
resp_data["msg"] = result.message or "上传失败"
|
||
logger.error(f"文件上传失败: {resp_data['msg']}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"上传文件时发生异常: {e}")
|
||
resp_data["status"] = False
|
||
resp_data["msg"] = str(e)
|
||
|
||
return resp_data
|
||
|
||
def tensor_to_io(self, tensor: torch.Tensor) -> io.BytesIO:
|
||
"""
|
||
将PyTorch张量转换为PNG格式的IO对象
|
||
|
||
Args:
|
||
tensor: PyTorch图像张量，支持多种格式
|
||
- (H, W) 灰度图
|
||
- (H, W, C) RGB图像
|
||
- (1, H, W, C) 批次图像
|
||
|
||
Returns:
|
||
io.BytesIO: PNG格式的字节流对象
|
||
|
||
Raises:
|
||
ValueError: 当张量格式不支持时
|
||
"""
|
||
try:
|
||
# 处理张量维度
|
||
if tensor.dim() == 4: # (1, H, W, C)
|
||
tensor = tensor.squeeze(0)
|
||
elif tensor.dim() == 2: # (H, W) 灰度图
|
||
pass # 保持原样
|
||
elif tensor.dim() == 3: # (H, W, C)
|
||
pass # 保持原样
|
||
else:
|
||
raise ValueError(f"不支持的张量维度: {tensor.dim()}D")
|
||
|
||
# 转换为numpy数组
|
||
numpy_array = tensor.cpu().numpy()
|
||
|
||
# 确保数值在有效范围内
|
||
numpy_array = np.clip(numpy_array, 0.0, 1.0)
|
||
|
||
# 转换为0-255范围的uint8
|
||
image_array = (numpy_array * 255).astype(np.uint8)
|
||
|
||
# 处理灰度图
|
||
if len(image_array.shape) == 2:
|
||
img = Image.fromarray(image_array, mode="L")
|
||
else:
|
||
img = Image.fromarray(image_array, mode="RGB")
|
||
|
||
# 保存为PNG格式的BytesIO
|
||
image_data = io.BytesIO()
|
||
img.save(image_data, format="PNG", optimize=True)
|
||
image_data.seek(0)
|
||
|
||
logger.debug(f"张量转换为PNG成功，大小: {len(image_data.getvalue())} bytes")
|
||
return image_data
|
||
|
||
except Exception as e:
|
||
logger.error(f"张量转换失败: {e}")
|
||
raise ValueError(f"张量转换为图像失败: {str(e)}")
|
||
|
||
def read_video_last_frame_to_tensor(self, video_path: str) -> torch.Tensor:
|
||
"""
|
||
读取视频文件的最后一帧并将其转换为BCHW格式的PyTorch张量。
|
||
|
||
参数:
|
||
video_path (str): 视频文件的路径。
|
||
|
||
返回:
|
||
torch.Tensor: 形状为[1, H, W, C]的张量，其中H和W分别是视频帧的高度和宽度，通道顺序为RGB。
|
||
|
||
异常:
|
||
FileNotFoundError: 如果指定的视频文件不存在。
|
||
ValueError: 如果视频文件为空或无法读取帧。
|
||
"""
|
||
# 打开视频文件
|
||
cap = cv2.VideoCapture(video_path)
|
||
|
||
if not cap.isOpened():
|
||
raise FileNotFoundError(f"无法打开视频文件: {video_path}")
|
||
|
||
# 获取视频总帧数
|
||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
|
||
if total_frames == 0:
|
||
cap.release()
|
||
raise ValueError("视频文件为空或无法确定帧数")
|
||
|
||
# 设置读取位置到最后一帧
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, total_frames - 1)
|
||
|
||
# 读取最后一帧
|
||
ret, frame = cap.read()
|
||
|
||
# 释放资源
|
||
cap.release()
|
||
|
||
if not ret or frame is None:
|
||
raise ValueError(f"无法读取视频的最后一帧，可能视频已损坏")
|
||
|
||
# 转换BGR到RGB (OpenCV默认读取为BGR)
|
||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
|
||
# 转换为PyTorch张量并调整维度为BCHW
|
||
transform = transforms.Compose(
|
||
[transforms.ToTensor()] # 转换为[C, H, W]格式的张量，值范围从0到1
|
||
)
|
||
|
||
tensor = (
|
||
transform(frame_rgb).unsqueeze(0).permute(0, 2, 3, 1)
|
||
) # 添加批次维度，变为[1, H, W, C]
|
||
|
||
return tensor
|
||
|
||
def download_video(self, url, timeout=30, retries=3, path=None):
|
||
"""下载视频到临时文件并返回文件路径"""
|
||
for attempt in range(retries):
|
||
try:
|
||
# 创建临时文件
|
||
if path:
|
||
temp_path = path
|
||
else:
|
||
temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
|
||
temp_path = temp_file.name
|
||
temp_file.close()
|
||
|
||
# 下载视频
|
||
print(f"开始下载视频 (尝试 {attempt + 1}/{retries})...")
|
||
response = requests.get(url, stream=True, timeout=timeout)
|
||
response.raise_for_status()
|
||
|
||
# 获取文件大小
|
||
total_size = int(response.headers.get("content-length", 0))
|
||
block_size = 1024 # 1 KB
|
||
|
||
# 使用tqdm显示下载进度
|
||
with open(temp_path, "wb") as f, tqdm(
|
||
desc=url.split("/")[-1],
|
||
total=total_size,
|
||
unit="B",
|
||
unit_scale=True,
|
||
unit_divisor=1024,
|
||
) as bar:
|
||
for data in response.iter_content(block_size):
|
||
size = f.write(data)
|
||
bar.update(size)
|
||
|
||
print(f"视频下载完成: {temp_path}")
|
||
return temp_path, self.read_video_last_frame_to_tensor(temp_path)
|
||
|
||
except Exception as e:
|
||
print(f"下载错误 (尝试 {attempt + 1}/{retries}): {str(e)}")
|
||
if attempt < retries - 1:
|
||
time.sleep((attempt + 1) * 2)
|
||
else:
|
||
raise
|
||
|
||
def jpg_to_tensor(self, image_path):
|
||
"""
|
||
将JPG图像转换为PyTorch张量
|
||
|
||
参数:
|
||
- image_path: JPG图像文件路径
|
||
|
||
返回:
|
||
- tensor: PyTorch张量
|
||
"""
|
||
try:
|
||
# 打开图像文件
|
||
image = Image.open(image_path).convert("RGB")
|
||
|
||
# 转换为张量
|
||
tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
|
||
|
||
return tensor
|
||
|
||
except Exception as e:
|
||
print(f"转换失败: {str(e)}")
|
||
raise
|
||
|
||
def get_last_15th_frame_tensor(self, video_url):
|
||
"""
|
||
从视频URL截取倒数第15帧并转换为Tensor
|
||
先下载视频到本地临时文件再处理
|
||
"""
|
||
try:
|
||
# 下载视频
|
||
video_path, _ = self.download_video(video_url)
|
||
|
||
# 获取视频总帧数
|
||
cmd_frames = [
|
||
"ffprobe",
|
||
"-v",
|
||
"error",
|
||
"-select_streams",
|
||
"v:0",
|
||
"-show_entries",
|
||
"stream=nb_frames",
|
||
"-of",
|
||
"default=nokey=1:noprint_wrappers=1",
|
||
video_path,
|
||
]
|
||
|
||
result = subprocess.run(
|
||
cmd_frames, capture_output=True, text=True, check=True
|
||
)
|
||
|
||
# 处理可能的非数字输出
|
||
frame_count = result.stdout.strip()
|
||
if not frame_count.isdigit():
|
||
# 备选方案：通过解码获取帧数
|
||
print("无法获取准确帧数，尝试直接解码...")
|
||
cmd_decode = ["ffmpeg", "-i", video_path, "-f", "null", "-"]
|
||
decode_result = subprocess.run(
|
||
cmd_decode, capture_output=True, text=True
|
||
)
|
||
|
||
for line in decode_result.stderr.split("\n"):
|
||
if "frame=" in line:
|
||
parts = line.split("frame=")[-1].split()[0]
|
||
if parts.isdigit():
|
||
frame_count = int(parts)
|
||
break
|
||
else:
|
||
raise ValueError("无法确定视频帧数")
|
||
else:
|
||
frame_count = int(frame_count)
|
||
|
||
# 计算目标帧
|
||
target_frame = max(0, frame_count - 15)
|
||
print(f"视频总帧数: {frame_count}, 目标帧: {target_frame}")
|
||
|
||
# 截取指定帧
|
||
with tempfile.NamedTemporaryFile(
|
||
suffix="%03d.jpg", delete=True
|
||
) as frame_file:
|
||
frame_path = frame_file.name
|
||
|
||
cmd_extract = [
|
||
"ffmpeg",
|
||
"-ss",
|
||
f"00:00:00",
|
||
"-i",
|
||
video_path,
|
||
"-vframes",
|
||
"1",
|
||
"-vf",
|
||
f"select=eq(n\,{target_frame})",
|
||
"-vsync",
|
||
"0",
|
||
"-an",
|
||
"-y",
|
||
frame_path,
|
||
]
|
||
|
||
subprocess.run(cmd_extract, capture_output=True, check=True)
|
||
|
||
# 转换为Tensor
|
||
tensor = self.jpg_to_tensor(frame_path.replace("%03d", "001"))
|
||
except Exception as e:
|
||
raise e
|
||
return tensor
|
||
|
||
|
||
class JMGestureCorrect:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {"required": {"image": ("IMAGE",), "resolution": (["720p", "1080p"])}}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("正面图",)
|
||
FUNCTION = "gen"
|
||
CATEGORY = "不忘科技-自定义节点🚩/图片/姿态"
|
||
|
||
def gen(self, image: torch.Tensor, resolution: str):
|
||
wait_time = 240
|
||
interval = 2
|
||
client = JMUtils()
|
||
image_io = client.tensor_to_io(image)
|
||
upload_data = client.upload_io_to_cos(image_io)
|
||
if upload_data["status"]:
|
||
image_url = upload_data["data"]
|
||
else:
|
||
raise Exception("上传失败")
|
||
prompt = "Stand straight ahead, facing the camera, showing your full body, maintaining a proper posture, keeping the camera still, and ensuring that your head and feet are all within the frame"
|
||
submit_data = client.submit_task(
|
||
prompt, image_url, duration="5", resolution=resolution
|
||
)
|
||
if submit_data["status"]:
|
||
job_id = submit_data["data"]
|
||
else:
|
||
raise Exception("即梦任务提交失败")
|
||
job_data = None
|
||
for idx, _ in enumerate(range(0, wait_time, interval)):
|
||
logger.info(f"查询即梦结果 {idx + 1}")
|
||
query = client.query_status(job_id)
|
||
if query["status"]:
|
||
job_data = query["data"]
|
||
break
|
||
else:
|
||
if (
|
||
"error" in query["msg"]
|
||
or "失败" in query["msg"]
|
||
or "fail" in query["msg"]
|
||
):
|
||
raise Exception("即梦任务失败 {}".format(query["msg"]))
|
||
sleep(interval)
|
||
if not job_data:
|
||
raise Exception("即梦任务等待超时")
|
||
return (client.get_last_15th_frame_tensor(job_data),)
|
||
|
||
|
||
class JMCustom:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"prompt": (
|
||
"STRING",
|
||
{
|
||
"default": "Stand straight ahead, facing the camera, showing your full body, maintaining a proper posture, keeping the camera still, and ensuring that your head and feet are all within the frame",
|
||
"multiline": True,
|
||
},
|
||
),
|
||
"duration": ("INT", {"default": 5, "min": 2, "max": 10}),
|
||
"resolution": (["720p", "1080p"]),
|
||
"wait_time": ("INT", {"default": 180, "min": 60, "max": 600}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (
|
||
"STRING",
|
||
"IMAGE",
|
||
)
|
||
RETURN_NAMES = ("视频存储路径", "视频最后一帧")
|
||
FUNCTION = "gen"
|
||
CATEGORY = "不忘科技-自定义节点🚩/视频/即梦"
|
||
|
||
def gen(
|
||
self,
|
||
image: torch.Tensor,
|
||
prompt: str,
|
||
duration: int,
|
||
resolution: str,
|
||
wait_time: int,
|
||
):
|
||
interval = 2
|
||
client = JMUtils()
|
||
image_io = client.tensor_to_io(image)
|
||
upload_data = client.upload_io_to_cos(image_io)
|
||
if upload_data["status"]:
|
||
image_url = upload_data["data"]
|
||
else:
|
||
raise Exception("上传失败")
|
||
submit_data = client.submit_task(
|
||
prompt, image_url, str(duration), resolution=resolution
|
||
)
|
||
if submit_data["status"]:
|
||
job_id = submit_data["data"]
|
||
else:
|
||
raise Exception("即梦任务提交失败")
|
||
job_data = None
|
||
for idx, _ in enumerate(range(0, wait_time, interval)):
|
||
logger.info(f"查询即梦结果 {idx + 1}")
|
||
query = client.query_status(job_id)
|
||
if query["status"]:
|
||
job_data = query["data"]
|
||
break
|
||
else:
|
||
if (
|
||
"error" in query["msg"]
|
||
or "失败" in query["msg"]
|
||
or "fail" in query["msg"]
|
||
):
|
||
raise Exception("即梦任务失败 {}".format(query["msg"]))
|
||
sleep(interval)
|
||
if not job_data:
|
||
raise Exception("即梦任务等待超时")
|
||
output_dir = folder_paths.get_output_directory()
|
||
|
||
video_path, last_scene = client.download_video(
|
||
job_data, path=os.path.join(output_dir, f"{uuid.uuid4()}.mp4")
|
||
)
|
||
return (
|
||
video_path,
|
||
last_scene,
|
||
)
|