551 lines
25 KiB
Python
551 lines
25 KiB
Python
import json
|
|
|
|
from agents.tool import function_tool
|
|
from ..utils.modelscope_gateway import ModelScopeGateway
|
|
from ..utils.request_context import get_session_id
|
|
|
|
from ..utils.comfy_gateway import get_object_info_by_class
|
|
from ..dao.workflow_table import get_workflow_data, save_workflow_data
|
|
from ..utils.logger import log
|
|
|
|
async def get_node_parameters(node_name: str, param_name: str = "") -> str:
|
|
"""获取节点的参数信息，如果param_name为空则返回所有参数"""
|
|
try:
|
|
node_info_dict = await get_object_info_by_class(node_name)
|
|
if not node_info_dict or node_name not in node_info_dict:
|
|
return json.dumps({"error": f"Node '{node_name}' not found"})
|
|
|
|
node_info = node_info_dict[node_name]
|
|
if 'input' not in node_info:
|
|
return json.dumps({"error": f"Node '{node_name}' has no input parameters"})
|
|
|
|
input_params = node_info['input']
|
|
|
|
if param_name:
|
|
# 检查特定参数
|
|
if input_params.get('required') and param_name in input_params['required']:
|
|
return json.dumps({
|
|
"parameter": param_name,
|
|
"type": "required",
|
|
"config": input_params['required'][param_name]
|
|
})
|
|
|
|
if input_params.get('optional') and param_name in input_params['optional']:
|
|
return json.dumps({
|
|
"parameter": param_name,
|
|
"type": "optional",
|
|
"config": input_params['optional'][param_name]
|
|
})
|
|
|
|
return json.dumps({"error": f"Parameter '{param_name}' not found in node '{node_name}'"})
|
|
else:
|
|
# 返回所有参数
|
|
return json.dumps({
|
|
"node": node_name,
|
|
"required": input_params.get('required', {}),
|
|
"optional": input_params.get('optional', {})
|
|
})
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to get node parameters: {str(e)}"})
|
|
|
|
@function_tool
|
|
async def find_matching_parameter_value(node_name: str, param_name: str, current_value: str, error_info: str = "") -> str:
|
|
"""根据错误信息找到匹配的参数值，支持多种参数类型的智能处理"""
|
|
try:
|
|
# 获取参数配置
|
|
param_info_str = await get_node_parameters(node_name, param_name)
|
|
param_info = json.loads(param_info_str)
|
|
|
|
if "error" in param_info:
|
|
return json.dumps(param_info)
|
|
|
|
param_config = param_info.get("config", [])
|
|
error_lower = error_info.lower()
|
|
|
|
# 检查错误类型并提供相应处理策略
|
|
error_analysis = {
|
|
"error_type": "unknown",
|
|
"is_model_related": False,
|
|
"is_file_related": False,
|
|
}
|
|
|
|
# 优先识别image文件相关错误（在model检测之前，因为image文件可能包含model关键词）
|
|
if (any(img_ext in current_value.lower() for img_ext in [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp"]) or
|
|
param_name.lower() in ["image", "img", "picture", "photo"] or
|
|
any(img_keyword in error_lower for img_keyword in ["invalid image", "image file", "image not found"])):
|
|
error_analysis["error_type"] = "image_file_missing"
|
|
error_analysis["is_file_related"] = True
|
|
error_analysis["can_auto_fix"] = True
|
|
|
|
# 如果参数配置是列表，查找其他可用的图片
|
|
if isinstance(param_config, list) and len(param_config) > 0 and param_config[0]:
|
|
available_images = [img for img in param_config[0] if any(ext in str(img).lower() for ext in [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp"])]
|
|
|
|
if available_images:
|
|
# 随机选择一张可用图片
|
|
import random
|
|
recommended_image = random.choice(available_images)
|
|
return json.dumps({
|
|
"found_match": True,
|
|
"error_type": "image_file_missing",
|
|
"solution_type": "auto_replace",
|
|
"recommended_value": recommended_image,
|
|
"match_type": "image_replacement",
|
|
"message": f"Replaced missing image '{current_value}' with available image '{recommended_image}'",
|
|
"can_auto_fix": True,
|
|
"next_action": "update_parameter"
|
|
})
|
|
|
|
return json.dumps({
|
|
"found_match": False,
|
|
"error_type": "image_file_missing",
|
|
"solution_type": "manual_fix",
|
|
"message": f"Missing image file: {current_value}",
|
|
"suggestion": "Please add a valid image file to your ComfyUI input folder or choose an existing one",
|
|
"can_auto_fix": False
|
|
})
|
|
|
|
# 识别model相关错误（在image检测之后）
|
|
elif (any(model_keyword in current_value.lower() for model_keyword in [
|
|
".ckpt", ".safetensors", ".pt", ".pth", ".bin", "checkpoint", "lora", "vae", "controlnet", "clip", "unet"
|
|
]) or any(model_keyword in param_name.lower() for model_keyword in [
|
|
"model", "checkpoint", "lora", "vae", "clip", "controlnet", "unet"
|
|
]) or any(model_keyword in error_lower for model_keyword in [
|
|
"model not found", "checkpoint", "missing model", "file not found"
|
|
])):
|
|
error_analysis["error_type"] = "model_missing"
|
|
error_analysis["is_model_related"] = True
|
|
|
|
return json.dumps({
|
|
"found_match": False,
|
|
"error_type": "model_missing",
|
|
"message": f"Missing model file: {current_value}",
|
|
"recommendation": "Check for available model replacements, otherwise download required.",
|
|
"next_action": "check_available_models_or_suggest_download",
|
|
"details": {
|
|
"node_name": node_name,
|
|
"param_name": param_name,
|
|
"missing_file": current_value,
|
|
"is_model_related": True,
|
|
"param_config": param_config
|
|
}
|
|
})
|
|
|
|
# 处理枚举类型的参数（原有逻辑，但增强）
|
|
elif isinstance(param_config, list) and len(param_config) > 0:
|
|
available_values = param_config
|
|
error_analysis["error_type"] = "enum_value_mismatch"
|
|
error_analysis["can_auto_fix"] = True
|
|
|
|
# 改进的匹配算法
|
|
current_lower = current_value.lower().replace("_", " ").replace("-", " ")
|
|
|
|
# 1. 完全匹配
|
|
for value in available_values:
|
|
if current_value == value:
|
|
return json.dumps({
|
|
"found_match": True,
|
|
"recommended_value": value,
|
|
"match_type": "exact",
|
|
"error_type": "enum_value_mismatch",
|
|
"solution_type": "exact_match",
|
|
"can_auto_fix": True,
|
|
"all_available": available_values
|
|
})
|
|
|
|
# 2. 忽略大小写和符号的匹配
|
|
for value in available_values:
|
|
value_lower = str(value).lower().replace("_", " ").replace("-", " ")
|
|
if current_lower == value_lower:
|
|
return json.dumps({
|
|
"found_match": True,
|
|
"recommended_value": value,
|
|
"match_type": "case_insensitive",
|
|
"error_type": "enum_value_mismatch",
|
|
"solution_type": "auto_replace",
|
|
"message": f"Found case-insensitive match: '{current_value}' -> '{value}'",
|
|
"can_auto_fix": True,
|
|
"next_action": "update_parameter",
|
|
"all_available": available_values
|
|
})
|
|
|
|
# 3. 包含关系匹配
|
|
best_match = None
|
|
best_score = 0
|
|
|
|
for value in available_values:
|
|
value_lower = str(value).lower()
|
|
value_parts = value_lower.replace("_", " ").replace("-", " ").split()
|
|
current_parts = current_lower.split()
|
|
|
|
# 计算匹配分数
|
|
score = 0
|
|
for part in current_parts:
|
|
if any(part in vp or vp in part for vp in value_parts):
|
|
score += 1
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
best_match = value
|
|
|
|
if best_match and best_score > 0:
|
|
return json.dumps({
|
|
"found_match": True,
|
|
"recommended_value": best_match,
|
|
"match_type": "partial",
|
|
"error_type": "enum_value_mismatch",
|
|
"solution_type": "auto_replace",
|
|
"match_score": best_score,
|
|
"message": f"Found partial match: '{current_value}' -> '{best_match}' (score: {best_score})",
|
|
"can_auto_fix": True,
|
|
"next_action": "update_parameter",
|
|
"all_available": available_values[:10],
|
|
"original_value": current_value
|
|
})
|
|
|
|
# 4. 没有匹配，但可以用第一个可用值替代
|
|
return json.dumps({
|
|
"found_match": False,
|
|
"recommended_value": available_values[0] if available_values else None,
|
|
"match_type": "no_match",
|
|
"error_type": "enum_value_mismatch",
|
|
"solution_type": "default_replace",
|
|
"message": f"No match found for '{current_value}'. Using default value '{available_values[0]}'",
|
|
"can_auto_fix": True,
|
|
"next_action": "update_parameter",
|
|
"all_available": available_values[:10],
|
|
"total_options": len(available_values),
|
|
"original_value": current_value,
|
|
"suggestion": f"No match found for '{current_value}'. Replacing with default option."
|
|
})
|
|
|
|
# 处理其他类型的参数
|
|
else:
|
|
error_analysis["error_type"] = "non_enum_parameter"
|
|
error_analysis["can_auto_fix"] = False
|
|
|
|
return json.dumps({
|
|
"found_match": False,
|
|
"error_type": "non_enum_parameter",
|
|
"solution_type": "manual_fix",
|
|
"parameter_type": type(param_config).__name__,
|
|
"config": param_config,
|
|
"can_auto_fix": False,
|
|
"message": f"Parameter '{param_name}' is not an enumerable type",
|
|
"suggestion": f"Parameter '{param_name}' requires manual configuration. Check the parameter requirements."
|
|
})
|
|
|
|
except Exception as e:
|
|
return json.dumps({
|
|
"error": f"Failed to find matching parameter value: {str(e)}",
|
|
"can_auto_fix": False,
|
|
"solution_type": "error"
|
|
})
|
|
|
|
@function_tool
|
|
async def get_model_files(model_type: str = "checkpoints") -> str:
|
|
"""获取可用的模型文件列表"""
|
|
try:
|
|
# 定义模型类型到节点的映射
|
|
model_type_mapping = {
|
|
"checkpoints": ["CheckpointLoaderSimple", "CheckpointLoader"],
|
|
"loras": ["LoraLoader", "LoraLoaderModelOnly"],
|
|
"vae": ["VAELoader"],
|
|
"clip": ["CLIPLoader", "DualCLIPLoader"],
|
|
"controlnet": ["ControlNetLoader", "ControlNetApply"],
|
|
"unet": ["UNETLoader"],
|
|
"ipadapter": ["IPAdapterModelLoader"]
|
|
}
|
|
|
|
# 查找对应的节点
|
|
model_files = {}
|
|
for node_name in model_type_mapping.get(model_type.lower(), []):
|
|
try:
|
|
# 使用 get_object_info_by_class 获取单个节点信息，减少数据量
|
|
node_data = await get_object_info_by_class(node_name)
|
|
|
|
# 处理不同的返回格式
|
|
node_info = None
|
|
if node_name in node_data:
|
|
# 格式：{"NodeName": {...}}
|
|
node_info = node_data[node_name]
|
|
elif 'input' in node_data:
|
|
# 格式：直接返回节点信息 {...}
|
|
node_info = node_data
|
|
|
|
if node_info and 'input' in node_info:
|
|
# 查找包含文件列表的参数
|
|
for input_type in ['required', 'optional']:
|
|
if input_type in node_info['input']:
|
|
for param_name, param_config in node_info['input'][input_type].items():
|
|
# 检查参数配置格式：[file_list, {...}] 或 [file_list]
|
|
if isinstance(param_config, list) and len(param_config) > 0:
|
|
if isinstance(param_config[0], list) and len(param_config[0]) > 0:
|
|
# 检查是否为文件列表（包含文件扩展名或路径）
|
|
file_list = param_config[0]
|
|
if any(isinstance(item, str) and ('.' in item or '/' in item) for item in file_list):
|
|
model_files[f"{node_name}.{param_name}"] = file_list
|
|
|
|
except Exception as e:
|
|
# 单个节点查询失败，继续处理其他节点
|
|
log.error(f"Failed to get info for node {node_name}: {e}")
|
|
continue
|
|
|
|
if model_files:
|
|
return json.dumps({
|
|
"model_type": model_type,
|
|
"available_models": model_files
|
|
})
|
|
else:
|
|
return json.dumps({
|
|
"model_type": model_type,
|
|
"available_models": {},
|
|
"message": f"No {model_type} models found. Please check your ComfyUI models folder."
|
|
})
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to get model files: {str(e)}"})
|
|
|
|
|
|
def suggest_model_download_by_modelscope(model_name_keyword: str) -> str:
|
|
"""建议下载缺失的模型，执行一次即可结束流程返回结果"""
|
|
modelscope_gateway = ModelScopeGateway()
|
|
return modelscope_gateway.suggest(name=model_name_keyword)
|
|
|
|
|
|
@function_tool
|
|
def suggest_model_download(models_list: str = "") -> str:
|
|
"""
|
|
建议下载缺失的模型，执行一次即可结束流程返回结果，支持批量处理
|
|
|
|
Args:
|
|
models_list: 缺失模型列表的JSON字符串，格式为：
|
|
'[{"model_type":"checkpoints","missing_model":"model.safetensors","model_name_keyword":"sd"}]'
|
|
其中model_type: 模型类型，取值范围为checkpoints, clip, clip_vision, configs, controlnet, diffusers, diffusion_models, embeddings, gligen, hypernetworks, loras, photomaker, style_models, text_encoders, unet, upscale_models, vae, vae_approx
|
|
missing_model: 缺失的模型名称
|
|
model_name_keyword: 模型名称的关键词，用于模糊查询，短词优先，不要携带文件类型后缀
|
|
|
|
Returns:
|
|
str: 建议下载的模型列表json字符串
|
|
|
|
"""
|
|
try:
|
|
if not models_list or not models_list.strip():
|
|
return json.dumps({"error": "models_list is required"})
|
|
|
|
# 解析输入的模型列表
|
|
try:
|
|
# 尝试解析为列表
|
|
if models_list.strip().startswith('['):
|
|
models_data = json.loads(models_list)
|
|
else:
|
|
# 单个模型对象，转换为列表
|
|
models_data = [json.loads(models_list)]
|
|
except json.JSONDecodeError as e:
|
|
return json.dumps({"error": f"Invalid JSON format: {str(e)}"})
|
|
|
|
if not isinstance(models_data, list):
|
|
return json.dumps({"error": "models_list should be a list or single model object"})
|
|
|
|
# 存储所有模型建议结果
|
|
all_model_suggestions = []
|
|
missing_models_info = []
|
|
failed_models = []
|
|
|
|
# 遍历每个模型进行查询
|
|
for model_item in models_data:
|
|
if not isinstance(model_item, dict):
|
|
failed_models.append({"model": str(model_item), "error": "Invalid model format"})
|
|
continue
|
|
|
|
model_type = model_item.get("model_type", "")
|
|
missing_model = model_item.get("missing_model", "")
|
|
model_name_keyword = model_item.get("model_name_keyword", "")
|
|
|
|
if not model_type or not missing_model:
|
|
failed_models.append({
|
|
"model": missing_model or "unknown",
|
|
"error": "model_type and missing_model are required"
|
|
})
|
|
continue
|
|
|
|
# 如果没有提供关键词，从模型名中提取
|
|
if not (model_name_keyword and model_name_keyword.strip()):
|
|
missing_model_parts = missing_model.split(".")
|
|
if len(missing_model_parts) > 1:
|
|
model_name_keyword = missing_model_parts[0]
|
|
else:
|
|
model_name_keyword = missing_model
|
|
|
|
# 记录缺失模型信息
|
|
missing_models_info.append({
|
|
"model_type": model_type,
|
|
"missing_model": missing_model,
|
|
"model_name_keyword": model_name_keyword
|
|
})
|
|
|
|
# 优先使用modelscope检索模型
|
|
try:
|
|
result = suggest_model_download_by_modelscope(model_name_keyword)
|
|
if result and result.get("data") and len(result.get("data")) > 0:
|
|
# 为每个结果添加上下文信息
|
|
for suggestion in result["data"]:
|
|
suggestion["source_model_type"] = model_type
|
|
suggestion["source_missing_model"] = missing_model
|
|
suggestion["source_keyword"] = model_name_keyword
|
|
for item in result["data"]:
|
|
item["source_model_type"] = model_type
|
|
all_model_suggestions.extend(result["data"])
|
|
else:
|
|
# ModelScope没有找到结果，记录失败信息
|
|
failed_models.append({
|
|
"model": missing_model,
|
|
"model_type": model_type,
|
|
"keyword": model_name_keyword,
|
|
"error": "No models found in ModelScope"
|
|
})
|
|
except Exception as e:
|
|
failed_models.append({
|
|
"model": missing_model,
|
|
"model_type": model_type,
|
|
"keyword": model_name_keyword,
|
|
"error": f"ModelScope query failed: {str(e)}"
|
|
})
|
|
|
|
# 先按模型分组，每个模型取top3，然后再去重
|
|
grouped_suggestions = {}
|
|
for suggestion in all_model_suggestions:
|
|
source_model = suggestion.get("source_missing_model", "unknown")
|
|
if source_model not in grouped_suggestions:
|
|
grouped_suggestions[source_model] = []
|
|
grouped_suggestions[source_model].append(suggestion)
|
|
|
|
# 每个模型取top3
|
|
top3_per_model = []
|
|
for model, suggestions in grouped_suggestions.items():
|
|
suggestions_top3 = suggestions[:3]
|
|
for item in suggestions_top3:
|
|
item["model_type"] = model_type
|
|
top3_per_model.extend(suggestions_top3)
|
|
|
|
# 基于模型名称去重
|
|
seen_models = set()
|
|
unique_suggestions = []
|
|
for suggestion in top3_per_model:
|
|
# 使用Path和Name组合作为唯一标识
|
|
model_id = f"{suggestion.get('Path', '')}/{suggestion.get('Name', '')}"
|
|
if model_id not in seen_models:
|
|
seen_models.add(model_id)
|
|
unique_suggestions.append(suggestion)
|
|
|
|
# 最终建议列表
|
|
top_suggestions = unique_suggestions
|
|
|
|
# 构建返回结果
|
|
result = {
|
|
"success": True,
|
|
"total_requested": len(models_data),
|
|
"found_suggestions": len(top_suggestions),
|
|
"failed_count": len(failed_models),
|
|
"missing_models": missing_models_info,
|
|
"message": f"Found {len(top_suggestions)} model suggestions from ModelScope (top 3, deduplicated)",
|
|
"ext": [{
|
|
"type": "param_update",
|
|
"data": {
|
|
"model_suggest": top_suggestions,
|
|
"failed_models": failed_models,
|
|
"summary": {
|
|
"total_requested": len(models_data),
|
|
"found_suggestions": len(top_suggestions),
|
|
"failed_count": len(failed_models)
|
|
}
|
|
}
|
|
}]
|
|
}
|
|
|
|
# 如果没有找到任何建议，添加兜底信息
|
|
if not top_suggestions:
|
|
result["fallback_message"] = "No models found in ModelScope. Consider manual download from Hugging Face or Civitai."
|
|
result["general_suggestions"] = [
|
|
"1. Check model names and keywords for typos",
|
|
"2. Try broader keywords (e.g., 'stable-diffusion' instead of specific version)",
|
|
"3. Download from Hugging Face: https://huggingface.co/models",
|
|
"4. Download from Civitai: https://civitai.com/",
|
|
"5. Place models in appropriate ComfyUI/models/ subfolders"
|
|
]
|
|
|
|
return json.dumps(result)
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to suggest model download: {str(e)}"})
|
|
|
|
@function_tool
|
|
def update_workflow_parameter(node_id: str, param_name: str, new_value: str) -> str:
|
|
"""更新工作流中的特定参数"""
|
|
try:
|
|
session_id = get_session_id()
|
|
if not session_id:
|
|
log.error("update_workflow_parameter: No session_id found in context")
|
|
return json.dumps({"error": "No session_id found in context"})
|
|
|
|
# 获取当前工作流
|
|
workflow_data = get_workflow_data(session_id)
|
|
if not workflow_data:
|
|
return json.dumps({"error": "No workflow data found for this session"})
|
|
|
|
# 检查节点是否存在
|
|
if node_id not in workflow_data:
|
|
return json.dumps({"error": f"Node {node_id} not found in workflow"})
|
|
|
|
# 更新参数
|
|
if "inputs" not in workflow_data[node_id]:
|
|
workflow_data[node_id]["inputs"] = {}
|
|
|
|
old_value = workflow_data[node_id]["inputs"].get(param_name, "not set")
|
|
workflow_data[node_id]["inputs"][param_name] = new_value
|
|
|
|
# 保存更新的工作流到数据库
|
|
save_workflow_data(
|
|
session_id,
|
|
workflow_data,
|
|
workflow_data_ui=None, # UI format not available here
|
|
attributes={
|
|
"action": "parameter_update",
|
|
"description": f"Updated {param_name} in node {node_id}",
|
|
"changes": {
|
|
"node_id": node_id,
|
|
"parameter": param_name,
|
|
"old_value": old_value,
|
|
"new_value": new_value
|
|
}
|
|
}
|
|
)
|
|
|
|
return json.dumps({
|
|
"success": True,
|
|
"answer": f"Successfully updated {param_name} from '{old_value}' to '{new_value}' in node {node_id}",
|
|
"node_id": node_id,
|
|
"parameter": param_name,
|
|
"old_value": old_value,
|
|
"new_value": new_value,
|
|
"message": f"Successfully updated {param_name} from '{old_value}' to '{new_value}' in node {node_id}",
|
|
# 添加ext数据用于前端实时更新画布
|
|
"ext": [{
|
|
"type": "param_update",
|
|
"data": {
|
|
"workflow_data": workflow_data,
|
|
"changes": [{ # 包装成数组格式，与前端MessageList期望的格式匹配
|
|
"node_id": node_id,
|
|
"parameter": param_name,
|
|
"old_value": old_value,
|
|
"new_value": new_value
|
|
}]
|
|
}
|
|
}]
|
|
})
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to update workflow parameter: {str(e)}"})
|
|
|