comi-logic/ComfyUI-CustomNode
7
0
Fork 0
Code Pull Requests Releases Activity

refactor

Browse Source
This commit is contained in:
iHeyTang
2025-08-25 14:42:35 +08:00
parent 467c8b2549
commit f41fd24316
17 changed files with 4353 additions and 943 deletions
Download Patch File Download Diff File Expand all files Collapse all files

506
nodes/image_gesture_nodes.py
View File

@@ -6,42 +6,90 @@ 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
import yaml
from PIL import Image
from loguru import logger
from qcloud_cos import CosConfig, CosS3Client
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):
if "aws_key_id" in list(os.environ.keys()):
yaml_config = os.environ
else:
with open(
os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yaml"
),
encoding="utf-8",
mode="r+",
) as f:
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
"""
初始化即梦工具实例
self.api_key = yaml_config["jm_api_key"]
self.cos_region = yaml_config["cos_region"]
self.cos_secret_id = yaml_config["cos_secret_id"]
self.cos_secret_key = yaml_config["cos_secret_key"]
self.cos_bucket_name = yaml_config["cos_sucai_bucket_name"]
def submit_task(self, prompt: str, img_url: str, duration: str = "10", resolution:str="720p"):
从配置中读取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}",
@@ -52,80 +100,244 @@ class JMUtils:
"content": [
{
"type": "text",
"text": f"{prompt} --resolution {resolution} --dur {duration} --camerafixed false",
"text": f"{prompt.strip()} --resolution {resolution} --dur {duration} --camerafixed false",
},
{
"type": "image_url",
"image_url": {
"url": img_url,
"url": img_url.strip(),
},
},
],
}
response = requests.post("https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks",
headers=headers, json=json_data)
logger.info(f"submit task: {json.dumps(response.json())}")
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:
return {"status": False, "data": img_url, "msg": resp_json["error"]["message"]}
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(e)
logger.error(f"即梦任务提交异常: {e}")
return {"data": None, "status": False, "msg": str(e)}
def query_status(self, job_id: str):
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}",
headers=headers)
resp_json = response.json()
resp_dict["status"] = resp_json["status"] == "succeeded"
resp_dict["msg"] = resp_json["status"]
resp_dict["data"] = resp_json["content"]["video_url"] if "content" in resp_json else None
except Exception as e:
logger.error(f"error:{str(e)}")
resp_dict["msg"] = str(e)
finally:
return resp_dict
def upload_io_to_cos(self, file: io.IOBase, mime_type: str = "image/png"):
resp_data = {'status': True, 'data': '', 'msg': ''}
category = mime_type.split('/')[0]
suffix = mime_type.split('/')[1]
try:
object_key = f'tk/{category}/{uuid.uuid4()}.{suffix}'
config = CosConfig(Region=self.cos_region, SecretId=self.cos_secret_id, SecretKey=self.cos_secret_key)
client = CosS3Client(config)
_ = client.upload_file_from_buffer(
Bucket=self.cos_bucket_name,
Key=object_key,
Body=file
response = requests.get(
f"https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks/{job_id.strip()}",
headers=headers,
timeout=15,
)
url = f'https://{self.cos_bucket_name}.cos.{self.cos_region}.myqcloud.com/{object_key}'
resp_data['data'] = url
resp_data['msg'] = '上传成功'
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(e)
resp_data['status'] = False
resp_data['msg'] = str(e)
logger.error(f"上传文件时发生异常: {e}")
resp_data["status"] = False
resp_data["msg"] = str(e)
return resp_data
def tensor_to_io(srlf, tensor: torch.Tensor):
# 转换为PIL图像
img = Image.fromarray(np.clip(255. * tensor.cpu().squeeze().numpy(), 0, 255).astype(np.uint8))
image_data = io.BytesIO()
img.save(image_data, format='PNG')
image_data.seek(0)
return image_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:
"""
@@ -170,11 +382,13 @@ class JMUtils:
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 转换为PyTorch张量并调整维度为BCHW
transform = transforms.Compose([
transforms.ToTensor() # 转换为[C, H, W]格式的张量，值范围从0到1
])
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]
tensor = (
transform(frame_rgb).unsqueeze(0).permute(0, 2, 3, 1)
) # 添加批次维度，变为[1, H, W, C]
return tensor
@@ -186,7 +400,7 @@ class JMUtils:
if path:
temp_path = path
else:
temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
temp_path = temp_file.name
temp_file.close()
@@ -196,16 +410,16 @@ class JMUtils:
response.raise_for_status()
# 获取文件大小
total_size = int(response.headers.get('content-length', 0))
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
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)
@@ -221,21 +435,19 @@ class JMUtils:
else:
raise
def jpg_to_tensor(self, image_path, channel_first=False):
def jpg_to_tensor(self, image_path):
"""
将JPG图像转换为PyTorch张量
参数:
- image_path: JPG图像文件路径
- normalize: 是否将像素值归一化到[0.0, 1.0]
- channel_first: 是否将通道维度放在前面 (C, H, W)
返回:
- tensor: PyTorch张量
"""
try:
# 打开图像文件
image = Image.open(image_path).convert('RGB')
image = Image.open(image_path).convert("RGB")
# 转换为张量
tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
@@ -246,7 +458,7 @@ class JMUtils:
print(f"转换失败: {str(e)}")
raise
def get_last_15th_frame_tensor(self, video_url, cleanup=True):
def get_last_15th_frame_tensor(self, video_url):
"""
从视频URL截取倒数第15帧并转换为Tensor
先下载视频到本地临时文件再处理
@@ -257,18 +469,20 @@ class JMUtils:
# 获取视频总帧数
cmd_frames = [
'ffprobe', '-v', 'error',
'-select_streams', 'v:0',
'-show_entries', 'stream=nb_frames',
'-of', 'default=nokey=1:noprint_wrappers=1',
video_path
"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
cmd_frames, capture_output=True, text=True, check=True
)
# 处理可能的非数字输出
@@ -276,19 +490,14 @@ class JMUtils:
if not frame_count.isdigit():
# 备选方案：通过解码获取帧数
print("无法获取准确帧数，尝试直接解码...")
cmd_decode = [
'ffmpeg', '-i', video_path,
'-f', 'null', '-'
]
cmd_decode = ["ffmpeg", "-i", video_path, "-f", "null", "-"]
decode_result = subprocess.run(
cmd_decode,
capture_output=True,
text=True
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]
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
@@ -302,25 +511,29 @@ class JMUtils:
print(f"视频总帧数: {frame_count}, 目标帧: {target_frame}")
# 截取指定帧
with tempfile.NamedTemporaryFile(suffix='%03d.jpg', delete=True) as frame_file:
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
"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
)
subprocess.run(cmd_extract, capture_output=True, check=True)
# 转换为Tensor
tensor = self.jpg_to_tensor(frame_path.replace("%03d", "001"))
@@ -332,12 +545,7 @@ class JMUtils:
class JMGestureCorrect:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"resolution":(["720p","1080p"])
}
}
return {"required": {"image": ("IMAGE",), "resolution": (["720p", "1080p"])}}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("正面图",)
@@ -355,7 +563,9 @@ class JMGestureCorrect:
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)
submit_data = client.submit_task(
prompt, image_url, duration="5", resolution=resolution
)
if submit_data["status"]:
job_id = submit_data["data"]
else:
@@ -368,7 +578,11 @@ class JMGestureCorrect:
job_data = query["data"]
break
else:
if "error" in query["msg"] or "失败" in query["msg"] or "fail" in query["msg"]:
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:
@@ -382,21 +596,35 @@ class JMCustom:
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}),
"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_TYPES = (
"STRING",
"IMAGE",
)
RETURN_NAMES = ("视频存储路径", "视频最后一帧")
FUNCTION = "gen"
CATEGORY = "不忘科技-自定义节点🚩/视频/即梦"
def gen(self, image: torch.Tensor, prompt: str, duration: int, resolution: str, wait_time: int):
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)
@@ -405,7 +633,9 @@ class JMCustom:
image_url = upload_data["data"]
else:
raise Exception("上传失败")
submit_data = client.submit_task(prompt, image_url, str(duration), resolution=resolution)
submit_data = client.submit_task(
prompt, image_url, str(duration), resolution=resolution
)
if submit_data["status"]:
job_id = submit_data["data"]
else:
@@ -418,11 +648,21 @@ class JMCustom:
job_data = query["data"]
break
else:
if "error" in query["msg"] or "失败" in query["msg"] or "fail" in query["msg"]:
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("即梦任务等待超时")
video_path, last_scene = client.download_video(job_data, path=os.path.join(folder_paths.get_output_directory(),
f"{uuid.uuid4()}.mp4"))
return (video_path, last_scene,)
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,
)