import torch import folder_paths import comfy.sd import comfy.utils import comfy.model_management as model_management import sys import gc import psutil from datetime import datetime import time from pathlib import Path class DDModelOptimizer: @classmethod def INPUT_TYPES(s): return { "required": { "模型文件": (folder_paths.get_filename_list("diffusion_models"), ), "智能模式": ("BOOLEAN", {"default": False}), "加载模式": ([ "标准加载", "分步加载", ],), "优化模式": ([ "禁用优化", # 改为更直观的名称 "FP8基础内存优化", # fp8_e4m3fn "FP8高速性能优化", # fp8_e4m3fn_fast "FP8稳定质量优化" # fp8_e5m2 ],), } } RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("优化模型",) FUNCTION = "optimize_model" CATEGORY = "🍺DD系列节点" def get_system_info(self): """获取系统配置信息""" system_info = { "total_memory": psutil.virtual_memory().total / (1024**3), # GB "available_memory": psutil.virtual_memory().available / (1024**3), # GB "cpu_count": psutil.cpu_count(logical=False), "gpu_info": None, } if torch.cuda.is_available(): system_info["gpu_info"] = { "name": torch.cuda.get_device_name(), "total_memory": torch.cuda.get_device_properties(0).total_memory / (1024**3), # GB "free_memory": torch.cuda.memory_allocated(0) / (1024**3), # GB } return system_info def determine_smart_options(self, system_info, model_path): """根据系统配置和模型智能确定最佳选项 Args: system_info: 系统信息 model_path: 模型路径 """ try: # 将字符串路径转换为Path对象 path_obj = Path(model_path) # 安全获取文件大小 if path_obj.exists(): model_size = path_obj.stat().st_size / (1024**3) # GB print(f"\n模型文件大小: {model_size:.2f}GB") else: print(f"警告: 模型文件不存在: {model_path}") model_size = 0 except Exception as e: print(f"警告: 无法获取模型文件信息: {str(e)}") model_size = 0 options = { "加载模式": "标准加载", "优化模式": "禁用优化" } # 获取显存和内存信息 vram = 0 ram = system_info["total_memory"] if system_info["gpu_info"]: vram = system_info["gpu_info"]["total_memory"] print(f"GPU总显存: {vram:.2f}GB") print(f"系统总内存: {ram:.2f}GB") # 设置缓冲系数 vram_scale_factor = 1.5 ram_scale_factor = 1.5 # 决定加载模式 # 如果模型大小大于可用显存的70%或大于显存总量，使用分步加载 if model_size > vram * 0.7: options["加载模式"] = "分步加载" print("由于模型大小超过可用显存的70%，选择分步加载模式") # 应用新的优化规则 if model_size <= vram: options["优化模式"] = "禁用优化" print("模型完全适配显存，获得最佳质量") elif model_size <= vram * vram_scale_factor and model_size <= ram: options["优化模式"] = "禁用优化" print("模型可通过显存+虚拟内存加载，保持最佳质量") elif model_size > vram * vram_scale_factor and model_size <= ram: options["优化模式"] = "FP8稳定质量优化" print("模型显著超出显存但适配内存，启动质量优先优化") elif model_size > ram and model_size <= ram * ram_scale_factor: options["优化模式"] = "FP8基础内存优化" print("模型超出物理内存，启动内存优化") else: # model_size > ram * ram_scale_factor options["优化模式"] = "FP8高速性能优化" print("模型严重超出物理内存，启动性能优先优化") return options def show_step_progress(self, step_name, progress=0): """显示步骤进度 Args: step_name: 步骤名称 progress: 进度值(0-100) """ bar_width = 30 filled = int(bar_width * progress / 100) bar = '#' * filled + '-' * (bar_width - filled) sys.stdout.write(f'\r{step_name} [{bar}] {progress}%') sys.stdout.flush() if progress >= 100: print() def optimize_model(self, 模型文件, 智能模式, 加载模式, 优化模式): """ 模型优化加载主函数 Args: 模型文件: 要加载的模型文件名 智能模式: 是否启用智能模式 加载模式: 标准加载或分步加载 优化模式: 禁用优化/FP8基础内存优化/FP8高速性能优化/FP8稳定质量优化 Returns: 优化后的模型 """ start_time = time.time() print(f"\n开始处理模型: {模型文件}") print(f"处理时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") try: # 获取模型路径 model_path = folder_paths.get_full_path_or_raise("diffusion_models", 模型文件) # 如果启用智能模式，重新确定加载模式和优化模式 if 智能模式: system_info = self.get_system_info() smart_options = self.determine_smart_options(system_info, model_path) 加载模式 = smart_options["加载模式"] 优化模式 = smart_options["优化模式"] print("\n智能模式已启用:") print(f"系统配置: {'GPU: ' + system_info['gpu_info']['name'] if system_info['gpu_info'] else 'CPU'}") print(f"自动选择 - 加载模式: {加载模式}") print(f"自动选择 - 优化模式: {优化模式}") # 设置优化选项 model_options = {} needs_optimization = True # 是否需要进行优化处理 if 优化模式 == "FP8基础内存优化": print("使用 FP8 基础内存优化模式") model_options["dtype"] = torch.float8_e4m3fn elif 优化模式 == "FP8高速性能优化": print("使用 FP8 高速性能优化模式") model_options["dtype"] = torch.float8_e4m3fn model_options["fp8_optimizations"] = True elif 优化模式 == "FP8稳定质量优化": print("使用 FP8 稳定质量优化模式") model_options["dtype"] = torch.float8_e5m2 else: print("优化已禁用，使用原始加载模式") needs_optimization = False if 加载模式 == "标准加载": print("使用标准加载模式...") self.show_step_progress("加载模型", 0) model = comfy.sd.load_diffusion_model(model_path, model_options=model_options) self.show_step_progress("加载模型", 100) else: print("使用分步加载模式...") # 第一步：加载模型文件 print("步骤1: 加载模型文件") self.show_step_progress("加载文件", 0) state_dict = comfy.utils.load_torch_file(model_path) self.show_step_progress("加载文件", 100) # 第二步：预处理权重 print("\n步骤2: 预处理权重") total_keys = len(state_dict) processed = 0 update_interval = max(1, total_keys // 100) # 确保至少显示100个更新点 if needs_optimization: # 只在需要优化时进行处理 for key, value in state_dict.items(): if torch.is_tensor(value): state_dict[key] = value.to(model_options["dtype"]) processed += 1 if processed % update_interval == 0: progress = int(processed / total_keys * 100) self.show_step_progress("处理权重", progress) self.show_step_progress("处理权重", 100) # 第三步：创建模型 print("\n步骤3: 构建模型") self.show_step_progress("构建模型", 0) model = comfy.sd.load_diffusion_model_state_dict(state_dict, model_options) self.show_step_progress("构建模型", 100) # 清理内存 gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() end_time = time.time() duration = end_time - start_time print(f"\n模型加载完成！用时: {duration:.2f}秒") return (model,) except Exception as e: print(f"\n模型加载失败: {str(e)}") raise e NODE_CLASS_MAPPINGS = { "DD-ModelOptimizer": DDModelOptimizer } NODE_DISPLAY_NAME_MAPPINGS = { "DD-ModelOptimizer": "DD Model Optimizer" }