fix: 修复ComfyUI队列API响应格式解析问题

- 修复WebSocket错误消息解析失败问题,支持更多错误字段格式
- 改进执行错误处理,收到错误时立即中断而不是等待超时
- 修复ComfyUI队列API响应格式不匹配问题,支持数组格式的队列数据
- 添加灵活的队列项解析逻辑,支持整数和字符串类型的prompt_id
- 在EventEmitter中添加错误存储功能,支持实时错误状态检查
- 更新所有相关的队列状态转换代码

解决的问题:
1. WebSocket消息解析错误:missing field 'message'
2. 执行错误后仍等待超时的问题
3. 队列API响应解析错误:invalid type integer/string expected
4. 批量处理中错误统计不准确的问题
This commit is contained in:
imeepos
2025-08-11 11:09:27 +08:00
parent b1e7191c10
commit 2075a731d2
9 changed files with 224 additions and 17 deletions

View File

@@ -361,7 +361,8 @@ pub struct QueueTaskInfo {
impl QueueModel {
/// 从 SDK 队列状态创建模型
pub fn from_queue_status(queue_status: &QueueStatus) -> Self {
let running_tasks: Vec<QueueTaskInfo> = queue_status.queue_running
let running_tasks: Vec<QueueTaskInfo> = queue_status
.get_running_items()
.iter()
.map(|item| QueueTaskInfo {
prompt_id: item.prompt_id.clone(),
@@ -373,7 +374,8 @@ impl QueueModel {
})
.collect();
let pending_tasks: Vec<QueueTaskInfo> = queue_status.queue_pending
let pending_tasks: Vec<QueueTaskInfo> = queue_status
.get_pending_items()
.iter()
.map(|item| QueueTaskInfo {
prompt_id: item.prompt_id.clone(),

View File

@@ -465,15 +465,18 @@ fn convert_system_info(info: comfyui_sdk::types::SystemStats) -> SystemInfoRespo
/// 转换队列状态
fn convert_queue_status(status: comfyui_sdk::types::QueueStatus) -> QueueStatusResponse {
let running_items = status.get_running_items();
let pending_items = status.get_pending_items();
QueueStatusResponse {
running_count: status.queue_running.len() as u32,
pending_count: status.queue_pending.len() as u32,
running_tasks: status.queue_running.into_iter().map(|item| QueueTaskResponse {
running_count: running_items.len() as u32,
pending_count: pending_items.len() as u32,
running_tasks: running_items.into_iter().map(|item| QueueTaskResponse {
prompt_id: item.prompt_id,
number: item.number,
status: "running".to_string(),
}).collect(),
pending_tasks: status.queue_pending.into_iter().map(|item| QueueTaskResponse {
pending_tasks: pending_items.into_iter().map(|item| QueueTaskResponse {
prompt_id: item.prompt_id,
number: item.number,
status: "pending".to_string(),

View File

@@ -417,6 +417,7 @@ pub async fn ai_model_face_hair_fix_batch_images(
match ai_model_face_hair_fix_internal(input_path, &server_url, &face_prompt, &face_denoise, &output_path_str, &window).await {
Ok(_) => {
let execution_time = file_start_time.elapsed().as_secs_f64();
println!("✅ 图片处理成功: {} -> {}", input_path, output_path_str);
results.push(ImageEnhancementResult {
success: true,
message: "图片增强成功".to_string(),
@@ -429,6 +430,8 @@ pub async fn ai_model_face_hair_fix_batch_images(
}
Err(e) => {
let execution_time = file_start_time.elapsed().as_secs_f64();
println!("❌ 图片处理失败: {} - 错误: {}", input_path, e);
eprintln!("详细错误信息: {}", e);
results.push(ImageEnhancementResult {
success: false,
message: format!("图片增强失败: {}", e),
@@ -444,6 +447,13 @@ pub async fn ai_model_face_hair_fix_batch_images(
let total_execution_time = start_time.elapsed().as_secs_f64();
// 打印最终统计结果
println!("📊 批量处理完成统计:");
println!(" 总计: {}", total_files);
println!(" 成功: {}", successful_count);
println!(" 失败: {}", failed_count);
println!(" 总耗时: {:.2}", total_execution_time);
let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
current: total_files,
total: total_files,
@@ -452,7 +462,7 @@ pub async fn ai_model_face_hair_fix_batch_images(
current_file: None,
});
Ok(BatchImageEnhancementResult {
let batch_result = BatchImageEnhancementResult {
success: failed_count == 0,
message: format!("批量处理完成: 总计 {}, 成功 {}, 失败 {}", total_files, successful_count, failed_count),
total_processed: total_files,
@@ -460,7 +470,10 @@ pub async fn ai_model_face_hair_fix_batch_images(
failed_count,
results,
total_execution_time,
})
};
println!("🎯 返回批量处理结果: {:?}", batch_result);
Ok(batch_result)
}
/// AI模型面部头发修复内部处理函数
@@ -472,6 +485,9 @@ async fn ai_model_face_hair_fix_internal(
output_image: &str,
window: &tauri::Window,
) -> Result<(), String> {
println!("🔄 开始处理图片: {}", input_image);
println!("📋 处理参数: server_url={}, face_prompt={}, face_denoise={}", server_url, face_prompt, face_denoise);
// 验证输入图片文件
if !Path::new(input_image).exists() {
return Err(format!("输入图片文件不存在: {}", input_image));
@@ -481,16 +497,19 @@ async fn ai_model_face_hair_fix_internal(
validate_image_file(input_image)?;
// 初始化 ComfyUI 客户端
println!("🔌 初始化ComfyUI客户端...");
let mut client = ComfyUIClient::new(ComfyUIClientConfig {
base_url: server_url.to_string(),
..Default::default()
}).map_err(|e| format!("创建ComfyUI客户端失败: {}", e))?;
// 连接到服务器
println!("🌐 连接到ComfyUI服务器: {}", server_url);
client.connect().await
.map_err(|e| format!("连接ComfyUI服务器失败: {}", e))?;
// 注册模板
println!("📝 注册AI模型面部头发修复模板...");
client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())
.map_err(|e| format!("注册模板失败: {}", e))?;
@@ -499,8 +518,10 @@ async fn ai_model_face_hair_fix_internal(
.ok_or("模板未找到")?;
// 上传图片
println!("📤 上传图片: {}", input_image);
let upload_response = client.upload_image(input_image, false).await
.map_err(|e| format!("上传图片失败: {}", e))?;
println!("✅ 图片上传成功,服务器文件名: {}", upload_response.name);
// 设置进度回调
let callbacks = Arc::new(

View File

@@ -56,7 +56,7 @@ const AiModelFaceHairFixTool: React.FC = () => {
// 通用设置
const [serverUrl, setServerUrl] = useState<string>('http://192.168.0.193:8188');
const [facePrompt, setFacePrompt] = useState<string>('beautiful woman, perfect skin, detailed facial features');
const [faceDenoiseStr, setFaceDenoiseStr] = useState<string>('0.25');
const [faceDenoiseStr, setFaceDenoiseStr] = useState<string>('0.35');
// 处理状态
const [isProcessing, setIsProcessing] = useState<boolean>(false);

View File

@@ -178,6 +178,9 @@ impl ComfyUIClient {
prompt_id: &str,
timeout: Option<Duration>,
) -> Result<HashMap<String, serde_json::Value>> {
// Clear any previous errors before starting
self.ws_client.clear_last_error().await;
let timeout_duration = timeout.unwrap_or(Duration::from_secs(300)); // 5 minutes default
let start_time = Instant::now();
let check_interval = Duration::from_millis(1000); // Check every second
@@ -190,6 +193,13 @@ impl ComfyUIClient {
)));
}
// Check for WebSocket execution errors first
if let Some(error) = self.ws_client.get_last_error().await {
return Err(ComfyUIError::execution(format!(
"Execution failed: {}", error.message
)));
}
// Check history for completion
match self.http_client.get_history_by_prompt(prompt_id).await {
Ok(history) => {

View File

@@ -167,8 +167,15 @@ impl WebSocketClient {
/// Handles incoming WebSocket messages
async fn handle_message(text: &str, event_emitter: &EventEmitter) -> Result<()> {
// 打印原始消息用于调试
log::debug!("Received WebSocket message: {}", text);
let message: WSMessage = serde_json::from_str(text)
.map_err(|e| ComfyUIError::new(format!("Failed to parse WebSocket message: {e}")))?;
.map_err(|e| {
log::error!("Failed to parse WebSocket message: {}", e);
log::error!("Raw message content: {}", text);
ComfyUIError::new(format!("Failed to parse WebSocket message: {e}"))
})?;
match message {
WSMessage::Progress { data } => {
@@ -195,12 +202,40 @@ impl WebSocketClient {
event_emitter.emit_executed(result).await;
}
WSMessage::ExecutionError { data } => {
// 构建错误消息,尝试多个可能的字段
let error_message = data.message
.or(data.error)
.or(data.exception_message)
.or_else(|| {
// 尝试从extra字段中提取错误信息
data.extra.get("exception_message")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
})
.or_else(|| {
data.extra.get("error")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
})
.unwrap_or_else(|| "Unknown execution error".to_string());
// 获取节点ID
let node_id = data.node_id
.or(data.node)
.or_else(|| {
data.extra.get("node")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
});
let error = ExecutionError {
node_id: data.node_id,
message: data.message,
node_id,
message: error_message,
details: data.details,
timestamp: Utc::now(),
};
log::error!("ComfyUI execution error: {}", error.message);
event_emitter.emit_error(error).await;
}
WSMessage::Unknown => {
@@ -215,6 +250,16 @@ impl WebSocketClient {
pub fn event_emitter(&self) -> &EventEmitter {
&self.event_emitter
}
/// Gets the last execution error and clears it
pub async fn get_last_error(&self) -> Option<ExecutionError> {
self.event_emitter.get_last_error().await
}
/// Clears the last error
pub async fn clear_last_error(&self) {
self.event_emitter.clear_last_error().await
}
}
impl Drop for WebSocketClient {

View File

@@ -6,18 +6,59 @@ use std::collections::HashMap;
/// Queue status response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueStatus {
pub queue_running: Vec<QueueItem>,
pub queue_pending: Vec<QueueItem>,
pub queue_running: Vec<serde_json::Value>,
pub queue_pending: Vec<serde_json::Value>,
}
/// Queue item
#[derive(Debug, Clone, Serialize, Deserialize)]
/// Queue item (parsed from array format)
#[derive(Debug, Clone)]
pub struct QueueItem {
pub prompt_id: String,
pub number: u32,
pub outputs: HashMap<String, serde_json::Value>,
}
impl QueueStatus {
/// Parse running queue items from raw JSON values
pub fn get_running_items(&self) -> Vec<QueueItem> {
self.queue_running.iter()
.filter_map(|item| Self::parse_queue_item(item))
.collect()
}
/// Parse pending queue items from raw JSON values
pub fn get_pending_items(&self) -> Vec<QueueItem> {
self.queue_pending.iter()
.filter_map(|item| Self::parse_queue_item(item))
.collect()
}
/// Parse a single queue item from JSON value
fn parse_queue_item(value: &serde_json::Value) -> Option<QueueItem> {
if let serde_json::Value::Array(arr) = value {
if arr.len() >= 2 {
let number = match &arr[0] {
serde_json::Value::Number(n) => n.as_u64()? as u32,
_ => return None,
};
let prompt_id = match &arr[1] {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Number(n) => n.to_string(),
_ => return None,
};
return Some(QueueItem {
prompt_id,
number,
outputs: HashMap::new(), // ComfyUI queue format doesn't include outputs
});
}
}
None
}
}
/// Prompt request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptRequest {
@@ -137,3 +178,59 @@ pub struct ViewMetadata {
#[serde(flatten)]
pub metadata: HashMap<String, serde_json::Value>,
}
/// Custom deserializer that can handle both string and number types
fn deserialize_string_or_number<'de, D>(deserializer: D) -> Result<String, D::Error>
where
D: serde::Deserializer<'de>,
{
use serde::de::{self, Visitor};
use std::fmt;
struct StringOrNumberVisitor;
impl<'de> Visitor<'de> for StringOrNumberVisitor {
type Value = String;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a string or number")
}
fn visit_str<E>(self, value: &str) -> Result<String, E>
where
E: de::Error,
{
Ok(value.to_string())
}
fn visit_string<E>(self, value: String) -> Result<String, E>
where
E: de::Error,
{
Ok(value)
}
fn visit_i64<E>(self, value: i64) -> Result<String, E>
where
E: de::Error,
{
Ok(value.to_string())
}
fn visit_u64<E>(self, value: u64) -> Result<String, E>
where
E: de::Error,
{
Ok(value.to_string())
}
fn visit_f64<E>(self, value: f64) -> Result<String, E>
where
E: de::Error,
{
Ok(value.to_string())
}
}
deserializer.deserialize_any(StringOrNumberVisitor)
}

View File

@@ -97,8 +97,17 @@ pub struct WSExecutedData {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSErrorData {
pub node_id: Option<String>,
pub message: String,
pub node: Option<String>,
pub message: Option<String>,
pub details: Option<serde_json::Value>,
pub error: Option<String>,
pub exception_message: Option<String>,
pub exception_type: Option<String>,
pub traceback: Option<Vec<String>>,
pub prompt_id: Option<String>,
// 使用 flatten 来捕获任何其他字段
#[serde(flatten)]
pub extra: HashMap<String, serde_json::Value>,
}
/// Execution options

View File

@@ -9,6 +9,7 @@ use crate::types::{ExecutionProgress, ExecutionResult, ExecutionError, Execution
#[derive(Clone)]
pub struct EventEmitter {
callbacks: Arc<RwLock<HashMap<String, Arc<dyn ExecutionCallbacks>>>>,
last_error: Arc<RwLock<Option<ExecutionError>>>,
}
impl EventEmitter {
@@ -16,6 +17,7 @@ impl EventEmitter {
pub fn new() -> Self {
Self {
callbacks: Arc::new(RwLock::new(HashMap::new())),
last_error: Arc::new(RwLock::new(None)),
}
}
@@ -57,6 +59,12 @@ impl EventEmitter {
/// Emits an error event
pub async fn emit_error(&self, error: ExecutionError) {
// Store the error for later retrieval
{
let mut last_error_guard = self.last_error.write().await;
*last_error_guard = Some(error.clone());
}
let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() {
callback.on_error(error.clone());
@@ -74,6 +82,18 @@ impl EventEmitter {
let mut callbacks_map = self.callbacks.write().await;
callbacks_map.clear();
}
/// Gets the last execution error and clears it
pub async fn get_last_error(&self) -> Option<ExecutionError> {
let mut error_guard = self.last_error.write().await;
error_guard.take()
}
/// Clears the last error
pub async fn clear_last_error(&self) {
let mut error_guard = self.last_error.write().await;
*error_guard = None;
}
}
impl Default for EventEmitter {