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

FIX job_id直接获取extra_data

Browse Source
This commit is contained in:
康宇佳
2025-03-03 11:45:14 +08:00
parent 16e842dc53
commit 092af340fb
1 changed files with 420 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

557
__init__.py
View File

@@ -3,126 +3,107 @@ import json
import os
import shutil
import traceback
import urllib.request
import uuid
from datetime import datetime
import server
import cv2
import ffmpy
import numpy as np
import torch
import yaml
from ultralytics import YOLO
from comfy import model_management
from qcloud_cos import CosConfig, CosClientError, CosServiceError
from qcloud_cos import CosS3Client
from sqlalchemy import Column, Integer, func, DateTime, ForeignKey, String, create_engine
from sqlalchemy.orm import sessionmaker
from ultralytics import YOLO
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
from .test_single_image import test_node
import ffmpy
video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov']
video_extensions = ["webm", "mp4", "mkv", "gif", "mov"]
class FaceDetect:
"""
A example node
Class methods
-------------
INPUT_TYPES (dict):
Tell the main program input parameters of nodes.
IS_CHANGED:
optional method to control when the node is re executed.
Attributes
----------
RETURN_TYPES (`tuple`):
The type of each element in the output tuple.
RETURN_NAMES (`tuple`):
Optional: The name of each output in the output tuple.
FUNCTION (`str`):
The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute()
OUTPUT_NODE ([`bool`]):
If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example.
The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected.
Assumed to be False if not present.
CATEGORY (`str`):
The category the node should appear in the UI.
DEPRECATED (`bool`):
Indicates whether the node is deprecated. Deprecated nodes are hidden by default in the UI, but remain
functional in existing workflows that use them.
EXPERIMENTAL (`bool`):
Indicates whether the node is experimental. Experimental nodes are marked as such in the UI and may be subject to
significant changes or removal in future versions. Use with caution in production workflows.
execute(s) -> tuple || None:
The entry point method. The name of this method must be the same as the value of property `FUNCTION`.
For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`.
人脸遮挡检测
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
"""
Return a dictionary which contains config for all input fields.
Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
Input types "INT", "STRING" or "FLOAT" are special values for fields on the node.
The type can be a list for selection.
Returns: `dict`:
- Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
- Value input_fields (`dict`): Contains input fields config:
* Key field_name (`string`): Name of a entry-point method's argument
* Value field_config (`tuple`):
+ First value is a string indicate the type of field or a list for selection.
+ Second value is a config for type "INT", "STRING" or "FLOAT".
"""
return {
"required": {
"image": ("IMAGE",),
"main_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"main_seed": (
"INT:seed",
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
),
"model": (["convnext_tiny", "convnext_base"],),
"length": ("INT", {"default": 10, "min": 3, "max": 60, "step": 1}),
"threshold": ("FLOAT", {"default": 94, "min": 55, "max": 99, "step": 0.1})
"threshold": (
"FLOAT",
{"default": 94, "min": 55, "max": 99, "step": 0.1},
),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "STRING", "STRING", "STRING", "STRING", "INT", "INT")
RETURN_NAMES = ("图像", "选中人脸", "分类", "概率", "采用帧序号", "全部帧序列", "剪辑配置", "起始帧序号", "帧数量")
RETURN_TYPES = (
"IMAGE",
"IMAGE",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"INT",
"INT",
)
RETURN_NAMES = (
"图像",
"选中人脸",
"分类",
"概率",
"采用帧序号",
"全部帧序列",
"剪辑配置",
"起始帧序号",
"帧数量",
)
FUNCTION = "predict"
# OUTPUT_NODE = False
CATEGORY = "自定义节点"
CATEGORY = "不忘科技-自定义节点🚩"
def predict(self, image, main_seed, model, length, threshold):
image, image_selected, cls, prob, nums, period = test_node(image, length=length, thres=threshold,
model_name=model)
image, image_selected, cls, prob, nums, period = test_node(
image, length=length, thres=threshold, model_name=model
)
print("全部帧序列", period)
if len(period) > 0:
start, end = period[main_seed % len(period)]
config = {"start": start, "end": end}
else:
config = {}
start = 0
end = 0
raise RuntimeError("未找到符合要求的视频片段")
return (image, image_selected, cls, prob, nums, str(period), json.dumps(config), start, end - start)
"""
The node will always be re executed if any of the inputs change but
this method can be used to force the node to execute again even when the inputs don't change.
You can make this node return a number or a string. This value will be compared to the one returned the last time the node was
executed, if it is different the node will be executed again.
This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash
changes between executions the LoadImage node is executed again.
"""
# @classmethod
# def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
# return ""
return (
image,
image_selected,
cls,
prob,
nums,
str(period),
json.dumps(config),
start,
end - start + 1,
)
class FaceExtract():
class FaceExtract:
"""人脸提取 By YOLO"""
@classmethod
def INPUT_TYPES(s):
return {
@@ -136,12 +117,18 @@ class FaceExtract():
FUNCTION = "crop"
CATEGORY = "自定义节点"
CATEGORY = "不忘科技-自定义节点🚩"
def crop(self, image):
device = model_management.get_torch_device()
image_np = 255. * image.cpu().numpy()
model = YOLO(model=os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", "yolov8n-face-lindevs.pt"))
image_np = 255.0 * image.cpu().numpy()
model = YOLO(
model=os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"model",
"yolov8n-face-lindevs.pt",
)
)
total_images = image_np.shape[0]
out_images = np.ndarray(shape=(total_images, 512, 512, 3))
print("shape", image_np.shape)
@@ -149,12 +136,7 @@ class FaceExtract():
idx = 0
for image_item in image_np:
results = model.predict(
image_item,
imgsz=640,
conf=0.75,
iou=0.7,
device=device,
verbose=False
image_item, imgsz=640, conf=0.75, iou=0.7, device=device, verbose=False
)
n = 512
r = results[0]
@@ -162,13 +144,21 @@ class FaceExtract():
y1, x1, y2, x2, p, cls = r.boxes.data.cpu().numpy()[0]
face_size = int(max(y2 - y1, x2 - x1))
center = (x1 + x2) // 2, (y1 + y2) // 2
x1, x2, y1, y2 = center[0] - face_size // 2, center[0] + face_size // 2, center[1] - face_size // 2, \
center[1] + face_size // 2
x1, x2, y1, y2 = (
center[0] - face_size // 2,
center[0] + face_size // 2,
center[1] - face_size // 2,
center[1] + face_size // 2,
)
template = np.ndarray(shape=(face_size, face_size, 3))
template.fill(20)
for a, a1 in zip(list(range(int(x1), int(x2))), list(range(face_size))):
for b, b1 in zip(list(range(int(y1), int(y2))), list(range(face_size))):
if (a >= 0 and a <= r.orig_img.shape[1]) and (b >= 0 and b <= r.orig_img.shape[0]):
for b, b1 in zip(
list(range(int(y1), int(y2))), list(range(face_size))
):
if (a >= 0 and a < r.orig_img.shape[0]) and (
b >= 0 and b < r.orig_img.shape[1]
):
template[a1][b1] = r.orig_img[a][b]
print(int(x1), int(x2), int(y1), int(y2))
img = cv2.resize(template, (n, n))
@@ -182,6 +172,9 @@ class FaceExtract():
class COSDownload:
# TODO 增加腾讯云VOD视频下载
"""腾讯云COS下载"""
@classmethod
def INPUT_TYPES(s):
return {
@@ -193,33 +186,62 @@ class COSDownload:
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("视频存储路径",)
FUNCTION = "download"
CATEGORY = "自定义节点"
CATEGORY = "不忘科技-自定义节点🚩"
def download(self, cos_key):
if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
os.makedirs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key)),
exist_ok=True)
os.makedirs(
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
os.path.dirname(cos_key),
),
exist_ok=True,
)
for i in range(0, 10):
try:
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8",
mode="r+") as f:
with open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)), "config.yaml"
),
encoding="utf-8",
mode="r+",
) as f:
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"],
SecretKey=yaml_config["secret_key"])
config = CosConfig(
Region=yaml_config["region"],
SecretId=yaml_config["secret_id"],
SecretKey=yaml_config["secret_key"],
)
client = CosS3Client(config)
response = client.download_file(
Bucket=yaml_config["bucket"],
Key=cos_key,
DestFilePath=os.path.join(os.path.dirname(os.path.abspath(__file__)), "download",
os.path.dirname(cos_key), os.path.basename(cos_key)))
DestFilePath=os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
os.path.dirname(cos_key),
os.path.basename(cos_key),
),
)
break
except CosClientError or CosServiceError as e:
print(f"下载失败 {e}")
return (os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key),
os.path.basename(cos_key)),)
return (
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
os.path.dirname(cos_key),
os.path.basename(cos_key),
),
)
class COSUpload:
"""腾讯云COS上传"""
# TODO 增加腾讯云VOD视频上传
@classmethod
def INPUT_TYPES(s):
return {
@@ -232,32 +254,141 @@ class COSUpload:
RETURN_NAMES = ("COS文件Key",)
FUNCTION = "upload"
CATEGORY = "自定义节点"
CATEGORY = "不忘科技-自定义节点🚩"
def upload(self, path):
for i in range(0, 10):
try:
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8",
mode="r+") as f:
with open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)), "config.yaml"
),
encoding="utf-8",
mode="r+",
) as f:
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"],
SecretKey=yaml_config["secret_key"])
config = CosConfig(
Region=yaml_config["region"],
SecretId=yaml_config["secret_id"],
SecretKey=yaml_config["secret_key"],
)
client = CosS3Client(config)
response = client.upload_file(
Bucket=yaml_config["bucket"],
Key="/".join(
[yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]),
LocalFilePath=path)
[
yaml_config["subfolder"],
(
path.split("/")[-1]
if "/" in path
else path.split("\\")[-1]
),
]
),
LocalFilePath=path,
)
break
except CosClientError or CosServiceError as e:
print(e)
return ("/".join([yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]),)
raise RuntimeError("上传失败")
data = {"prompt_id": "",
"video_url": "https://{}.cos.{}.myqcloud.com/{}".format(yaml_config['bucket'], yaml_config['region'],
'/'.join([yaml_config['subfolder'],
path.split('/')[
-1] if '/' in path else
path.split('\\')[-1], ]))
}
headers = {'Content-Type': 'application/json'}
try:
req = urllib.request.Request("", data=json.dumps(data).encode("utf-8"), headers=headers)
response = urllib.request.urlopen(req)
except:
raise RuntimeError("上报MQ状态失败")
return (
"/".join(
[
yaml_config["subfolder"],
path.split("/")[-1] if "/" in path else path.split("\\")[-1],
]
),
)
class Task(Base):
__tablename__ = 'task'
id = Column(Integer, primary_key=True)
gmt_create = Column(DateTime(timezone=True), server_default=func.now())
gmt_modified = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())
prompt_id = Column(String, index=True, nullable=False, unique=True)
result = Column(String, nullable=True)
job_id = Column(Integer, index=True, nullable=False, unique=True)
status = Column(Integer)
def __repr__(self):
return f"{self.id},{self.gmt_create},{self.gmt_modified},{self.prompt_id},{self.result}"
class LogToDB:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"job_id": ("STRING",{"forceInput": True}),
"log": ("STRING",{"forceInput": True}),
"status": ("INT",{"default": 1, "max": 1}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "log2db"
OUTPUT_NODE = True
OUTPUT_IS_LIST = (True,)
# OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩"
def log2db(self, log, status, unique_id):
# 获取comfy服务器队列信息
(_, prompt_id, prompt, extra_data, outputs_to_execute) = next(
iter(server.PromptServer.instance.prompt_queue.currently_running.values()))
job_id = extra_data["client_id"]
engine = create_engine(
"mysql+pymysql://root:*k3&5xxG6oqHJM@sh-cdb-1xspb808.sql.tencentcdb.com:28795/comfy",
echo=True
)
# Base.metadata.create_all(engine)
session = sessionmaker(bind=engine)()
# 查询
tasks = session.query(Task).filter(Task.prompt_id == prompt_id).all()
print(prompt)
result = {
"curr_node_id": str(unique_id),
"last_node_id": list(prompt.keys())[-1],
"node_output": str(log)
}
if len(tasks) == 0:
# 不存在插入
task = Task(prompt_id=prompt_id, job_id=job_id, result=json.dumps(result), status=status)
session.add(task)
elif len(tasks) == 1:
# 存在更新
session.query(Task).filter(Task.prompt_id == prompt_id).update({"result": json.dumps(result),
"status": status})
else:
# 异常报错
raise RuntimeError("状态数据库prompt_id不唯一, 无法记录状态！")
session.commit()
return {"ui": {"text": json.dumps(result)}, "result": (json.dumps(result),)}
# 有问题
class VideoCut:
def __init__(self):
pass
"""FFMPEG视频剪辑 -- ！有卡顿问题 暂废弃"""
@classmethod
def INPUT_TYPES(s):
@@ -267,7 +398,16 @@ class VideoCut:
"video_path": ("STRING",),
"mod": ("INT",),
"fps": ("FLOAT",),
"period_length": ("INT", {"default": 10, "min": 4, "max": 100, "step": 1, "forceInput": True})
"period_length": (
"INT",
{
"default": 10,
"min": 4,
"max": 100,
"step": 1,
"forceInput": True,
},
),
},
}
@@ -278,7 +418,7 @@ class VideoCut:
# OUTPUT_NODE = False
CATEGORY = "自定义节点"
CATEGORY = "不忘科技-自定义节点🚩"
def cut(self, config, video_path, mod, fps, period_length):
# 原文件名
@@ -292,31 +432,54 @@ class VideoCut:
start, end = config["start"], config["end"]
# 新文件名 复制改名适配ffmpeg
uid = uuid.uuid1()
temp_fname = os.sep.join([*video_path.split(os.sep)[:-1], "%s.%s" % (str(uid), video_path.split(".")[-1])])
temp_fname = os.sep.join(
[
*video_path.split(os.sep)[:-1],
"%s.%s" % (str(uid), video_path.split(".")[-1]),
]
)
try:
shutil.copy(video_path, temp_fname)
except:
return ("请检查输入文件权限",)
video_path = temp_fname
# 组装输出文件名
output_name = (".".join([*video_path.split(os.sep)[-1].split(".")[:-2],
video_path.split(os.sep)[-1].split(".")[-2] +
"_output_%%03d_%s" % datetime.now().strftime('%Y%m%d_%H%M%S'),
video_path.split(os.sep)[-1].split(".")[-1]]))
output = (os.sep.join([*video_path.split(os.sep)[:-1], output_name])
.replace(os.sep.join(["ComfyUI", "input"]), os.sep.join(["ComfyUI", "output"])).replace(" ", ""))
#调用ffmpeg
output_name = ".".join(
[
*video_path.split(os.sep)[-1].split(".")[:-2],
video_path.split(os.sep)[-1].split(".")[-2]
+ "_output_%%03d_%s" % datetime.now().strftime("%Y%m%d_%H%M%S"),
video_path.split(os.sep)[-1].split(".")[-1],
]
)
output = (
os.sep.join([*video_path.split(os.sep)[:-1], output_name])
.replace(
os.sep.join(["ComfyUI", "input"]), os.sep.join(["ComfyUI", "output"])
)
.replace(" ", "")
)
# 调用ffmpeg
ff = ffmpy.FFmpeg(
inputs={video_path: ['-accurate_seek']},
outputs={output: [
'-f', 'segment',
'-ss', str(round(start * mul, 3)),
'-to', str(round(end * mul, 3)),
'-segment_times', str(period_length),
'-c', 'copy',
'-map', '0',
'-avoid_negative_ts', '1'
]}
inputs={video_path: ["-accurate_seek"]},
outputs={
output: [
"-f",
"segment",
"-ss",
str(round(start * mul, 3)),
"-to",
str(round(end * mul, 3)),
"-segment_times",
str(period_length),
"-c",
"copy",
"-map",
"0",
"-avoid_negative_ts",
"1",
]
},
)
print(ff.cmd)
ff.run()
@@ -329,7 +492,9 @@ class VideoCut:
files = glob.glob(output.replace("%03d", "*"))
for file in files:
shutil.move(file, file.replace(str(uid), origin_fname))
files = glob.glob(output.replace(str(uid), origin_fname).replace("%03d", "*"))
files = glob.glob(
output.replace(str(uid), origin_fname).replace("%03d", "*")
)
return (str(files),)
except:
files = glob.glob(output.replace("%03d", "*"))
@@ -347,6 +512,120 @@ async def get_hello(request):
return web.json_response("hello")
# 腾讯云 VOD
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.common import credential
from tencentcloud.vod.v20180717 import vod_client, models
import requests
from pathlib import Path
class VodToLocalNode:
def __init__(self):
self.secret_id = "AKIDsrihIyjZOBsjimt8TsN8yvv1AMh5dB44"
self.secret_key = "CPZcxdk6W39Jd4cGY95wvupoyMd0YFqW"
self.vod_client = self.init_vod_client()
def init_vod_client(self):
"""初始化VOD客户端"""
try:
http_profile = HttpProfile(endpoint="vod.tencentcloudapi.com")
client_profile = ClientProfile(httpProfile=http_profile)
cred = credential.Credential(self.secret_id, self.secret_key)
return vod_client.VodClient(cred, "ap-shanghai", client_profile)
except Exception as e:
raise RuntimeError(f"VOD client initialization failed: {e}")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"file_id": ("STRING", {"default": ""}),
"sub_app_id": ("STRING", {"default": ""}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("local_path",)
FUNCTION = "execute"
CATEGORY = "video"
def execute(self, file_id, sub_app_id):
# 调用下载逻辑
local_path = self.download_vod(file_id, sub_app_id)
print(f"下载成功: {local_path}")
return (local_path,)
def _get_download_url(self, file_id, sub_app_id):
"""获取媒体文件下载地址"""
try:
req = models.DescribeMediaInfosRequest()
req.FileIds = [file_id]
req.SubAppId = int(sub_app_id)
resp = self.vod_client.DescribeMediaInfos(req)
if not resp.MediaInfoSet:
raise ValueError("File not found")
media_info = resp.MediaInfoSet[0]
if not media_info.BasicInfo.MediaUrl:
raise ValueError("No download URL available")
return media_info.BasicInfo.MediaUrl
except Exception as e:
raise RuntimeError(f"Tencent API error: {e}")
def create_directory(self, path):
p = Path(path)
if not p.exists():
p.mkdir(
parents=True, exist_ok=True
) # parents=True会自动创建所有必需的父目录，exist_ok=True表示如果目录已存在则不会引发异常
print(f"目录已创建: {path}")
else:
print(f"目录已存在: {path}")
def download_vod(self, file_id, sub_app_id):
"""
需要补充腾讯云VOD SDK调用逻辑
返回本地文件路径
"""
media_url = self._get_download_url(file_id=file_id, sub_app_id=sub_app_id)
print(f"download from url: {media_url}")
# 生成一个临时目录路径名并创建该目录
self.create_directory(
os.path.join(
os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}"
)
)
output_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
f"{sub_app_id}",
f"{file_id}.mp4",
)
# 判断文件是否存在
if os.path.exists(output_dir):
return output_dir
return self._download_file(url=media_url, save_path=output_dir, timeout=60 * 10)
def _download_file(self, url: str, save_path: str, timeout: int = 30):
"""下载文件到本地"""
try:
with requests.get(url, stream=True, timeout=timeout) as response:
response.raise_for_status()
with open(save_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
return save_path
except Exception as e:
raise RuntimeError(f"Download error: {e}")
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
@@ -354,7 +633,9 @@ NODE_CLASS_MAPPINGS = {
"FaceExtract": FaceExtract,
"COSUpload": COSUpload,
"COSDownload": COSDownload,
"VideoCutCustom": VideoCut
"VideoCutCustom": VideoCut,
"VodToLocal": VodToLocalNode,
"LogToDB": LogToDB
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -363,5 +644,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"FaceExtract": "面部提取",
"COSUpload": "COS上传",
"COSDownload": "COS下载",
"VideoCutCustom": "视频剪裁"
"VideoCutCustom": "视频剪裁",
"VodToLocal": "腾讯云VOD下载",
"LogToDB": "状态持久化DB"
}