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

ADD 增加视频帧位计算、字符串判空节点

Browse Source 
PERF 重构项目结构
This commit is contained in:
康宇佳
2025-03-04 16:50:54 +08:00
parent e1849bb8b9
commit cafa224374
14 changed files with 718 additions and 650 deletions
Download Patch File Download Diff File Expand all files Collapse all files

643
__init__.py
View File

@@ -1,630 +1,11 @@
import glob
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 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
video_extensions = ["webm", "mp4", "mkv", "gif", "mov"]
class FaceDetect:
"""
人脸遮挡检测
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"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},
),
},
}
RETURN_TYPES = (
"IMAGE",
"IMAGE",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"INT",
"INT",
)
RETURN_NAMES = (
"图像",
"选中人脸",
"分类",
"概率",
"采用帧序号",
"全部帧序列",
"剪辑配置",
"起始帧序号",
"帧数量",
)
FUNCTION = "predict"
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
)
print("全部帧序列", period)
if len(period) > 0:
start, end = period[main_seed % len(period)]
config = {"start": start, "end": end}
else:
raise RuntimeError("未找到符合要求的视频片段")
return (
image,
image_selected,
cls,
prob,
nums,
str(period),
json.dumps(config),
start,
end - start + 1,
)
class FaceExtract:
"""人脸提取 By YOLO"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("图片",)
FUNCTION = "crop"
CATEGORY = "不忘科技-自定义节点🚩"
def crop(self, image):
device = model_management.get_torch_device()
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)
print("aaaaa")
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
)
n = 512
r = results[0]
if len(r.boxes.data.cpu().numpy()) == 1:
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,
)
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[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))
out_images[idx] = img
idx += 1
else:
idx += 1
cropped_face = np.array(out_images).astype(np.float32) / 255.0
cropped_face = torch.from_numpy(cropped_face)
return (cropped_face,)
class COSDownload:
# TODO 增加腾讯云VOD视频下载
"""腾讯云COS下载"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"cos_key": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("视频存储路径",)
FUNCTION = "download"
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,
)
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:
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
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),
),
)
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),
),
)
class COSUpload:
"""腾讯云COS上传"""
# TODO 增加腾讯云VOD视频上传
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"path": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("COS文件Key",)
FUNCTION = "upload"
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:
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
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,
)
break
except CosClientError or CosServiceError as e:
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:
"""FFMPEG视频剪辑 -- ！有卡顿问题 暂废弃"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"config": ("STRING",),
"video_path": ("STRING",),
"mod": ("INT",),
"fps": ("FLOAT",),
"period_length": (
"INT",
{
"default": 10,
"min": 4,
"max": 100,
"step": 1,
"forceInput": True,
},
),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("视频路径",)
FUNCTION = "cut"
# OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩"
def cut(self, config, video_path, mod, fps, period_length):
# 原文件名
origin_fname = ".".join(video_path.split(os.sep)[-1].split(".")[:-1])
# 配置获取
mul = mod / fps
print("fps", fps)
config = json.loads(config)
if len(config.keys()) == 0:
return ("无法生成符合要求的片段",)
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]),
]
)
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
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",
]
},
)
print(ff.cmd)
ff.run()
# uuid填充改回原文件名
try:
os.remove(temp_fname)
except:
pass
try:
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", "*")
)
return (str(files),)
except:
files = glob.glob(output.replace("%03d", "*"))
traceback.print_exc()
return (str(files),)
# Add custom API routes, using router
from aiohttp import web
from server import PromptServer
@PromptServer.instance.routes.get("/hello")
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}")
from .nodes.string_empty_judgement import StringEmptyJudgement
from .nodes.compute_video_point import VideoStartPointDurationCompute
from .nodes.cos import COSUpload, COSDownload
from .nodes.face_detect import FaceDetect
from .nodes.face_extract import FaceExtract
from .nodes.log2db import LogToDB
from .nodes.videocut import VideoCut
from .nodes.vod2local import VodToLocalNode
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
@@ -635,7 +16,9 @@ NODE_CLASS_MAPPINGS = {
"COSDownload": COSDownload,
"VideoCutCustom": VideoCut,
"VodToLocal": VodToLocalNode,
"LogToDB": LogToDB
"LogToDB": LogToDB,
"VideoPointCompute": VideoStartPointDurationCompute,
"StringEmptyJudgement": StringEmptyJudgement,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -646,5 +29,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"COSDownload": "COS下载",
"VideoCutCustom": "视频剪裁",
"VodToLocal": "腾讯云VOD下载",
"LogToDB": "状态持久化DB"
"LogToDB": "状态持久化DB",
"VideoPointCompute": "视频帧位计算",
"StringEmptyJudgement": "字符串是否为空",
}

44
nodes/compute_video_point.py Normal file
View File

@@ -0,0 +1,44 @@
import re
from datetime import datetime
from math import ceil
def validate_time_format(time_str):
pattern = r'^([0-1][0-9]|2[0-3]):([0-5][0-9]):([0-5][0-9]|\d{1,2})$'
return bool(re.match(pattern, time_str))
def get_duration_wave(audio):
waveform, sample_rate = audio["waveform"], audio["sample_rate"]
return waveform.shape[2] / sample_rate
class VideoStartPointDurationCompute:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"start_time": ("STRING", {"forceInput": True}),
"audio": ("AUDIO", {"forceInput": True}),
"fps": ("INT", {"default": 25, "step": 1}),
},
}
RETURN_TYPES = ("INT", "INT",)
RETURN_NAMES = ("起始帧位", "帧数")
FUNCTION = "compute"
CATEGORY = "不忘科技-自定义节点🚩"
def compute(self, start_time, audio, fps):
if not validate_time_format(start_time):
raise ValueError("start_time或者end_time时间格式不对（start_time or end_time is not in time format）")
time_format = "%H:%M:%S"
start_dt = datetime.strptime(start_time, time_format)
start_sec = (start_dt - datetime(1900, 1, 1)).total_seconds()
start_point = int(start_sec * fps)
print("audio duration %.2f s"%get_duration_wave(audio))
duration = ceil(get_duration_wave(audio) * fps)
return (start_point, duration,)

145
nodes/cos.py Normal file
View File

@@ -0,0 +1,145 @@
import json
import os
import urllib
import yaml
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
class COSDownload:
"""腾讯云COS下载"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"cos_key": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("视频存储路径",)
FUNCTION = "download"
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,
)
for i in range(0, 10):
try:
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)
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),
),
)
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),
),
)
class COSUpload:
"""腾讯云COS上传"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"path": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("COS文件Key",)
FUNCTION = "upload"
CATEGORY = "不忘科技-自定义节点🚩"
def upload(self, path):
for i in range(0, 10):
try:
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)
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,
)
break
except CosClientError or CosServiceError as e:
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],
]
),
)

76
nodes/face_detect.py Normal file
View File

@@ -0,0 +1,76 @@
import json
from .test_single_image import test_node
class FaceDetect:
"""
人脸遮挡检测
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"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},
),
},
}
RETURN_TYPES = (
"IMAGE",
"IMAGE",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"INT",
"INT",
)
RETURN_NAMES = (
"图像",
"选中人脸",
"分类",
"概率",
"采用帧序号",
"全部帧序列",
"剪辑配置",
"起始帧序号",
"帧数量",
)
FUNCTION = "predict"
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
)
print("全部帧序列", period)
if len(period) > 0:
start, end = period[main_seed % len(period)]
config = {"start": start, "end": end}
else:
raise RuntimeError("未找到符合要求的视频片段")
return (
image,
image_selected,
cls,
prob,
nums,
str(period),
json.dumps(config),
start,
end - start + 1,
)

77
nodes/face_extract.py Normal file
View File

@@ -0,0 +1,77 @@
import os
import cv2
import numpy as np
import torch
from comfy import model_management
from ultralytics import YOLO
class FaceExtract:
"""人脸提取 By YOLO"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("图片",)
FUNCTION = "crop"
CATEGORY = "不忘科技-自定义节点🚩"
def crop(self, image):
device = model_management.get_torch_device()
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)
print("aaaaa")
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
)
n = 512
r = results[0]
if len(r.boxes.data.cpu().numpy()) == 1:
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,
)
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[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))
out_images[idx] = img
idx += 1
else:
idx += 1
cropped_face = np.array(out_images).astype(np.float32) / 255.0
cropped_face = torch.from_numpy(cropped_face)
return (cropped_face,)

66
nodes/log2db.py Normal file
View File

@@ -0,0 +1,66 @@
import json
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import server
from .table import Task
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),)}

1
model.py → nodes/model.py
View File

@@ -3,7 +3,6 @@ from PIL import ImageFile
from .utils import get_model
ImageFile.LOAD_TRUNCATED_IMAGES = True

21
nodes/string_empty_judgement.py Normal file
View File

@@ -0,0 +1,21 @@
class StringEmptyJudgement:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"input": ("STRING", {"forceInput": True}),
},
}
RETURN_TYPES = ("BOOLEAN", )
RETURN_NAMES = ("是否为空", )
FUNCTION = "compute"
CATEGORY = "不忘科技-自定义节点🚩"
def compute(self, input):
if len(input) == 0:
return (True, )
else:
return (False, )

18
nodes/table.py Normal file
View File

@@ -0,0 +1,18 @@
from sqlalchemy import Column, Integer, DateTime, String, func
from sqlalchemy.orm import declarative_base
Base = declarative_base()
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}"

6
test_single_image.py → nodes/test_single_image.py
View File

@@ -10,7 +10,6 @@ from torchvision.transforms import Resize
from .utils import load_weight
from .model import Model
# CONSTANT
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
@@ -56,6 +55,7 @@ def test_image(opt):
opt.image, CLASSES[pred.item()], prob.item() * 100
))
def test_node(image: torch.Tensor, length=10, thres=95, model_name="convnext_tiny"):
weight_dic = {
"convnext_tiny": "best_convnext_tiny.pth",
@@ -102,8 +102,8 @@ def test_node(image:torch.Tensor,length=10,thres=95,model_name="convnext_tiny"):
while a + length - 1 <= i[1]:
temp_period.append([a, a + length - 1])
a = a + length
return (image.permute(0,2,3,1), image.permute(0,2,3,1)[nums,:,:,:], str(preds), str(probs), str(nums), temp_period)
return (
image.permute(0, 2, 3, 1), image.permute(0, 2, 3, 1)[nums, :, :, :], str(preds), str(probs), str(nums), temp_period)
if __name__ == "__main__":

0
utils.py → nodes/utils.py
View File

124
nodes/videocut.py Normal file
View File

@@ -0,0 +1,124 @@
import glob
import json
import os
import shutil
import traceback
import uuid
from datetime import datetime
import ffmpy
class VideoCut:
"""FFMPEG视频剪辑 -- ！有卡顿问题 暂废弃"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"config": ("STRING",),
"video_path": ("STRING",),
"mod": ("INT",),
"fps": ("FLOAT",),
"period_length": (
"INT",
{
"default": 10,
"min": 4,
"max": 100,
"step": 1,
"forceInput": True,
},
),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("视频路径",)
FUNCTION = "cut"
# OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩"
def cut(self, config, video_path, mod, fps, period_length):
# 原文件名
origin_fname = ".".join(video_path.split(os.sep)[-1].split(".")[:-1])
# 配置获取
mul = mod / fps
print("fps", fps)
config = json.loads(config)
if len(config.keys()) == 0:
return ("无法生成符合要求的片段",)
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]),
]
)
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
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",
]
},
)
print(ff.cmd)
ff.run()
# uuid填充改回原文件名
try:
os.remove(temp_fname)
except:
pass
try:
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", "*")
)
return (str(files),)
except:
files = glob.glob(output.replace("%03d", "*"))
traceback.print_exc()
return (str(files),)

112
nodes/vod2local.py Normal file
View File

@@ -0,0 +1,112 @@
import os
from pathlib import Path
import requests
from tencentcloud.common import credential
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.vod.v20180717 import vod_client, models
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}")

1
requirements.txt
View File

@@ -5,3 +5,4 @@ Pillow
opencv-python
ultralytics
cos-python-sdk-v5
sqlalchemy