fix: 修复批量任务状态更新错乱问题
- 添加 comfyui_prompt_id 字段到 OutfitImageRecord 和 OutfitPhotoGeneration 模型 - 创建数据库迁移添加 prompt_id 字段和索引 - 修改 ComfyUI 服务添加基于 prompt_id 的工作流执行方法 - 更新仓储层添加根据 prompt_id 查找记录的方法 - 修改任务执行逻辑保存 ComfyUI prompt_id 到数据库 - 创建基于 prompt_id 的进度回调函数替代基于索引的回调 - 修改批量处理逻辑使用精确的 ID 匹配而非数组索引 - 标记旧的基于索引的批量方法为已弃用 解决了当选择2个商品1个模特时,任务完成后更新错误任务状态的问题。 现在通过 ComfyUI 的 prompt_id 建立精确的任务映射关系。
This commit is contained in:
@@ -611,6 +611,32 @@ impl ComfyUIService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 通过 WebSocket 跟踪工作流执行进度(带 prompt_id 的回调)
|
||||
pub async fn track_workflow_progress_with_prompt_id<F>(
|
||||
&self,
|
||||
prompt_id: &str,
|
||||
mut progress_callback: F,
|
||||
) -> Result<()>
|
||||
where
|
||||
F: FnMut(String, WorkflowProgress) + Send + 'static,
|
||||
{
|
||||
let prompt_id_clone = prompt_id.to_string();
|
||||
let mut wrapped_callback = move |progress: WorkflowProgress| {
|
||||
progress_callback(prompt_id_clone.clone(), progress);
|
||||
};
|
||||
|
||||
let result = timeout(
|
||||
self.execution_config.websocket_timeout,
|
||||
self.track_workflow_progress_internal(prompt_id, &mut wrapped_callback)
|
||||
).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(())) => Ok(()),
|
||||
Ok(Err(e)) => Err(e),
|
||||
Err(_) => Err(anyhow!("工作流执行超时")),
|
||||
}
|
||||
}
|
||||
|
||||
/// 内部进度跟踪实现
|
||||
async fn track_workflow_progress_internal<F>(
|
||||
&self,
|
||||
@@ -1187,6 +1213,179 @@ impl ComfyUIService {
|
||||
Ok(upload_results)
|
||||
}
|
||||
|
||||
/// 执行工作流并自动上传结果到云端(带商品编号),返回 prompt_id 和上传结果
|
||||
pub async fn execute_workflow_with_upload_indexed_with_prompt_id<F>(
|
||||
&self,
|
||||
cloud_service: &CloudUploadService,
|
||||
workflow_path: &str,
|
||||
model_image_url: &str,
|
||||
product_image_url: &str,
|
||||
prompt: &str,
|
||||
negative_prompt: Option<&str>,
|
||||
remote_key_prefix: Option<&str>,
|
||||
progress_callback: F,
|
||||
product_index: Option<usize>,
|
||||
) -> Result<(String, Vec<UploadResult>)>
|
||||
where
|
||||
F: FnMut(WorkflowProgress) + Send + 'static,
|
||||
{
|
||||
info!("开始执行工作流并上传结果: {} (商品编号: {:?})", workflow_path, product_index.map(|i| i + 1));
|
||||
let start_time = Instant::now();
|
||||
|
||||
// 1. 验证服务器连接
|
||||
if !self.check_connection().await? {
|
||||
return Err(anyhow!("ComfyUI 服务器不可用"));
|
||||
}
|
||||
|
||||
// 2. 加载工作流
|
||||
let workflow = self.load_workflow(workflow_path).await
|
||||
.map_err(|e| anyhow!("加载工作流失败: {}", e))?;
|
||||
|
||||
// 3. 替换节点值
|
||||
let updated_workflow = self.auto_replace_input_nodes(
|
||||
workflow,
|
||||
model_image_url,
|
||||
product_image_url,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
).map_err(|e| anyhow!("替换工作流节点失败: {}", e))?;
|
||||
|
||||
// 4. 提交工作流并获取 prompt_id
|
||||
let prompt_id = self.submit_workflow(updated_workflow).await
|
||||
.map_err(|e| anyhow!("提交工作流失败: {}", e))?;
|
||||
|
||||
info!("工作流提交成功,提示 ID: {}", prompt_id);
|
||||
|
||||
// 5. 跟踪进度
|
||||
self.track_workflow_progress(&prompt_id, progress_callback).await
|
||||
.map_err(|e| anyhow!("跟踪工作流进度失败: {}", e))?;
|
||||
|
||||
// 6. 验证工作流完成状态
|
||||
let mut attempts = 0;
|
||||
let max_attempts = 10;
|
||||
|
||||
while attempts < max_attempts {
|
||||
if self.is_workflow_completed(&prompt_id).await? {
|
||||
break;
|
||||
}
|
||||
|
||||
attempts += 1;
|
||||
if attempts >= max_attempts {
|
||||
return Err(anyhow!("工作流执行超时,已等待 {} 次检查", max_attempts));
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
// 7. 获取结果
|
||||
let image_filenames = self.get_history(&prompt_id).await
|
||||
.map_err(|e| anyhow!("获取工作流结果失败: {}", e))?;
|
||||
|
||||
if image_filenames.is_empty() {
|
||||
return Err(anyhow!("工作流执行完成但未生成任何图片"));
|
||||
}
|
||||
|
||||
info!("工作流执行完成,开始上传 {} 张图片", image_filenames.len());
|
||||
|
||||
// 8. 下载并上传图片
|
||||
let upload_results = self.download_and_upload_images(
|
||||
cloud_service,
|
||||
&image_filenames,
|
||||
remote_key_prefix,
|
||||
).await?;
|
||||
|
||||
let successful_uploads = upload_results.iter().filter(|r| r.success).count();
|
||||
info!("图片上传完成: {}/{} 成功", successful_uploads, upload_results.len());
|
||||
|
||||
Ok((prompt_id, upload_results))
|
||||
}
|
||||
|
||||
/// 执行工作流并自动上传结果到云端(带 prompt_id 的进度回调),返回 prompt_id 和上传结果
|
||||
pub async fn execute_workflow_with_upload_indexed_with_prompt_id_callback<F>(
|
||||
&self,
|
||||
cloud_service: &CloudUploadService,
|
||||
workflow_path: &str,
|
||||
model_image_url: &str,
|
||||
product_image_url: &str,
|
||||
prompt: &str,
|
||||
negative_prompt: Option<&str>,
|
||||
remote_key_prefix: Option<&str>,
|
||||
progress_callback: F,
|
||||
product_index: Option<usize>,
|
||||
) -> Result<(String, Vec<UploadResult>)>
|
||||
where
|
||||
F: FnMut(String, WorkflowProgress) + Send + 'static,
|
||||
{
|
||||
info!("开始执行工作流并上传结果: {} (商品编号: {:?})", workflow_path, product_index.map(|i| i + 1));
|
||||
|
||||
// 1. 验证服务器连接
|
||||
if !self.check_connection().await? {
|
||||
return Err(anyhow!("ComfyUI 服务器不可用"));
|
||||
}
|
||||
|
||||
// 2. 加载工作流
|
||||
let workflow = self.load_workflow(workflow_path).await
|
||||
.map_err(|e| anyhow!("加载工作流失败: {}", e))?;
|
||||
|
||||
// 3. 替换节点值
|
||||
let updated_workflow = self.auto_replace_input_nodes(
|
||||
workflow,
|
||||
model_image_url,
|
||||
product_image_url,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
).map_err(|e| anyhow!("替换工作流节点失败: {}", e))?;
|
||||
|
||||
// 4. 提交工作流并获取 prompt_id
|
||||
let prompt_id = self.submit_workflow(updated_workflow).await
|
||||
.map_err(|e| anyhow!("提交工作流失败: {}", e))?;
|
||||
|
||||
info!("工作流提交成功,提示 ID: {}", prompt_id);
|
||||
|
||||
// 5. 跟踪进度(使用带 prompt_id 的回调)
|
||||
self.track_workflow_progress_with_prompt_id(&prompt_id, progress_callback).await
|
||||
.map_err(|e| anyhow!("跟踪工作流进度失败: {}", e))?;
|
||||
|
||||
// 6. 验证工作流完成状态
|
||||
let mut attempts = 0;
|
||||
let max_attempts = 10;
|
||||
|
||||
while attempts < max_attempts {
|
||||
if self.is_workflow_completed(&prompt_id).await? {
|
||||
break;
|
||||
}
|
||||
|
||||
attempts += 1;
|
||||
if attempts >= max_attempts {
|
||||
return Err(anyhow!("工作流执行超时,已等待 {} 次检查", max_attempts));
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
// 7. 获取结果
|
||||
let image_filenames = self.get_history(&prompt_id).await
|
||||
.map_err(|e| anyhow!("获取工作流结果失败: {}", e))?;
|
||||
|
||||
if image_filenames.is_empty() {
|
||||
return Err(anyhow!("工作流执行完成但未生成任何图片"));
|
||||
}
|
||||
|
||||
info!("工作流执行完成,开始上传 {} 张图片", image_filenames.len());
|
||||
|
||||
// 8. 下载并上传图片
|
||||
let upload_results = self.download_and_upload_images(
|
||||
cloud_service,
|
||||
&image_filenames,
|
||||
remote_key_prefix,
|
||||
).await?;
|
||||
|
||||
let successful_uploads = upload_results.iter().filter(|r| r.success).count();
|
||||
info!("图片上传完成: {}/{} 成功", successful_uploads, upload_results.len());
|
||||
|
||||
Ok((prompt_id, upload_results))
|
||||
}
|
||||
|
||||
/// 保存调试用的工作流文件
|
||||
async fn save_debug_workflow(&self, workflow: &Value, original_path: &str, product_index: Option<usize>) -> Result<()> {
|
||||
use std::path::Path;
|
||||
|
||||
@@ -2,7 +2,8 @@ use anyhow::{Result, anyhow};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
use tracing::{info, warn, error};
|
||||
use tauri::{AppHandle, Emitter};
|
||||
|
||||
use crate::config::AppConfig;
|
||||
use crate::data::models::outfit_photo_generation::{
|
||||
@@ -124,7 +125,7 @@ impl OutfitPhotoGenerationService {
|
||||
// 执行 ComfyUI 工作流并自动上传结果
|
||||
let remote_key_prefix = format!("outfit-photos/{}/{}", generation.project_id, generation.id);
|
||||
|
||||
match comfyui_service.execute_workflow_with_upload(
|
||||
match comfyui_service.execute_workflow_with_upload_indexed_with_prompt_id(
|
||||
&self.cloud_upload_service,
|
||||
&workflow_config.workflow_file_path,
|
||||
&model_image_url,
|
||||
@@ -133,8 +134,13 @@ impl OutfitPhotoGenerationService {
|
||||
generation.negative_prompt.as_deref(),
|
||||
Some(&remote_key_prefix),
|
||||
progress_callback,
|
||||
None, // product_index
|
||||
).await {
|
||||
Ok(upload_results) => {
|
||||
Ok((prompt_id, upload_results)) => {
|
||||
// 保存 ComfyUI prompt_id
|
||||
generation.set_comfyui_prompt_id(prompt_id.clone());
|
||||
info!("✅ 已保存 ComfyUI prompt_id: {} 到生成记录: {}", prompt_id, generation.id);
|
||||
|
||||
// 处理上传结果
|
||||
let mut successful_uploads = 0;
|
||||
|
||||
@@ -185,6 +191,10 @@ impl OutfitPhotoGenerationService {
|
||||
}
|
||||
|
||||
/// 批量生成穿搭照片
|
||||
///
|
||||
/// ⚠️ 已弃用:此方法使用基于索引的进度回调,在并发场景下可能导致任务状态更新错乱。
|
||||
/// 请使用 `generate_batch_with_prompt_id_callback` 方法,它使用基于 prompt_id 的精确匹配。
|
||||
#[deprecated(note = "使用 generate_batch_with_prompt_id_callback 替代")]
|
||||
pub async fn generate_batch<F>(
|
||||
&self,
|
||||
requests: Vec<OutfitPhotoGenerationRequest>,
|
||||
@@ -243,6 +253,201 @@ impl OutfitPhotoGenerationService {
|
||||
Ok(responses)
|
||||
}
|
||||
|
||||
/// 批量生成穿搭照片(使用基于 prompt_id 的进度回调)
|
||||
pub async fn generate_batch_with_prompt_id_callback<F>(
|
||||
&self,
|
||||
requests: Vec<OutfitPhotoGenerationRequest>,
|
||||
progress_callback: F,
|
||||
) -> Result<Vec<OutfitPhotoGenerationResponse>>
|
||||
where
|
||||
F: Fn(String, WorkflowProgress) + Send + Sync + 'static,
|
||||
{
|
||||
let mut responses = Vec::new();
|
||||
let total_requests = requests.len();
|
||||
|
||||
info!("开始批量生成穿搭照片(基于 prompt_id),共 {} 个请求", total_requests);
|
||||
|
||||
// 将回调包装在 Arc 中以便在循环中共享
|
||||
let callback_arc = Arc::new(progress_callback);
|
||||
|
||||
for (index, request) in requests.into_iter().enumerate() {
|
||||
info!("处理批量请求 {}/{}", index + 1, total_requests);
|
||||
|
||||
// 为每个请求创建基于 prompt_id 的进度回调
|
||||
let callback_clone = Arc::clone(&callback_arc);
|
||||
let individual_callback = move |prompt_id: String, progress: WorkflowProgress| {
|
||||
callback_clone(prompt_id, progress);
|
||||
};
|
||||
|
||||
// 执行单个生成请求(使用新的方法)
|
||||
match self.generate_outfit_photo_with_prompt_id_callback(request, individual_callback).await {
|
||||
Ok(response) => {
|
||||
info!("批量请求 {} 完成", index + 1);
|
||||
responses.push(response);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("批量请求 {} 失败: {}", index + 1, e);
|
||||
// 创建失败响应
|
||||
responses.push(OutfitPhotoGenerationResponse {
|
||||
id: format!("batch_error_{}", index),
|
||||
status: GenerationStatus::Failed,
|
||||
result_image_urls: vec![],
|
||||
error_message: Some(e.to_string()),
|
||||
generation_time_ms: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 批量处理间隔,避免服务器过载
|
||||
if index < total_requests - 1 {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
|
||||
info!("批量生成完成(基于 prompt_id),成功: {}/{}",
|
||||
responses.iter().filter(|r| r.status == GenerationStatus::Completed).count(),
|
||||
total_requests
|
||||
);
|
||||
|
||||
Ok(responses)
|
||||
}
|
||||
|
||||
/// 生成穿搭照片(使用基于 prompt_id 的进度回调)
|
||||
pub async fn generate_outfit_photo_with_prompt_id_callback<F>(
|
||||
&self,
|
||||
request: OutfitPhotoGenerationRequest,
|
||||
progress_callback: F,
|
||||
) -> Result<OutfitPhotoGenerationResponse>
|
||||
where
|
||||
F: FnMut(String, WorkflowProgress) + Send + 'static,
|
||||
{
|
||||
// 创建生成任务
|
||||
let generation = self.create_generation_task(request).await?;
|
||||
|
||||
// 执行生成任务(使用基于 prompt_id 的回调)
|
||||
self.execute_generation_with_prompt_id_callback(&generation.id, progress_callback).await
|
||||
}
|
||||
|
||||
/// 执行穿搭照片生成(使用基于 prompt_id 的进度回调)
|
||||
pub async fn execute_generation_with_prompt_id_callback<F>(
|
||||
&self,
|
||||
generation_id: &str,
|
||||
progress_callback: F,
|
||||
) -> Result<OutfitPhotoGenerationResponse>
|
||||
where
|
||||
F: FnMut(String, WorkflowProgress) + Send + 'static,
|
||||
{
|
||||
let mut generation = self.get_generation_record(generation_id).await?;
|
||||
|
||||
// 检查 ComfyUI 是否启用
|
||||
let config = self.config.lock().await;
|
||||
if !config.is_comfyui_enabled() {
|
||||
return Err(anyhow!("ComfyUI 功能未启用"));
|
||||
}
|
||||
|
||||
let comfyui_service = ComfyUIService::new(config.get_comfyui_settings().clone());
|
||||
drop(config);
|
||||
|
||||
// 更新状态为处理中
|
||||
generation.update_status(GenerationStatus::Processing);
|
||||
self.update_generation_record(&generation).await?;
|
||||
|
||||
// 检查 ComfyUI 连接
|
||||
if !comfyui_service.check_connection().await? {
|
||||
generation.set_error("ComfyUI 服务器连接失败".to_string());
|
||||
self.update_generation_record(&generation).await?;
|
||||
return Ok(self.create_response(&generation));
|
||||
}
|
||||
|
||||
// 上传商品图片到云端
|
||||
let product_image_url = if generation.product_image_path.starts_with("https://") {
|
||||
generation.product_image_path.clone()
|
||||
} else {
|
||||
match self.cloud_upload_service.upload_file(&generation.product_image_path, None, None).await {
|
||||
Ok(upload_result) => {
|
||||
if upload_result.success {
|
||||
if let Some(url) = upload_result.remote_url {
|
||||
generation.product_image_url = Some(url.clone());
|
||||
self.update_generation_record(&generation).await?;
|
||||
url
|
||||
} else {
|
||||
generation.set_error("商品图片上传失败:未获得远程URL".to_string());
|
||||
self.update_generation_record(&generation).await?;
|
||||
return Ok(self.create_response(&generation));
|
||||
}
|
||||
} else {
|
||||
generation.set_error(format!("商品图片上传失败: {}",
|
||||
upload_result.error_message.unwrap_or_default()));
|
||||
self.update_generation_record(&generation).await?;
|
||||
return Ok(self.create_response(&generation));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
generation.set_error(format!("商品图片上传失败: {}", e));
|
||||
self.update_generation_record(&generation).await?;
|
||||
return Ok(self.create_response(&generation));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 获取模特图片URL
|
||||
let model_image_url = self.get_model_image_url(&generation.model_id).await?;
|
||||
|
||||
// 获取工作流配置
|
||||
let workflow_config = self.get_default_workflow_config().await?;
|
||||
|
||||
// 执行 ComfyUI 工作流并自动上传结果(使用基于 prompt_id 的回调)
|
||||
let remote_key_prefix = format!("outfit-photos/{}/{}", generation.project_id, generation.id);
|
||||
|
||||
match comfyui_service.execute_workflow_with_upload_indexed_with_prompt_id_callback(
|
||||
&self.cloud_upload_service,
|
||||
&workflow_config.workflow_file_path,
|
||||
&model_image_url,
|
||||
&product_image_url,
|
||||
&generation.prompt,
|
||||
generation.negative_prompt.as_deref(),
|
||||
Some(&remote_key_prefix),
|
||||
progress_callback,
|
||||
None, // product_index
|
||||
).await {
|
||||
Ok((prompt_id, upload_results)) => {
|
||||
// 保存 ComfyUI prompt_id
|
||||
generation.set_comfyui_prompt_id(prompt_id.clone());
|
||||
info!("✅ 已保存 ComfyUI prompt_id: {} 到生成记录: {}", prompt_id, generation.id);
|
||||
|
||||
// 处理上传结果
|
||||
let mut successful_uploads = 0;
|
||||
|
||||
for upload_result in upload_results {
|
||||
if upload_result.success {
|
||||
if let Some(remote_url) = upload_result.remote_url {
|
||||
generation.result_image_urls.push(remote_url.clone());
|
||||
successful_uploads += 1;
|
||||
info!("图片上传成功: {}", remote_url);
|
||||
}
|
||||
} else {
|
||||
warn!("图片上传失败: {}", upload_result.error_message.unwrap_or_default());
|
||||
}
|
||||
}
|
||||
|
||||
if successful_uploads == 0 {
|
||||
generation.set_error("所有图片上传失败".to_string());
|
||||
} else {
|
||||
generation.update_status(GenerationStatus::Completed);
|
||||
info!("成功上传 {} 张图片", successful_uploads);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
generation.set_error(format!("ComfyUI 工作流执行或上传失败: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// 更新数据库记录
|
||||
self.update_generation_record(&generation).await?;
|
||||
|
||||
Ok(self.create_response(&generation))
|
||||
}
|
||||
|
||||
/// 重新上传失败的图片
|
||||
pub async fn retry_failed_uploads(&self, generation_id: &str) -> Result<OutfitPhotoGenerationResponse> {
|
||||
// 从数据库获取生成记录
|
||||
@@ -455,3 +660,53 @@ impl OutfitPhotoGenerationService {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建基于 prompt_id 的进度回调函数(用于 OutfitPhotoGeneration)
|
||||
/// 这个函数用于批量处理场景,通过 prompt_id 精确匹配任务
|
||||
pub fn create_outfit_photo_prompt_id_based_progress_callback(
|
||||
database: Arc<Database>,
|
||||
app_handle: AppHandle,
|
||||
) -> impl Fn(String, WorkflowProgress) + Send + Sync + 'static {
|
||||
move |prompt_id: String, progress: WorkflowProgress| {
|
||||
let repo = match OutfitPhotoGenerationRepository::new(database.clone()) {
|
||||
Ok(repo) => repo,
|
||||
Err(e) => {
|
||||
error!("❌ 创建 OutfitPhotoGenerationRepository 失败: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 根据 prompt_id 查找对应的记录
|
||||
match repo.get_by_comfyui_prompt_id(&prompt_id) {
|
||||
Ok(Some(mut generation)) => {
|
||||
let generation_id = generation.id.clone();
|
||||
|
||||
// 更新数据库中的进度(OutfitPhotoGeneration 没有 progress 字段,所以我们只发送事件)
|
||||
info!("✅ 通过 prompt_id {} 更新生成记录 {} 进度: {}%",
|
||||
prompt_id, generation_id, progress.progress_percentage);
|
||||
|
||||
// 发送进度事件到前端
|
||||
let progress_event = serde_json::json!({
|
||||
"type": "outfit_photo_generation_progress",
|
||||
"generation_id": generation_id,
|
||||
"prompt_id": prompt_id,
|
||||
"progress_percentage": progress.progress_percentage,
|
||||
"current_step": progress.current_step,
|
||||
"total_steps": progress.total_steps,
|
||||
"status_message": progress.status_message,
|
||||
"current_node_id": progress.current_node_id
|
||||
});
|
||||
|
||||
if let Err(e) = app_handle.emit("outfit_photo_generation_progress", &progress_event) {
|
||||
warn!("发送进度事件失败: {}", e);
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
warn!("⚠️ 未找到 prompt_id {} 对应的生成记录", prompt_id);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("❌ 查找 prompt_id {} 对应生成记录失败: {}", prompt_id, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,7 +50,8 @@ pub struct OutfitImageRecord {
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub duration_ms: Option<u64>,
|
||||
|
||||
pub comfyui_prompt_id: Option<String>, // ComfyUI 任务ID,用于精确匹配任务状态更新
|
||||
|
||||
// 关联数据
|
||||
pub product_images: Vec<ProductImage>,
|
||||
pub outfit_images: Vec<OutfitImage>,
|
||||
@@ -72,6 +73,7 @@ impl OutfitImageRecord {
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
comfyui_prompt_id: None,
|
||||
product_images: Vec::new(),
|
||||
outfit_images: Vec::new(),
|
||||
}
|
||||
@@ -120,9 +122,15 @@ impl OutfitImageRecord {
|
||||
self.started_at = None;
|
||||
self.completed_at = None;
|
||||
self.duration_ms = None;
|
||||
self.comfyui_prompt_id = None; // 重置 ComfyUI 任务ID
|
||||
// 保留 result_urls,如果之前有部分成功的结果
|
||||
// self.result_urls.clear(); // 可选:是否清除之前的结果
|
||||
}
|
||||
|
||||
/// 设置 ComfyUI 任务ID
|
||||
pub fn set_comfyui_prompt_id(&mut self, prompt_id: String) {
|
||||
self.comfyui_prompt_id = Some(prompt_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// 商品图片
|
||||
|
||||
@@ -200,6 +200,12 @@ impl OutfitPhotoGeneration {
|
||||
self.update_status(GenerationStatus::Failed);
|
||||
}
|
||||
|
||||
/// 设置 ComfyUI prompt ID
|
||||
pub fn set_comfyui_prompt_id(&mut self, prompt_id: String) {
|
||||
self.comfyui_prompt_id = Some(prompt_id);
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
/// 添加生成结果
|
||||
pub fn add_result(&mut self, image_path: String, image_url: Option<String>) {
|
||||
self.result_image_paths.push(image_path);
|
||||
|
||||
@@ -48,6 +48,7 @@ impl OutfitImageRepository {
|
||||
started_at DATETIME,
|
||||
completed_at DATETIME,
|
||||
duration_ms INTEGER,
|
||||
comfyui_prompt_id TEXT,
|
||||
FOREIGN KEY (model_id) REFERENCES models (id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (model_image_id) REFERENCES model_photos (id) ON DELETE CASCADE
|
||||
)
|
||||
@@ -104,6 +105,11 @@ impl OutfitImageRepository {
|
||||
[],
|
||||
).map_err(|e| anyhow!("创建索引失败: {}", e))?;
|
||||
|
||||
pooled_conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_outfit_image_records_comfyui_prompt_id ON outfit_image_records (comfyui_prompt_id)",
|
||||
[],
|
||||
).map_err(|e| anyhow!("创建索引失败: {}", e))?;
|
||||
|
||||
pooled_conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_product_images_outfit_record_id ON product_images (outfit_record_id)",
|
||||
[],
|
||||
@@ -134,8 +140,8 @@ impl OutfitImageRepository {
|
||||
r#"
|
||||
INSERT INTO outfit_image_records (
|
||||
id, model_id, model_image_id, generation_prompt, status, progress,
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms, comfyui_prompt_id
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)
|
||||
"#,
|
||||
params![
|
||||
record.id,
|
||||
@@ -150,6 +156,7 @@ impl OutfitImageRepository {
|
||||
record.started_at.map(|dt| dt.to_rfc3339()),
|
||||
record.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
record.duration_ms,
|
||||
record.comfyui_prompt_id,
|
||||
],
|
||||
).map_err(|e| anyhow!("创建穿搭图片记录失败: {}", e))?;
|
||||
|
||||
@@ -372,7 +379,7 @@ impl OutfitImageRepository {
|
||||
let mut stmt = conn.prepare(
|
||||
r#"
|
||||
SELECT id, model_id, model_image_id, generation_prompt, status, progress,
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms, comfyui_prompt_id
|
||||
FROM outfit_image_records WHERE id = ?1
|
||||
"#,
|
||||
).map_err(|e| anyhow!("准备查询语句失败: {}", e))?;
|
||||
@@ -394,6 +401,41 @@ impl OutfitImageRepository {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// 根据 ComfyUI prompt_id 获取穿搭图片生成记录
|
||||
pub fn get_record_by_comfyui_prompt_id(&self, prompt_id: &str) -> Result<Option<OutfitImageRecord>> {
|
||||
if !self.database.has_pool() {
|
||||
return Err(anyhow!("连接池未启用,无法安全执行数据库操作"));
|
||||
}
|
||||
|
||||
let conn = self
|
||||
.database
|
||||
.acquire_from_pool()
|
||||
.map_err(|e| anyhow!("获取连接池连接失败: {}", e))?;
|
||||
let mut stmt = conn.prepare(
|
||||
r#"
|
||||
SELECT id, model_id, model_image_id, generation_prompt, status, progress,
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms, comfyui_prompt_id
|
||||
FROM outfit_image_records WHERE comfyui_prompt_id = ?1
|
||||
"#,
|
||||
).map_err(|e| anyhow!("准备查询语句失败: {}", e))?;
|
||||
|
||||
let record_iter = stmt.query_map(params![prompt_id], |row| {
|
||||
self.row_to_record(row)
|
||||
}).map_err(|e| anyhow!("查询穿搭图片记录失败: {}", e))?;
|
||||
|
||||
for record_result in record_iter {
|
||||
let mut record = record_result.map_err(|e| anyhow!("解析记录失败: {}", e))?;
|
||||
|
||||
// 加载关联的商品图片和穿搭图片
|
||||
record.product_images = self.get_product_images_by_record_id(&record.id)?;
|
||||
record.outfit_images = self.get_outfit_images_by_record_id(&record.id)?;
|
||||
|
||||
return Ok(Some(record));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// 根据模特ID获取穿搭图片生成记录列表(强制使用连接池)
|
||||
pub fn get_records_by_model_id(&self, model_id: &str) -> Result<Vec<OutfitImageRecord>> {
|
||||
|
||||
@@ -408,7 +450,7 @@ impl OutfitImageRepository {
|
||||
let mut stmt = pooled_conn.prepare(
|
||||
r#"
|
||||
SELECT id, model_id, model_image_id, generation_prompt, status, progress,
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms, comfyui_prompt_id
|
||||
FROM outfit_image_records
|
||||
WHERE model_id = ?1
|
||||
ORDER BY created_at DESC
|
||||
@@ -448,7 +490,7 @@ impl OutfitImageRepository {
|
||||
let mut stmt = pooled_conn.prepare(
|
||||
r#"
|
||||
SELECT id, model_id, model_image_id, generation_prompt, status, progress,
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms
|
||||
result_urls, error_message, created_at, started_at, completed_at, duration_ms, comfyui_prompt_id
|
||||
FROM outfit_image_records
|
||||
WHERE model_id = ?1
|
||||
ORDER BY created_at DESC
|
||||
@@ -750,6 +792,7 @@ impl OutfitImageRepository {
|
||||
started_at,
|
||||
completed_at,
|
||||
duration_ms: row.get(11)?,
|
||||
comfyui_prompt_id: row.get(12)?,
|
||||
product_images: Vec::new(), // 将在调用方加载
|
||||
outfit_images: Vec::new(), // 将在调用方加载
|
||||
})
|
||||
|
||||
@@ -85,6 +85,32 @@ impl OutfitPhotoGenerationRepository {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// 根据 ComfyUI prompt ID 获取穿搭照片生成记录
|
||||
pub fn get_by_comfyui_prompt_id(&self, prompt_id: &str) -> Result<Option<OutfitPhotoGeneration>> {
|
||||
// 优先使用只读连接,提高并发性能
|
||||
let conn_handle = self.database.get_best_read_connection()?;
|
||||
|
||||
let mut stmt = conn_handle.prepare(
|
||||
r#"
|
||||
SELECT id, project_id, model_id, product_image_path, product_image_url,
|
||||
prompt, negative_prompt, status, workflow_id, comfyui_prompt_id,
|
||||
result_image_paths, result_image_urls, error_message,
|
||||
started_at, completed_at, generation_time_ms, created_at, updated_at
|
||||
FROM outfit_photo_generations WHERE comfyui_prompt_id = ?1
|
||||
"#
|
||||
)?;
|
||||
|
||||
let generation_iter = stmt.query_map([prompt_id], |row| {
|
||||
self.row_to_generation(row)
|
||||
})?;
|
||||
|
||||
for generation in generation_iter {
|
||||
return Ok(Some(generation?));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// 根据项目ID获取穿搭照片生成记录
|
||||
pub fn get_by_project_id(&self, project_id: &str) -> Result<Vec<OutfitPhotoGeneration>> {
|
||||
// 优先使用只读连接,提高并发性能
|
||||
|
||||
@@ -248,7 +248,7 @@ impl MigrationManager {
|
||||
|
||||
// 迁移 26: 创建穿搭照片生成记录表
|
||||
self.add_migration(Migration {
|
||||
version: 28,
|
||||
version: 26,
|
||||
description: "创建穿搭照片生成记录表".to_string(),
|
||||
up_sql: include_str!("migrations/026_create_outfit_photo_generations_table.sql").to_string(),
|
||||
down_sql: Some(include_str!("migrations/026_create_outfit_photo_generations_table_down.sql").to_string()),
|
||||
@@ -261,6 +261,14 @@ impl MigrationManager {
|
||||
up_sql: include_str!("migrations/027_create_video_generation_records_table.sql").to_string(),
|
||||
down_sql: Some(include_str!("migrations/027_create_video_generation_records_table_down.sql").to_string()),
|
||||
});
|
||||
|
||||
// 迁移 28: 添加 ComfyUI prompt_id 字段到穿搭图片生成记录表
|
||||
self.add_migration(Migration {
|
||||
version: 28,
|
||||
description: "添加 ComfyUI prompt_id 字段到穿搭图片生成记录表".to_string(),
|
||||
up_sql: include_str!("migrations/028_add_comfyui_prompt_id_to_outfit_image_records.sql").to_string(),
|
||||
down_sql: None, // 暂时不提供回滚脚本
|
||||
});
|
||||
}
|
||||
|
||||
/// 添加迁移
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
-- 添加 ComfyUI prompt_id 字段到穿搭图片生成记录表
|
||||
-- 用于建立与 ComfyUI 任务的精确映射关系,解决批量任务状态更新错乱问题
|
||||
|
||||
ALTER TABLE outfit_image_records ADD COLUMN comfyui_prompt_id TEXT;
|
||||
|
||||
-- 创建索引以提高根据 prompt_id 查询的性能
|
||||
CREATE INDEX IF NOT EXISTS idx_outfit_image_records_comfyui_prompt_id ON outfit_image_records (comfyui_prompt_id);
|
||||
@@ -1067,7 +1067,7 @@ async fn perform_outfit_image_generation(
|
||||
}
|
||||
};
|
||||
|
||||
// 创建进度回调
|
||||
// 创建进度回调(基于 record_id 的直接匹配)
|
||||
let app_handle_clone = app_handle.clone();
|
||||
let record_id_clone = record_id.clone();
|
||||
let outfit_repo_clone = OutfitImageRepository::new(database.clone());
|
||||
@@ -1111,7 +1111,7 @@ async fn perform_outfit_image_generation(
|
||||
// 执行 ComfyUI 工作流并自动上传结果
|
||||
let remote_key_prefix = format!("outfit-images/{}", record_id);
|
||||
|
||||
match comfyui_service.execute_workflow_with_upload_indexed(
|
||||
match comfyui_service.execute_workflow_with_upload_indexed_with_prompt_id(
|
||||
&cloud_upload_service,
|
||||
&workflow_path_str,
|
||||
&model_image_url,
|
||||
@@ -1122,7 +1122,17 @@ async fn perform_outfit_image_generation(
|
||||
progress_callback,
|
||||
request.product_index, // 使用请求中的商品编号
|
||||
).await {
|
||||
Ok(upload_results) => {
|
||||
Ok((prompt_id, upload_results)) => {
|
||||
// 保存 ComfyUI prompt_id 到数据库记录
|
||||
if let Ok(Some(mut record)) = outfit_repo.get_record_by_id(&record_id) {
|
||||
record.set_comfyui_prompt_id(prompt_id.clone());
|
||||
if let Err(e) = outfit_repo.update_record(&record) {
|
||||
warn!("保存 ComfyUI prompt_id 失败: {}", e);
|
||||
} else {
|
||||
info!("✅ 已保存 ComfyUI prompt_id: {} 到记录: {}", prompt_id, record_id);
|
||||
}
|
||||
}
|
||||
|
||||
// 处理上传结果
|
||||
let mut successful_uploads = 0;
|
||||
|
||||
@@ -1231,3 +1241,52 @@ async fn perform_outfit_image_generation(
|
||||
info!("✅ 后台穿搭图片生成完成,耗时: {}ms", total_duration);
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// 创建基于 prompt_id 的进度回调函数
|
||||
/// 这个函数用于批量处理场景,通过 prompt_id 精确匹配任务
|
||||
pub fn create_prompt_id_based_progress_callback(
|
||||
database: Arc<Database>,
|
||||
app_handle: AppHandle,
|
||||
) -> impl Fn(String, WorkflowProgress) + Send + Sync + 'static {
|
||||
move |prompt_id: String, progress: WorkflowProgress| {
|
||||
let outfit_repo = OutfitImageRepository::new(database.clone());
|
||||
|
||||
// 根据 prompt_id 查找对应的记录
|
||||
match outfit_repo.get_record_by_comfyui_prompt_id(&prompt_id) {
|
||||
Ok(Some(mut record)) => {
|
||||
let record_id = record.id.clone();
|
||||
|
||||
// 更新数据库中的进度
|
||||
record.progress = progress.progress_percentage / 100.0; // 转换为0-1范围
|
||||
if let Err(e) = outfit_repo.update_record(&record) {
|
||||
warn!("更新记录进度失败: {}", e);
|
||||
} else {
|
||||
info!("✅ 通过 prompt_id {} 更新记录 {} 进度: {}%",
|
||||
prompt_id, record_id, progress.progress_percentage);
|
||||
}
|
||||
|
||||
// 发送进度事件到前端
|
||||
let progress_event = serde_json::json!({
|
||||
"type": "outfit_generation_progress",
|
||||
"record_id": record_id,
|
||||
"prompt_id": prompt_id,
|
||||
"progress_percentage": progress.progress_percentage,
|
||||
"current_step": progress.current_step,
|
||||
"total_steps": progress.total_steps,
|
||||
"status_message": progress.status_message,
|
||||
"current_node_id": progress.current_node_id
|
||||
});
|
||||
|
||||
if let Err(e) = app_handle.emit("outfit_generation_progress", &progress_event) {
|
||||
warn!("发送进度事件失败: {}", e);
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
warn!("⚠️ 未找到 prompt_id {} 对应的记录", prompt_id);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("❌ 查找 prompt_id {} 对应记录失败: {}", prompt_id, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,9 @@ use crate::app_state::AppState;
|
||||
use crate::data::models::outfit_photo_generation::{
|
||||
OutfitPhotoGenerationRequest, OutfitPhotoGenerationResponse, GenerationStatus, WorkflowProgress
|
||||
};
|
||||
use crate::business::services::outfit_photo_generation_service::OutfitPhotoGenerationService;
|
||||
use crate::business::services::outfit_photo_generation_service::{
|
||||
OutfitPhotoGenerationService, create_outfit_photo_prompt_id_based_progress_callback
|
||||
};
|
||||
use crate::business::services::cloud_upload_service::CloudUploadService;
|
||||
use crate::config::AppConfig;
|
||||
|
||||
@@ -163,7 +165,7 @@ pub async fn execute_outfit_photo_generation_batch(
|
||||
|
||||
// 创建穿搭照片生成服务
|
||||
let service = match OutfitPhotoGenerationService::new(
|
||||
database,
|
||||
database.clone(),
|
||||
config,
|
||||
cloud_upload_service,
|
||||
) {
|
||||
@@ -174,31 +176,14 @@ pub async fn execute_outfit_photo_generation_batch(
|
||||
}
|
||||
};
|
||||
|
||||
// 创建批量进度回调函数
|
||||
let app_handle_clone = app_handle.clone();
|
||||
let total_requests = requests.len();
|
||||
let progress_callback = move |current: usize, workflow_progress: WorkflowProgress| {
|
||||
let progress_event = serde_json::json!({
|
||||
"type": "batch_progress",
|
||||
"current": current,
|
||||
"total": total_requests,
|
||||
"percentage": (current as f64 / total_requests as f64) * 100.0,
|
||||
"workflow_progress": {
|
||||
"current_step": workflow_progress.current_step,
|
||||
"total_steps": workflow_progress.total_steps,
|
||||
"current_node_id": workflow_progress.current_node_id,
|
||||
"progress_percentage": workflow_progress.progress_percentage,
|
||||
"status_message": workflow_progress.status_message
|
||||
}
|
||||
});
|
||||
// 创建基于 prompt_id 的进度回调函数
|
||||
let progress_callback = create_outfit_photo_prompt_id_based_progress_callback(
|
||||
database.clone(),
|
||||
app_handle.clone(),
|
||||
);
|
||||
|
||||
if let Err(e) = app_handle_clone.emit("outfit_generation_batch_progress", &progress_event) {
|
||||
error!("发送批量进度事件失败: {}", e);
|
||||
}
|
||||
};
|
||||
|
||||
// 执行批量生成
|
||||
match service.generate_batch(requests, progress_callback).await {
|
||||
// 执行批量生成(使用基于 prompt_id 的方法)
|
||||
match service.generate_batch_with_prompt_id_callback(requests, progress_callback).await {
|
||||
Ok(responses) => {
|
||||
info!("批量穿搭照片生成完成,共 {} 个响应", responses.len());
|
||||
|
||||
|
||||
@@ -597,9 +597,10 @@ pub async fn batch_download_audio_to_directory(
|
||||
let mut failed_downloads = Vec::new();
|
||||
|
||||
for (index, record) in records.iter().enumerate() {
|
||||
let default_id = format!("unknown_{}", index);
|
||||
let record_id = record.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&format!("unknown_{}", index));
|
||||
.unwrap_or(&default_id);
|
||||
|
||||
let audio_url = record.get("audio_url")
|
||||
.and_then(|v| v.as_str())
|
||||
|
||||
@@ -20,7 +20,8 @@ export interface OutfitImageRecord {
|
||||
started_at?: string;
|
||||
completed_at?: string;
|
||||
duration_ms?: number;
|
||||
|
||||
comfyui_prompt_id?: string; // ComfyUI 任务ID,用于精确匹配任务状态更新
|
||||
|
||||
// 关联数据
|
||||
product_images: ProductImage[];
|
||||
outfit_images: OutfitImage[];
|
||||
|
||||
Reference in New Issue
Block a user