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 { impl QueueModel {
/// 从 SDK 队列状态创建模型 /// 从 SDK 队列状态创建模型
pub fn from_queue_status(queue_status: &QueueStatus) -> Self { 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() .iter()
.map(|item| QueueTaskInfo { .map(|item| QueueTaskInfo {
prompt_id: item.prompt_id.clone(), prompt_id: item.prompt_id.clone(),
@@ -373,7 +374,8 @@ impl QueueModel {
}) })
.collect(); .collect();
let pending_tasks: Vec<QueueTaskInfo> = queue_status.queue_pending let pending_tasks: Vec<QueueTaskInfo> = queue_status
.get_pending_items()
.iter() .iter()
.map(|item| QueueTaskInfo { .map(|item| QueueTaskInfo {
prompt_id: item.prompt_id.clone(), 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 { 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 { QueueStatusResponse {
running_count: status.queue_running.len() as u32, running_count: running_items.len() as u32,
pending_count: status.queue_pending.len() as u32, pending_count: pending_items.len() as u32,
running_tasks: status.queue_running.into_iter().map(|item| QueueTaskResponse { running_tasks: running_items.into_iter().map(|item| QueueTaskResponse {
prompt_id: item.prompt_id, prompt_id: item.prompt_id,
number: item.number, number: item.number,
status: "running".to_string(), status: "running".to_string(),
}).collect(), }).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, prompt_id: item.prompt_id,
number: item.number, number: item.number,
status: "pending".to_string(), 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 { match ai_model_face_hair_fix_internal(input_path, &server_url, &face_prompt, &face_denoise, &output_path_str, &window).await {
Ok(_) => { Ok(_) => {
let execution_time = file_start_time.elapsed().as_secs_f64(); let execution_time = file_start_time.elapsed().as_secs_f64();
println!("✅ 图片处理成功: {} -> {}", input_path, output_path_str);
results.push(ImageEnhancementResult { results.push(ImageEnhancementResult {
success: true, success: true,
message: "图片增强成功".to_string(), message: "图片增强成功".to_string(),
@@ -429,6 +430,8 @@ pub async fn ai_model_face_hair_fix_batch_images(
} }
Err(e) => { Err(e) => {
let execution_time = file_start_time.elapsed().as_secs_f64(); let execution_time = file_start_time.elapsed().as_secs_f64();
println!("❌ 图片处理失败: {} - 错误: {}", input_path, e);
eprintln!("详细错误信息: {}", e);
results.push(ImageEnhancementResult { results.push(ImageEnhancementResult {
success: false, success: false,
message: format!("图片增强失败: {}", e), 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(); 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 { let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress {
current: total_files, current: total_files,
total: total_files, total: total_files,
@@ -452,7 +462,7 @@ pub async fn ai_model_face_hair_fix_batch_images(
current_file: None, current_file: None,
}); });
Ok(BatchImageEnhancementResult { let batch_result = BatchImageEnhancementResult {
success: failed_count == 0, success: failed_count == 0,
message: format!("批量处理完成: 总计 {}, 成功 {}, 失败 {}", total_files, successful_count, failed_count), message: format!("批量处理完成: 总计 {}, 成功 {}, 失败 {}", total_files, successful_count, failed_count),
total_processed: total_files, total_processed: total_files,
@@ -460,7 +470,10 @@ pub async fn ai_model_face_hair_fix_batch_images(
failed_count, failed_count,
results, results,
total_execution_time, total_execution_time,
}) };
println!("🎯 返回批量处理结果: {:?}", batch_result);
Ok(batch_result)
} }
/// AI模型面部头发修复内部处理函数 /// AI模型面部头发修复内部处理函数
@@ -472,6 +485,9 @@ async fn ai_model_face_hair_fix_internal(
output_image: &str, output_image: &str,
window: &tauri::Window, window: &tauri::Window,
) -> Result<(), String> { ) -> Result<(), String> {
println!("🔄 开始处理图片: {}", input_image);
println!("📋 处理参数: server_url={}, face_prompt={}, face_denoise={}", server_url, face_prompt, face_denoise);
// 验证输入图片文件 // 验证输入图片文件
if !Path::new(input_image).exists() { if !Path::new(input_image).exists() {
return Err(format!("输入图片文件不存在: {}", input_image)); return Err(format!("输入图片文件不存在: {}", input_image));
@@ -481,16 +497,19 @@ async fn ai_model_face_hair_fix_internal(
validate_image_file(input_image)?; validate_image_file(input_image)?;
// 初始化 ComfyUI 客户端 // 初始化 ComfyUI 客户端
println!("🔌 初始化ComfyUI客户端...");
let mut client = ComfyUIClient::new(ComfyUIClientConfig { let mut client = ComfyUIClient::new(ComfyUIClientConfig {
base_url: server_url.to_string(), base_url: server_url.to_string(),
..Default::default() ..Default::default()
}).map_err(|e| format!("创建ComfyUI客户端失败: {}", e))?; }).map_err(|e| format!("创建ComfyUI客户端失败: {}", e))?;
// 连接到服务器 // 连接到服务器
println!("🌐 连接到ComfyUI服务器: {}", server_url);
client.connect().await client.connect().await
.map_err(|e| format!("连接ComfyUI服务器失败: {}", e))?; .map_err(|e| format!("连接ComfyUI服务器失败: {}", e))?;
// 注册模板 // 注册模板
println!("📝 注册AI模型面部头发修复模板...");
client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone()) client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())
.map_err(|e| format!("注册模板失败: {}", e))?; .map_err(|e| format!("注册模板失败: {}", e))?;
@@ -499,8 +518,10 @@ async fn ai_model_face_hair_fix_internal(
.ok_or("模板未找到")?; .ok_or("模板未找到")?;
// 上传图片 // 上传图片
println!("📤 上传图片: {}", input_image);
let upload_response = client.upload_image(input_image, false).await let upload_response = client.upload_image(input_image, false).await
.map_err(|e| format!("上传图片失败: {}", e))?; .map_err(|e| format!("上传图片失败: {}", e))?;
println!("✅ 图片上传成功,服务器文件名: {}", upload_response.name);
// 设置进度回调 // 设置进度回调
let callbacks = Arc::new( 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 [serverUrl, setServerUrl] = useState<string>('http://192.168.0.193:8188');
const [facePrompt, setFacePrompt] = useState<string>('beautiful woman, perfect skin, detailed facial features'); 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); const [isProcessing, setIsProcessing] = useState<boolean>(false);

View File

@@ -178,6 +178,9 @@ impl ComfyUIClient {
prompt_id: &str, prompt_id: &str,
timeout: Option<Duration>, timeout: Option<Duration>,
) -> Result<HashMap<String, serde_json::Value>> { ) -> 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 timeout_duration = timeout.unwrap_or(Duration::from_secs(300)); // 5 minutes default
let start_time = Instant::now(); let start_time = Instant::now();
let check_interval = Duration::from_millis(1000); // Check every second 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 // Check history for completion
match self.http_client.get_history_by_prompt(prompt_id).await { match self.http_client.get_history_by_prompt(prompt_id).await {
Ok(history) => { Ok(history) => {

View File

@@ -167,8 +167,15 @@ impl WebSocketClient {
/// Handles incoming WebSocket messages /// Handles incoming WebSocket messages
async fn handle_message(text: &str, event_emitter: &EventEmitter) -> Result<()> { async fn handle_message(text: &str, event_emitter: &EventEmitter) -> Result<()> {
// 打印原始消息用于调试
log::debug!("Received WebSocket message: {}", text);
let message: WSMessage = serde_json::from_str(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 { match message {
WSMessage::Progress { data } => { WSMessage::Progress { data } => {
@@ -195,12 +202,40 @@ impl WebSocketClient {
event_emitter.emit_executed(result).await; event_emitter.emit_executed(result).await;
} }
WSMessage::ExecutionError { data } => { 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 { let error = ExecutionError {
node_id: data.node_id, node_id,
message: data.message, message: error_message,
details: data.details, details: data.details,
timestamp: Utc::now(), timestamp: Utc::now(),
}; };
log::error!("ComfyUI execution error: {}", error.message);
event_emitter.emit_error(error).await; event_emitter.emit_error(error).await;
} }
WSMessage::Unknown => { WSMessage::Unknown => {
@@ -215,6 +250,16 @@ impl WebSocketClient {
pub fn event_emitter(&self) -> &EventEmitter { pub fn event_emitter(&self) -> &EventEmitter {
&self.event_emitter &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 { impl Drop for WebSocketClient {

View File

@@ -6,18 +6,59 @@ use std::collections::HashMap;
/// Queue status response /// Queue status response
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueStatus { pub struct QueueStatus {
pub queue_running: Vec<QueueItem>, pub queue_running: Vec<serde_json::Value>,
pub queue_pending: Vec<QueueItem>, pub queue_pending: Vec<serde_json::Value>,
} }
/// Queue item /// Queue item (parsed from array format)
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone)]
pub struct QueueItem { pub struct QueueItem {
pub prompt_id: String, pub prompt_id: String,
pub number: u32, pub number: u32,
pub outputs: HashMap<String, serde_json::Value>, 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 /// Prompt request
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptRequest { pub struct PromptRequest {
@@ -137,3 +178,59 @@ pub struct ViewMetadata {
#[serde(flatten)] #[serde(flatten)]
pub metadata: HashMap<String, serde_json::Value>, 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)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSErrorData { pub struct WSErrorData {
pub node_id: Option<String>, pub node_id: Option<String>,
pub message: String, pub node: Option<String>,
pub message: Option<String>,
pub details: Option<serde_json::Value>, 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 /// Execution options

View File

@@ -9,6 +9,7 @@ use crate::types::{ExecutionProgress, ExecutionResult, ExecutionError, Execution
#[derive(Clone)] #[derive(Clone)]
pub struct EventEmitter { pub struct EventEmitter {
callbacks: Arc<RwLock<HashMap<String, Arc<dyn ExecutionCallbacks>>>>, callbacks: Arc<RwLock<HashMap<String, Arc<dyn ExecutionCallbacks>>>>,
last_error: Arc<RwLock<Option<ExecutionError>>>,
} }
impl EventEmitter { impl EventEmitter {
@@ -16,6 +17,7 @@ impl EventEmitter {
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
callbacks: Arc::new(RwLock::new(HashMap::new())), callbacks: Arc::new(RwLock::new(HashMap::new())),
last_error: Arc::new(RwLock::new(None)),
} }
} }
@@ -57,6 +59,12 @@ impl EventEmitter {
/// Emits an error event /// Emits an error event
pub async fn emit_error(&self, error: ExecutionError) { 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; let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() { for callback in callbacks_map.values() {
callback.on_error(error.clone()); callback.on_error(error.clone());
@@ -74,6 +82,18 @@ impl EventEmitter {
let mut callbacks_map = self.callbacks.write().await; let mut callbacks_map = self.callbacks.write().await;
callbacks_map.clear(); 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 { impl Default for EventEmitter {