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:
imeepos
2025-07-31 15:24:48 +08:00
parent ff5cccfb05
commit 4b20f69560
12 changed files with 639 additions and 41 deletions

View File

@@ -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;

View File

@@ -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);
}
}
}
}

View File

@@ -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);
}
}
/// 商品图片

View File

@@ -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);

View File

@@ -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(), // 将在调用方加载
})

View File

@@ -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>> {
// 优先使用只读连接,提高并发性能

View File

@@ -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, // 暂时不提供回滚脚本
});
}
/// 添加迁移

View File

@@ -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);

View File

@@ -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);
}
}
}
}

View File

@@ -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());

View File

@@ -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())

View File

@@ -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[];