diff --git a/apps/desktop/src-tauri/src/lib.rs b/apps/desktop/src-tauri/src/lib.rs index 1d2fa14..1f2d4b6 100644 --- a/apps/desktop/src-tauri/src/lib.rs +++ b/apps/desktop/src-tauri/src/lib.rs @@ -302,6 +302,8 @@ pub fn run() { commands::debug_commands::validate_template_structure, // 便捷工具命令 commands::tools_commands::clean_jsonl_data, + commands::tools_commands::ai_model_face_hair_fix_single_image, + commands::tools_commands::ai_model_face_hair_fix_batch_images, // 服装搭配搜索命令 commands::outfit_search_commands::analyze_outfit_image, commands::outfit_search_commands::search_similar_outfits, diff --git a/apps/desktop/src-tauri/src/presentation/commands/tools_commands.rs b/apps/desktop/src-tauri/src/presentation/commands/tools_commands.rs index 621396d..f7c0888 100644 --- a/apps/desktop/src-tauri/src/presentation/commands/tools_commands.rs +++ b/apps/desktop/src-tauri/src/presentation/commands/tools_commands.rs @@ -1,9 +1,18 @@ use tauri::{command, Emitter}; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; +use std::collections::{HashSet, HashMap}; use std::fs::File; use std::io::{BufRead, BufReader, Write}; -use std::path::Path; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tokio::fs; +use anyhow::Result; +use comfyui_sdk::{ + ComfyUIClient, ComfyUIClientConfig, ExecutionOptions, + AI_MODEL_FACE_HAIR_FIX_TEMPLATE +}; +use comfyui_sdk::utils::SimpleCallbacks; +use serde_json::json; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DataCleaningProgress { @@ -220,3 +229,438 @@ fn count_lines(file_path: &str) -> Result { let count = reader.lines().count(); Ok(count) } + +// ============================================================================ +// 图片增强相关结构体和命令 +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageEnhancementProgress { + pub current: usize, + pub total: usize, + pub percentage: f64, + pub status: String, + pub current_file: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageEnhancementResult { + pub success: bool, + pub message: String, + pub input_path: String, + pub output_path: Option, + pub execution_time: f64, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BatchImageEnhancementResult { + pub success: bool, + pub message: String, + pub total_processed: usize, + pub successful_count: usize, + pub failed_count: usize, + pub results: Vec, + pub total_execution_time: f64, +} + +/// 单张图片AI模型面部头发修复增强处理 +#[command] +pub async fn ai_model_face_hair_fix_single_image( + input_image: String, + server_url: String, + face_prompt: String, + face_denoise: String, + output_image: String, + window: tauri::Window, +) -> Result { + let start_time = std::time::Instant::now(); + + // 发送开始处理的进度 + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current: 0, + total: 1, + percentage: 0.0, + status: "开始图片增强处理...".to_string(), + current_file: Some(input_image.clone()), + }); + + match ai_model_face_hair_fix_internal(&input_image, &server_url, &face_prompt, &face_denoise, &output_image, &window).await { + Ok(_) => { + let execution_time = start_time.elapsed().as_secs_f64(); + + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current: 1, + total: 1, + percentage: 100.0, + status: "图片增强完成".to_string(), + current_file: Some(input_image.clone()), + }); + + Ok(ImageEnhancementResult { + success: true, + message: "图片增强成功".to_string(), + input_path: input_image, + output_path: Some(output_image), + execution_time, + error: None, + }) + } + Err(e) => { + let execution_time = start_time.elapsed().as_secs_f64(); + + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current: 1, + total: 1, + percentage: 100.0, + status: format!("图片增强失败: {}", e), + current_file: Some(input_image.clone()), + }); + + Ok(ImageEnhancementResult { + success: false, + message: format!("图片增强失败: {}", e), + input_path: input_image, + output_path: None, + execution_time, + error: Some(e), + }) + } + } +} + +/// 批量图片AI模型面部头发修复增强处理 +#[command] +pub async fn ai_model_face_hair_fix_batch_images( + input_dir: String, + server_url: String, + face_prompt: String, + face_denoise: String, + output_dir: String, + window: tauri::Window, +) -> Result { + let start_time = std::time::Instant::now(); + + // 发送开始处理的进度 + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current: 0, + total: 0, + percentage: 0.0, + status: "扫描图片文件...".to_string(), + current_file: None, + }); + + // 扫描输入目录中的所有图片文件 + let image_files = match scan_image_files(&input_dir).await { + Ok(files) => files, + Err(e) => return Err(format!("扫描图片文件失败: {}", e)), + }; + + if image_files.is_empty() { + return Ok(BatchImageEnhancementResult { + success: true, + message: "未找到图片文件".to_string(), + total_processed: 0, + successful_count: 0, + failed_count: 0, + results: vec![], + total_execution_time: start_time.elapsed().as_secs_f64(), + }); + } + + let total_files = image_files.len(); + let mut results = Vec::new(); + let mut successful_count = 0; + let mut failed_count = 0; + + // 确保输出目录存在 + if let Err(e) = fs::create_dir_all(&output_dir).await { + return Err(format!("创建输出目录失败: {}", e)); + } + + // 处理每个图片文件 + for (index, (input_path, relative_path)) in image_files.iter().enumerate() { + let current = index + 1; + let percentage = (current as f64 / total_files as f64) * 100.0; + + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current, + total: total_files, + percentage, + status: format!("处理图片 {}/{}", current, total_files), + current_file: Some(input_path.clone()), + }); + + // 构建输出路径，保持目录结构 + let output_path = Path::new(&output_dir).join(relative_path); + + // 确保输出文件的目录存在 + if let Some(parent) = output_path.parent() { + if let Err(e) = fs::create_dir_all(parent).await { + let error_msg = format!("创建输出目录失败: {}", e); + results.push(ImageEnhancementResult { + success: false, + message: error_msg.clone(), + input_path: input_path.clone(), + output_path: None, + execution_time: 0.0, + error: Some(error_msg), + }); + failed_count += 1; + continue; + } + } + + let output_path_str = output_path.to_string_lossy().to_string(); + let file_start_time = std::time::Instant::now(); + + 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(); + results.push(ImageEnhancementResult { + success: true, + message: "图片增强成功".to_string(), + input_path: input_path.clone(), + output_path: Some(output_path_str), + execution_time, + error: None, + }); + successful_count += 1; + } + Err(e) => { + let execution_time = file_start_time.elapsed().as_secs_f64(); + results.push(ImageEnhancementResult { + success: false, + message: format!("图片增强失败: {}", e), + input_path: input_path.clone(), + output_path: None, + execution_time, + error: Some(e), + }); + failed_count += 1; + } + } + } + + let total_execution_time = start_time.elapsed().as_secs_f64(); + + let _ = window.emit("image-enhancement-progress", ImageEnhancementProgress { + current: total_files, + total: total_files, + percentage: 100.0, + status: format!("批量处理完成: 成功 {}, 失败 {}", successful_count, failed_count), + current_file: None, + }); + + Ok(BatchImageEnhancementResult { + success: failed_count == 0, + message: format!("批量处理完成: 总计 {}, 成功 {}, 失败 {}", total_files, successful_count, failed_count), + total_processed: total_files, + successful_count, + failed_count, + results, + total_execution_time, + }) +} + +/// AI模型面部头发修复内部处理函数 +async fn ai_model_face_hair_fix_internal( + input_image: &str, + server_url: &str, + face_prompt: &str, + face_denoise: &str, + output_image: &str, + window: &tauri::Window, +) -> Result<(), String> { + // 验证输入图片文件 + if !Path::new(input_image).exists() { + return Err(format!("输入图片文件不存在: {}", input_image)); + } + + // 验证图片格式 + validate_image_file(input_image)?; + + // 初始化 ComfyUI 客户端 + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: server_url.to_string(), + ..Default::default() + }).map_err(|e| format!("创建ComfyUI客户端失败: {}", e))?; + + // 连接到服务器 + client.connect().await + .map_err(|e| format!("连接ComfyUI服务器失败: {}", e))?; + + // 注册模板 + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone()) + .map_err(|e| format!("注册模板失败: {}", e))?; + + let template = client.templates_ref() + .get_by_id("ai-model-face-hair-fix") + .ok_or("模板未找到")?; + + // 上传图片 + let upload_response = client.upload_image(input_image, false).await + .map_err(|e| format!("上传图片失败: {}", e))?; + + // 设置进度回调 + let callbacks = Arc::new( + SimpleCallbacks::new() + .with_progress(|progress| { + let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32; + // 这里可以发送更详细的进度信息 + }) + .with_executing(|_node_id| { + // 节点执行回调 + }) + .with_error(|error| { + eprintln!("ComfyUI执行错误: {}", error.message); + }) + ); + + // 设置参数 + let mut parameters = HashMap::new(); + parameters.insert("input_image".to_string(), json!(upload_response.name)); + parameters.insert("face_prompt".to_string(), json!(face_prompt)); + parameters.insert("face_denoise".to_string(), json!(face_denoise)); + + // 执行增强 + let result = client.execute_template_with_callbacks( + template, + parameters, + ExecutionOptions { + timeout: Some(std::time::Duration::from_secs(180)), + priority: None, + }, + callbacks, + ).await.map_err(|e| format!("执行图片增强失败: {}", e))?; + + if !result.success { + let error_msg = result.error + .map(|e| e.message) + .unwrap_or_else(|| "未知错误".to_string()); + return Err(format!("图片增强失败: {}", error_msg)); + } + + // 获取输出图片URL并下载 + if let Some(outputs) = &result.outputs { + let image_urls = client.outputs_to_urls(outputs); + + if let Some(first_url) = image_urls.first() { + // 下载增强后的图片 + download_image_from_url(first_url, output_image).await + .map_err(|e| format!("下载增强图片失败: {}", e))?; + } else { + return Err("未找到输出图片URL".to_string()); + } + } else { + return Err("未找到输出结果".to_string()); + } + + // 断开连接 + client.disconnect().await + .map_err(|e| format!("断开连接失败: {}", e))?; + + Ok(()) +} + +/// 扫描目录中的所有图片文件 +async fn scan_image_files(dir: &str) -> Result, String> { + let mut image_files = Vec::new(); + let base_path = Path::new(dir); + + if !base_path.exists() { + return Err(format!("目录不存在: {}", dir)); + } + + scan_directory_recursive(base_path, base_path, &mut image_files).await?; + + Ok(image_files) +} + +/// 递归扫描目录 +fn scan_directory_recursive<'a>( + current_dir: &'a Path, + base_dir: &'a Path, + image_files: &'a mut Vec<(String, String)>, +) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let mut entries = fs::read_dir(current_dir).await + .map_err(|e| format!("读取目录失败: {}", e))?; + + while let Some(entry) = entries.next_entry().await + .map_err(|e| format!("读取目录项失败: {}", e))? { + + let path = entry.path(); + + if path.is_dir() { + // 递归处理子目录 + scan_directory_recursive(&path, base_dir, image_files).await?; + } else if path.is_file() { + // 检查是否为图片文件 + if is_image_file(&path) { + let absolute_path = path.to_string_lossy().to_string(); + let relative_path = path.strip_prefix(base_dir) + .map_err(|e| format!("计算相对路径失败: {}", e))? + .to_string_lossy() + .to_string(); + + image_files.push((absolute_path, relative_path)); + } + } + } + + Ok(()) + }) +} + +/// 检查文件是否为图片文件 +fn is_image_file(path: &Path) -> bool { + if let Some(extension) = path.extension() { + let ext = extension.to_string_lossy().to_lowercase(); + matches!(ext.as_str(), "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp") + } else { + false + } +} + +/// 验证图片文件 +fn validate_image_file(path: &str) -> Result<(), String> { + let path = Path::new(path); + + if !path.exists() { + return Err(format!("文件不存在: {}", path.display())); + } + + if !path.is_file() { + return Err(format!("路径不是文件: {}", path.display())); + } + + // 检查文件扩展名 + if let Some(extension) = path.extension() { + let ext = extension.to_string_lossy().to_lowercase(); + match ext.as_str() { + "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp" => Ok(()), + _ => Err(format!("不支持的图片格式: {}", ext)), + } + } else { + Err("文件没有扩展名".to_string()) + } +} + +/// 从URL下载图片到本地文件 +async fn download_image_from_url(url: &str, output_path: &str) -> Result<(), String> { + let response = reqwest::get(url).await + .map_err(|e| format!("请求图片失败: {}", e))?; + + if !response.status().is_success() { + return Err(format!("下载图片失败，状态码: {}", response.status())); + } + + let bytes = response.bytes().await + .map_err(|e| format!("读取图片数据失败: {}", e))?; + + fs::write(output_path, bytes).await + .map_err(|e| format!("保存图片失败: {}", e))?; + + Ok(()) +} diff --git a/apps/desktop/src/App.tsx b/apps/desktop/src/App.tsx index 95f9dbd..2726634 100644 --- a/apps/desktop/src/App.tsx +++ b/apps/desktop/src/App.tsx @@ -33,6 +33,7 @@ import SimpleHedraLipSyncTool from './pages/tools/SimpleHedraLipSyncTool'; import HedraLipSyncRecords from './pages/tools/HedraLipSyncRecords'; import OmniHumanDetectionTool from './pages/tools/OmniHumanDetectionTool'; import { EnrichedAnalysisDemo } from './pages/tools/EnrichedAnalysisDemo'; +import AiModelFaceHairFixTool from './pages/tools/AiModelFaceHairFixTool'; import MaterialCenter from './pages/MaterialCenter'; import VideoGeneration from './pages/VideoGeneration'; import { OutfitPhotoGenerationPage } from './pages/OutfitPhotoGeneration'; @@ -176,6 +177,7 @@ function App() { } /> } /> } /> + } /> diff --git a/apps/desktop/src/components/Navigation.tsx b/apps/desktop/src/components/Navigation.tsx index c8ee8dd..6ad62e5 100644 --- a/apps/desktop/src/components/Navigation.tsx +++ b/apps/desktop/src/components/Navigation.tsx @@ -7,11 +7,6 @@ import { DocumentDuplicateIcon, WrenchScrewdriverIcon, SparklesIcon, - Cog6ToothIcon, - RectangleStackIcon, - ServerIcon, - PlayIcon, - ChartBarIcon, ChevronDownIcon, } from '@heroicons/react/24/outline'; @@ -74,43 +69,6 @@ const Navigation: React.FC = () => { icon: SparklesIcon, description: 'AI穿搭方案推荐与素材检索' }, - { - name: 'ComfyUI', - icon: RectangleStackIcon, - description: 'AI工作流管理平台', - children: [ - { - name: 'V2 仪表板', - href: '/comfyui-v2-dashboard', - icon: ChartBarIcon, - description: '现代化AI工作流管理仪表板' - }, - { - name: '集群管理', - href: '/comfyui-management', - icon: ServerIcon, - description: '分布式ComfyUI集群管理' - }, - { - name: '工作流测试', - href: '/comfyui-workflow-test', - icon: PlayIcon, - description: '工作流测试和调试' - }, - { - name: '模板创建器测试', - href: '/workflow-template-creator-test', - icon: DocumentDuplicateIcon, - description: '工作流模板创建器功能测试' - } - ] - }, - { - name: 'AI工作流', - href: '/workflows', - icon: Cog6ToothIcon, - description: '管理和执行各种AI生成任务工作流' - }, { name: '工具', href: '/tools', diff --git a/apps/desktop/src/components/comfyui/ExecutionMonitor.tsx b/apps/desktop/src/components/comfyui/ExecutionMonitor.tsx index 512876d..4d8bfbc 100644 --- a/apps/desktop/src/components/comfyui/ExecutionMonitor.tsx +++ b/apps/desktop/src/components/comfyui/ExecutionMonitor.tsx @@ -26,7 +26,7 @@ interface ExecutionMonitorProps { export const ExecutionMonitor: React.FC = ({ className = '', - showCompleted = true, + showCompleted: initialShowCompleted = true, maxItems = 10, }) => { const { @@ -43,6 +43,7 @@ export const ExecutionMonitor: React.FC = ({ const filteredExecutions = useFilteredExecutions(); const [autoRefresh, setAutoRefresh] = useState(true); + const [showCompleted, setShowCompleted] = useState(initialShowCompleted); const [refreshInterval, setRefreshInterval] = useState(); // 自动刷新执行列表 diff --git a/apps/desktop/src/components/comfyui/WorkflowManager.tsx b/apps/desktop/src/components/comfyui/WorkflowManager.tsx index f627be6..d171911 100644 --- a/apps/desktop/src/components/comfyui/WorkflowManager.tsx +++ b/apps/desktop/src/components/comfyui/WorkflowManager.tsx @@ -171,7 +171,7 @@ export const WorkflowManager: React.FC = ({ // 读取文件内容 const fileContent = await invoke('import_workflow_package'); - const result = await invoke('comfyui_v2_import_workflows', { + const result = await invoke<{imported_count: number, skipped_count: number}>('comfyui_v2_import_workflows', { request: { workflow_data: fileContent, overwrite_existing: false diff --git a/apps/desktop/src/components/comfyui/WorkflowV2Creator.tsx b/apps/desktop/src/components/comfyui/WorkflowV2Creator.tsx index cba6b30..8edfba0 100644 --- a/apps/desktop/src/components/comfyui/WorkflowV2Creator.tsx +++ b/apps/desktop/src/components/comfyui/WorkflowV2Creator.tsx @@ -275,7 +275,7 @@ export const WorkflowV2Creator: React.FC = ({

- {editingWorkflow ? '编辑工作流' : '创建工作流'} + {editingTemplate ? '编辑工作流' : '创建工作流'}

@@ -332,8 +332,11 @@ export const WorkflowV2Creator: React.FC = ({ updateField('name', e.target.value)} + value={templateData.metadata.name} + onChange={(e) => setTemplateData(prev => ({ + ...prev, + metadata: { ...prev.metadata, name: e.target.value } + }))} className={`w-full px-3 py-2 border rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent ${ errors.name ? 'border-red-300' : 'border-gray-300' }`} @@ -352,8 +355,11 @@ export const WorkflowV2Creator: React.FC = ({ 描述