diff --git a/nodes/image_face_nodes.py b/nodes/image_face_nodes.py index 2e379fc..c58bea8 100644 --- a/nodes/image_face_nodes.py +++ b/nodes/image_face_nodes.py @@ -7,6 +7,7 @@ import torch from comfy import model_management from ultralytics import YOLO +from ..utils.download_utils import download_file from ..utils.face_occu_detect import face_occu_detect @@ -21,7 +22,7 @@ class FaceDetect: "required": { "image": ("IMAGE",), "main_seed": ( - "INT:seed", + "INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}, ), "model": (["convnext_tiny", "convnext_base"],), @@ -104,12 +105,15 @@ class FaceExtract: 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_path = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "model", "yolov8n-face-lindevs.pt", ) + if not os.path.exists(model_path): + download_file("https://github.com/lindevs/yolov8-face/releases/latest/download/yolov8n-face-lindevs.pt",model_path) + model = YOLO( + model=model_path ) total_images = image_np.shape[0] out_images = np.ndarray(shape=(total_images, 512, 512, 3)) diff --git a/nodes/video_nodes.py b/nodes/video_nodes.py index 428a108..6c2f218 100644 --- a/nodes/video_nodes.py +++ b/nodes/video_nodes.py @@ -119,7 +119,7 @@ class VideoCut: os.remove(files[0]) except: pass - return (video, {"waveform": audio, "sample_rate": info["audio_fps"]},) + return (video/255.0, {"waveform": audio, "sample_rate": info["audio_fps"]} if "audio_fps" in info else None,) except: traceback.print_exc() raise Exception("Cut Failed") @@ -234,7 +234,7 @@ class VideoCutByFramePoint: os.remove(output) except: pass - return (video, {"waveform": audio, "sample_rate": info["audio_fps"]},) + return (video/255.0, {"waveform": audio, "sample_rate": info["audio_fps"]},) except: traceback.print_exc() raise Exception("Cut Failed") diff --git a/utils/download_utils.py b/utils/download_utils.py new file mode 100644 index 0000000..cb3bc41 --- /dev/null +++ b/utils/download_utils.py @@ -0,0 +1,146 @@ +import os +import requests +import threading +from tqdm import tqdm +from urllib.parse import urlparse + + +def download_file(url, output_path=None, num_threads=8, chunk_size=1024 * 1024): + """ + 多线程下载文件 + + 参数: + url (str): 下载URL + output_path (str, optional): 输出文件路径，默认为URL中的文件名 + num_threads (int, optional): 线程数，默认为8 + chunk_size (int, optional): 每个线程下载的块大小(字节)，默认为1MB + + 返回: + bool: 下载成功返回True，失败返回False + """ + try: + # 获取文件名 + if not output_path: + output_path = os.path.basename(urlparse(url).path) + if not output_path: + output_path = "downloaded_file" + + # 检查服务器是否支持范围请求 + response = requests.head(url) + response.raise_for_status() + + # 检查是否支持断点续传 + supports_range = 'Accept-Ranges' in response.headers and response.headers['Accept-Ranges'] == 'bytes' + + if not supports_range: + print("服务器不支持多线程下载，将使用单线程下载") + return _download_single_thread(url, output_path) + + # 获取文件大小 + file_size = int(response.headers.get('Content-Length', 0)) + if not file_size: + print("无法获取文件大小，将使用单线程下载") + return _download_single_thread(url, output_path) + + # 创建临时文件 + temp_files = [f"{output_path}.part{i}" for i in range(num_threads)] + + # 计算每个线程下载的范围 + ranges = [] + for i in range(num_threads): + start = i * (file_size // num_threads) + end = start + (file_size // num_threads) - 1 if i < num_threads - 1 else file_size - 1 + ranges.append((start, end)) + + # 创建进度条 + progress = tqdm(total=file_size, unit='B', unit_scale=True, desc="下载中") + + # 启动线程 + threads = [] + for i in range(num_threads): + start, end = ranges[i] + thread = threading.Thread(target=_download_chunk, args=(url, temp_files[i], start, end, progress)) + thread.start() + threads.append(thread) + + # 等待所有线程完成 + for thread in threads: + thread.join() + + progress.close() + + # 合并临时文件 + if not _merge_files(temp_files, output_path): + print("合并文件失败") + return False + + # 删除临时文件 + for temp_file in temp_files: + os.remove(temp_file) + + print(f"文件下载完成: {output_path}") + return True + + except Exception as e: + print(f"下载过程中发生错误: {e}") + # 清理临时文件 + for temp_file in temp_files: + if os.path.exists(temp_file): + os.remove(temp_file) + return False + + +def _download_single_thread(url, output_path): + """单线程下载文件""" + try: + response = requests.get(url, stream=True) + response.raise_for_status() + + with open(output_path, 'wb') as f, \ + tqdm(desc="下载中", total=int(response.headers.get('Content-Length', 0)), + unit='B', unit_scale=True) as progress: + for chunk in response.iter_content(chunk_size=1024): + if chunk: + f.write(chunk) + progress.update(len(chunk)) + + print(f"文件下载完成: {output_path}") + return True + except Exception as e: + print(f"单线程下载失败: {e}") + return False + + +def _download_chunk(url, output_path, start, end, progress): + """下载文件的指定块""" + headers = {'Range': f'bytes={start}-{end}'} + try: + response = requests.get(url, headers=headers, stream=True) + response.raise_for_status() + + with open(output_path, 'wb') as f: + for chunk in response.iter_content(chunk_size=1024): + if chunk: + f.write(chunk) + progress.update(len(chunk)) + except Exception as e: + print(f"下载块失败 {start}-{end}: {e}") + + +def _merge_files(temp_files, output_path): + """合并临时文件""" + try: + with open(output_path, 'wb') as outfile: + for temp_file in temp_files: + with open(temp_file, 'rb') as infile: + outfile.write(infile.read()) + return True + except Exception as e: + print(f"合并文件失败: {e}") + return False + + +# 使用示例 +if __name__ == "__main__": + download_url = "https://example.com/large_file.zip" # 替换为实际的下载URL + download_file(download_url, num_threads=4) \ No newline at end of file diff --git a/utils/face_occu_detect.py b/utils/face_occu_detect.py index 67446ad..40c0fd4 100644 --- a/utils/face_occu_detect.py +++ b/utils/face_occu_detect.py @@ -63,7 +63,9 @@ def face_occu_detect(image: torch.Tensor, length=10, thres=95, model_name="convn } device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = Model(model_name, 2, False).to(device) - weight = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", weight_dic[model_name]) + weight = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "model", weight_dic[model_name]) + if not os.path.exists(weight): + raise Exception("请前往https://github.com/LamKser/face-occlusion-classification下载所选权重文件到model文件夹！") model = load_weight(model, weight) model.eval() image = image.permute(0, 3, 1, 2)