From 0f13428101356c44c89cd31df9e401ff20fca2e7 Mon Sep 17 00:00:00 2001 From: imeepos Date: Fri, 8 Aug 2025 14:33:00 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=A0=B8=E5=BF=83=E6=9C=8D=E5=8A=A1?= =?UTF-8?q?=E9=87=8D=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- TODO.md | 533 +++++++++++++++ .../src/data/migrations/comfyui_tables.sql | 158 +++++ .../src-tauri/src/data/models/comfyui.rs | 523 +++++++++++++- .../data/repositories/comfyui_repository.rs | 646 ++++++++++++++++++ .../src-tauri/src/data/repositories/mod.rs | 1 + 5 files changed, 1843 insertions(+), 18 deletions(-) create mode 100644 TODO.md create mode 100644 apps/desktop/src-tauri/src/data/migrations/comfyui_tables.sql create mode 100644 apps/desktop/src-tauri/src/data/repositories/comfyui_repository.rs diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..ac3fdc3 --- /dev/null +++ b/TODO.md @@ -0,0 +1,533 @@ +# ComfyUI SDK 完全重写集成计划 + +## 📋 项目概述 + +**目标**: 完全基于 `cargos/comfyui-sdk` 重新实现 MixVideo 项目的 ComfyUI 集成，构建现代化、高性能的 AI 工作流管理系统。 + +**核心优势**: +- 🚀 充分利用 SDK 的所有高级功能 +- 🔄 WebSocket 实时通信和进度跟踪 +- 📝 强大的模板系统和参数验证 +- 🛡️ 类型安全和错误处理 +- ⚡ 连接池和性能优化 +- 🎯 统一的架构设计 + +**预计工期**: 3-4 周 +**风险等级**: 中等 (需要充分测试) + +--- + +## 🏗️ 架构设计 + +### 新架构概览 +``` +Frontend (React + TypeScript) + ↓ Tauri Invoke +Backend (Rust + ComfyUI SDK) + ├── Presentation Layer (Commands) + ├── Business Layer (Services) + ├── Data Layer (Models + Database) + └── Infrastructure Layer (ComfyUI SDK) +``` + +### 核心组件 +1. **ComfyUIManager**: 统一的 SDK 管理器 +2. **TemplateEngine**: 工作流模板引擎 +3. **ExecutionEngine**: 工作流执行引擎 +4. **RealtimeMonitor**: 实时状态监控 +5. **CacheManager**: 智能缓存管理 + +--- + +## 📅 详细开发计划 + +## 第一阶段: 基础架构重构 (第1周) + +### 🎯 目标 +建立全新的基于 SDK 的基础架构，替换现有的传统实现。 + +### 📋 任务清单 + +#### 1.1 数据模型重新设计 (2天) +- [ ] **删除旧模型**: 移除 `apps/desktop/src-tauri/src/data/models/comfyui.rs` +- [ ] **创建新模型**: 基于 SDK 类型创建统一数据模型 + - [ ] `WorkflowModel`: 工作流数据模型 + - [ ] `TemplateModel`: 模板数据模型 + - [ ] `ExecutionModel`: 执行记录模型 + - [ ] `QueueModel`: 队列状态模型 +- [ ] **数据库迁移**: 设计新的数据库表结构 + - [ ] `comfyui_workflows` 表 + - [ ] `comfyui_templates` 表 + - [ ] `comfyui_executions` 表 + - [ ] `comfyui_queue_history` 表 + +#### 1.2 核心服务重构 (3天) +- [ ] **删除旧服务**: 移除现有的 ComfyUI 服务实现 + - [ ] 删除 `comfyui_service.rs` + - [ ] 删除 `comfyui_sdk_service.rs` +- [ ] **创建 ComfyUIManager**: 统一的 SDK 管理器 + - [ ] 连接管理和健康检查 + - [ ] 配置管理和验证 + - [ ] 错误处理和重试机制 +- [ ] **创建 TemplateEngine**: 模板管理引擎 + - [ ] 模板加载和验证 + - [ ] 参数解析和验证 + - [ ] 模板实例化 +- [ ] **创建 ExecutionEngine**: 执行引擎 + - [ ] 工作流执行管理 + - [ ] 队列管理 + - [ ] 结果处理 + +#### 1.3 配置系统重构 (1天) +- [ ] **统一配置结构**: 基于 SDK 配置重新设计 +- [ ] **配置验证**: 添加配置有效性检查 +- [ ] **环境适配**: 支持开发/生产环境配置 + +#### 1.4 错误处理系统 (1天) +- [ ] **统一错误类型**: 基于 SDK 错误类型设计 +- [ ] **错误映射**: SDK 错误到应用错误的映射 +- [ ] **错误恢复**: 自动重试和降级策略 + +### 🧪 测试要求 +- [ ] 单元测试覆盖率 > 80% +- [ ] 集成测试验证基础功能 +- [ ] 性能基准测试 + +--- + +## 第二阶段: 命令层重写 (第2周) + +### 🎯 目标 +重新实现所有 Tauri 命令，提供完整的前端 API 接口。 + +### 📋 任务清单 + +#### 2.1 基础命令重写 (2天) +- [ ] **连接管理命令** + - [ ] `comfyui_connect`: 连接到 ComfyUI 服务 + - [ ] `comfyui_disconnect`: 断开连接 + - [ ] `comfyui_health_check`: 健康检查 + - [ ] `comfyui_get_system_info`: 获取系统信息 + +#### 2.2 工作流管理命令 (2天) +- [ ] **工作流 CRUD 命令** + - [ ] `comfyui_list_workflows`: 获取工作流列表 + - [ ] `comfyui_get_workflow`: 获取单个工作流 + - [ ] `comfyui_create_workflow`: 创建工作流 + - [ ] `comfyui_update_workflow`: 更新工作流 + - [ ] `comfyui_delete_workflow`: 删除工作流 + +#### 2.3 模板管理命令 (2天) +- [ ] **模板系统命令** + - [ ] `comfyui_list_templates`: 获取模板列表 + - [ ] `comfyui_get_template`: 获取模板详情 + - [ ] `comfyui_validate_template`: 验证模板 + - [ ] `comfyui_create_template_instance`: 创建模板实例 + +#### 2.4 执行管理命令 (1天) +- [ ] **执行控制命令** + - [ ] `comfyui_execute_workflow`: 执行工作流 + - [ ] `comfyui_execute_template`: 执行模板 + - [ ] `comfyui_cancel_execution`: 取消执行 + - [ ] `comfyui_get_execution_status`: 获取执行状态 + +### 🧪 测试要求 +- [ ] 每个命令的单元测试 +- [ ] 命令集成测试 +- [ ] 错误场景测试 + +--- + +## 第三阶段: 实时通信与高级功能 (第3周) + +### 🎯 目标 +实现 WebSocket 实时通信和高级功能，提供完整的用户体验。 + +### 📋 任务清单 + +#### 3.1 WebSocket 实时通信 (3天) +- [ ] **实时监控服务** + - [ ] WebSocket 连接管理 + - [ ] 事件订阅和分发 + - [ ] 进度更新推送 +- [ ] **前端实时更新** + - [ ] 实时状态显示 + - [ ] 进度条更新 + - [ ] 错误实时提醒 + +#### 3.2 队列管理系统 (2天) +- [ ] **队列监控** + - [ ] 队列状态实时更新 + - [ ] 任务优先级管理 + - [ ] 队列统计信息 +- [ ] **批量操作** + - [ ] 批量提交任务 + - [ ] 批量取消任务 + - [ ] 批量状态查询 + +#### 3.3 缓存和性能优化 (2天) +- [ ] **智能缓存** + - [ ] 工作流结果缓存 + - [ ] 模板缓存 + - [ ] 连接池管理 +- [ ] **性能监控** + - [ ] 执行时间统计 + - [ ] 资源使用监控 + - [ ] 性能指标收集 + +### 🧪 测试要求 +- [ ] WebSocket 连接稳定性测试 +- [ ] 高并发场景测试 +- [ ] 性能压力测试 + +--- + +## 第四阶段: 前端重构与UI优化 (第4周) + +### 🎯 目标 +重构前端代码，提供现代化的用户界面和交互体验。 + +### 📋 任务清单 + +#### 4.1 类型定义更新 (1天) +- [ ] **TypeScript 类型重写** + - [ ] 基于新后端 API 更新类型定义 + - [ ] 添加 SDK 特有的类型 + - [ ] 类型安全验证 + +#### 4.2 服务层重构 (2天) +- [ ] **前端服务重写** + - [ ] 重写 `ComfyuiService` 类 + - [ ] 添加实时通信支持 + - [ ] 错误处理优化 +- [ ] **状态管理优化** + - [ ] Zustand store 重构 + - [ ] 实时状态同步 + - [ ] 缓存策略实现 + +#### 4.3 UI 组件优化 (3天) +- [ ] **核心组件重构** + - [ ] 工作流管理界面 + - [ ] 模板选择器 + - [ ] 执行监控面板 + - [ ] 实时进度显示 +- [ ] **用户体验优化** + - [ ] 加载状态优化 + - [ ] 错误提示改进 + - [ ] 响应式设计完善 + +#### 4.4 高级功能界面 (1天) +- [ ] **新功能界面** + - [ ] 模板编辑器 + - [ ] 批量操作界面 + - [ ] 性能监控面板 + - [ ] 队列管理界面 + +### 🧪 测试要求 +- [ ] 组件单元测试 +- [ ] E2E 测试覆盖 +- [ ] 用户体验测试 + +--- + +## 🗂️ 文件结构规划 + +### 后端文件结构 +``` +apps/desktop/src-tauri/src/ +├── business/ +│ ├── services/ +│ │ ├── comfyui_manager.rs # 核心管理器 +│ │ ├── template_engine.rs # 模板引擎 +│ │ ├── execution_engine.rs # 执行引擎 +│ │ └── realtime_monitor.rs # 实时监控 +│ └── models/ +│ └── comfyui/ # 新的数据模型 +├── data/ +│ ├── repositories/ +│ │ └── comfyui_repository.rs # 数据访问层 +│ └── migrations/ +│ └── comfyui_tables.sql # 数据库迁移 +├── presentation/ +│ └── commands/ +│ └── comfyui_commands.rs # 重写的命令 +└── infrastructure/ + └── comfyui/ # SDK 集成层 +``` + +### 前端文件结构 +``` +apps/desktop/src/ +├── services/ +│ └── comfyuiService.ts # 重写的服务 +├── types/ +│ └── comfyui.ts # 更新的类型定义 +├── components/ +│ ├── ComfyUI/ # ComfyUI 相关组件 +│ │ ├── WorkflowManager.tsx +│ │ ├── TemplateSelector.tsx +│ │ ├── ExecutionMonitor.tsx +│ │ └── RealtimeProgress.tsx +│ └── common/ # 通用组件 +├── pages/ +│ ├── ComfyUIManagement.tsx # 重构的管理页面 +│ └── ComfyUIWorkflowTest.tsx # 重构的测试页面 +└── store/ + └── comfyuiStore.ts # 重构的状态管理 +``` + +--- + +## ⚠️ 风险评估与应对策略 + +### 高风险项 +1. **数据迁移风险** + - 风险: 现有数据丢失或损坏 + - 应对: 完整的数据备份和迁移脚本测试 + +2. **功能兼容性风险** + - 风险: 新实现与现有功能不兼容 + - 应对: 详细的功能对比测试 + +3. **性能回退风险** + - 风险: 新实现性能不如旧版本 + - 应对: 性能基准测试和优化 + +### 中等风险项 +1. **开发周期延长** + - 风险: 复杂度超出预期 + - 应对: 分阶段交付，及时调整计划 + +2. **用户体验变化** + - 风险: 用户不适应新界面 + - 应对: 渐进式界面更新，保留熟悉元素 + +### 应对措施 +- [ ] 每周进行风险评估 +- [ ] 建立回滚机制 +- [ ] 充分的测试覆盖 +- [ ] 用户反馈收集 + +--- + +## 📊 成功指标 + +### 技术指标 +- [ ] 代码覆盖率 > 85% +- [ ] API 响应时间 < 200ms +- [ ] WebSocket 连接稳定性 > 99% +- [ ] 内存使用优化 > 20% + +### 功能指标 +- [ ] 支持所有现有功能 +- [ ] 新增实时监控功能 +- [ ] 新增模板管理功能 +- [ ] 新增批量操作功能 + +### 用户体验指标 +- [ ] 界面响应速度提升 > 30% +- [ ] 错误率降低 > 50% +- [ ] 用户操作步骤减少 > 20% + +--- + +## 🚀 开始执行 + +### 立即行动项 +1. [ ] 确认开发计划和时间安排 +2. [ ] 备份现有代码和数据 +3. [ ] 创建开发分支 `feature/comfyui-sdk-rewrite` +4. [ ] 开始第一阶段开发 + +### 每日检查项 +- [ ] 代码提交和备份 +- [ ] 测试执行和结果记录 +- [ ] 进度更新和风险评估 +- [ ] 团队沟通和问题解决 + +--- + +## 📝 实施细节补充 + +### 关键技术决策 + +#### 1. 数据库设计 +```sql +-- 工作流表 +CREATE TABLE comfyui_workflows ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + workflow_data TEXT NOT NULL, -- JSON + version TEXT DEFAULT '1.0', + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP +); + +-- 模板表 +CREATE TABLE comfyui_templates ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + category TEXT, + description TEXT, + template_data TEXT NOT NULL, -- JSON + parameter_schema TEXT NOT NULL, -- JSON + created_at DATETIME DEFAULT CURRENT_TIMESTAMP +); + +-- 执行记录表 +CREATE TABLE comfyui_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT, + template_id TEXT, + prompt_id TEXT, + status TEXT NOT NULL, -- pending, running, completed, failed + parameters TEXT, -- JSON + results TEXT, -- JSON + error_message TEXT, + execution_time INTEGER, -- milliseconds + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + completed_at DATETIME +); +``` + +#### 2. 核心服务接口设计 +```rust +// ComfyUIManager 核心接口 +pub trait ComfyUIManagerTrait { + async fn connect(&mut self) -> Result<()>; + async fn disconnect(&mut self) -> Result<()>; + async fn is_connected(&self) -> bool; + async fn get_system_info(&self) -> Result; + async fn health_check(&self) -> Result; +} + +// TemplateEngine 接口 +pub trait TemplateEngineTrait { + async fn load_template(&self, id: &str) -> Result; + async fn validate_parameters(&self, template_id: &str, params: &ParameterValues) -> Result; + async fn create_instance(&self, template_id: &str, params: ParameterValues) -> Result; +} + +// ExecutionEngine 接口 +pub trait ExecutionEngineTrait { + async fn execute_workflow(&self, workflow: &WorkflowInstance) -> Result; + async fn execute_template(&self, template_id: &str, params: ParameterValues) -> Result; + async fn cancel_execution(&self, execution_id: &str) -> Result<()>; + async fn get_execution_status(&self, execution_id: &str) -> Result; +} +``` + +#### 3. 错误处理策略 +```rust +#[derive(Debug, thiserror::Error)] +pub enum ComfyUIError { + #[error("Connection error: {0}")] + Connection(String), + + #[error("Template error: {0}")] + Template(String), + + #[error("Execution error: {0}")] + Execution(String), + + #[error("Validation error: {0}")] + Validation(String), + + #[error("SDK error: {0}")] + Sdk(#[from] comfyui_sdk::error::ComfyUIError), +} +``` + +### 开发环境准备 + +#### 必需工具 +- [ ] Rust 1.70+ +- [ ] Node.js 18+ +- [ ] PNPM 8+ +- [ ] SQLite 3 +- [ ] ComfyUI 服务器 (用于测试) + +#### 开发配置 +```toml +# Cargo.toml 依赖更新 +[dependencies] +comfyui-sdk = { path = "../../../cargos/comfyui-sdk", features = ["websocket", "templates"] } +tokio = { version = "1.0", features = ["full"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +anyhow = "1.0" +thiserror = "1.0" +tracing = "0.1" +sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite", "chrono"] } +``` + +### 测试策略 + +#### 单元测试 +```rust +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_comfyui_manager_connect() { + // 测试连接功能 + } + + #[tokio::test] + async fn test_template_validation() { + // 测试模板验证 + } + + #[tokio::test] + async fn test_workflow_execution() { + // 测试工作流执行 + } +} +``` + +#### 集成测试 +- [ ] 端到端工作流测试 +- [ ] WebSocket 连接测试 +- [ ] 数据库操作测试 +- [ ] 错误恢复测试 + +### 部署和发布 + +#### 发布检查清单 +- [ ] 所有测试通过 +- [ ] 性能指标达标 +- [ ] 文档更新完成 +- [ ] 用户反馈收集 +- [ ] 回滚方案准备 + +#### 发布策略 +1. **Alpha 版本**: 内部测试 (第2周末) +2. **Beta 版本**: 小范围用户测试 (第3周末) +3. **正式版本**: 全量发布 (第4周末) + +--- + +## 📞 联系和支持 + +### 技术支持 +- **开发问题**: 查看项目 Issues +- **SDK 问题**: 参考 `cargos/comfyui-sdk/README.md` +- **架构问题**: 联系架构师 + +### 文档资源 +- [ ] API 文档: `docs/api/comfyui.md` +- [ ] 开发指南: `docs/development/comfyui-integration.md` +- [ ] 故障排除: `docs/troubleshooting/comfyui.md` + +--- + +**最后更新**: 2025-01-08 +**负责人**: 开发团队 +**审核人**: 项目经理 +**版本**: v1.0 + +> 💡 **提示**: 这是一个完全重写的计划，需要充分的准备和测试。建议在开始前进行详细的技术评审和团队讨论。 + +> ⚠️ **重要**: 开始开发前请务必备份现有代码和数据，并创建独立的开发分支。 diff --git a/apps/desktop/src-tauri/src/data/migrations/comfyui_tables.sql b/apps/desktop/src-tauri/src/data/migrations/comfyui_tables.sql new file mode 100644 index 0000000..59a43a3 --- /dev/null +++ b/apps/desktop/src-tauri/src/data/migrations/comfyui_tables.sql @@ -0,0 +1,158 @@ +-- ComfyUI 数据库表结构迁移脚本 +-- 基于新的 SDK 数据模型设计 + +-- 工作流表 +CREATE TABLE IF NOT EXISTS comfyui_workflows ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + workflow_data TEXT NOT NULL, -- JSON 格式的工作流数据 + version TEXT NOT NULL DEFAULT '1.0', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + enabled BOOLEAN NOT NULL DEFAULT 1, + tags TEXT, -- JSON 数组格式的标签 + category TEXT +); + +-- 工作流表索引 +CREATE INDEX IF NOT EXISTS idx_comfyui_workflows_name ON comfyui_workflows(name); +CREATE INDEX IF NOT EXISTS idx_comfyui_workflows_category ON comfyui_workflows(category); +CREATE INDEX IF NOT EXISTS idx_comfyui_workflows_created_at ON comfyui_workflows(created_at); +CREATE INDEX IF NOT EXISTS idx_comfyui_workflows_enabled ON comfyui_workflows(enabled); + +-- 模板表 +CREATE TABLE IF NOT EXISTS comfyui_templates ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + category TEXT, + description TEXT, + template_data TEXT NOT NULL, -- JSON 格式的模板数据 + parameter_schema TEXT NOT NULL, -- JSON 格式的参数定义 + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + enabled BOOLEAN NOT NULL DEFAULT 1, + tags TEXT, -- JSON 数组格式的标签 + author TEXT, + version TEXT NOT NULL DEFAULT '1.0' +); + +-- 模板表索引 +CREATE INDEX IF NOT EXISTS idx_comfyui_templates_name ON comfyui_templates(name); +CREATE INDEX IF NOT EXISTS idx_comfyui_templates_category ON comfyui_templates(category); +CREATE INDEX IF NOT EXISTS idx_comfyui_templates_author ON comfyui_templates(author); +CREATE INDEX IF NOT EXISTS idx_comfyui_templates_created_at ON comfyui_templates(created_at); +CREATE INDEX IF NOT EXISTS idx_comfyui_templates_enabled ON comfyui_templates(enabled); + +-- 执行记录表 +CREATE TABLE IF NOT EXISTS comfyui_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT, -- 关联的工作流 ID（可选） + template_id TEXT, -- 关联的模板 ID（可选） + prompt_id TEXT NOT NULL, -- ComfyUI 提示 ID + status TEXT NOT NULL CHECK (status IN ('pending', 'running', 'completed', 'failed', 'cancelled')), + parameters TEXT, -- JSON 格式的执行参数 + results TEXT, -- JSON 格式的执行结果 + output_urls TEXT, -- JSON 数组格式的输出文件 URLs + error_message TEXT, + execution_time INTEGER, -- 执行时间（毫秒） + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + completed_at DATETIME, + client_id TEXT, + node_outputs TEXT, -- JSON 格式的节点输出详情 + FOREIGN KEY (workflow_id) REFERENCES comfyui_workflows(id) ON DELETE SET NULL, + FOREIGN KEY (template_id) REFERENCES comfyui_templates(id) ON DELETE SET NULL +); + +-- 执行记录表索引 +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_workflow_id ON comfyui_executions(workflow_id); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_template_id ON comfyui_executions(template_id); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_prompt_id ON comfyui_executions(prompt_id); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_status ON comfyui_executions(status); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_created_at ON comfyui_executions(created_at); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_completed_at ON comfyui_executions(completed_at); +CREATE INDEX IF NOT EXISTS idx_comfyui_executions_client_id ON comfyui_executions(client_id); + +-- 队列历史表 +CREATE TABLE IF NOT EXISTS comfyui_queue_history ( + id TEXT PRIMARY KEY, + running_tasks TEXT NOT NULL, -- JSON 数组格式的正在运行任务 + pending_tasks TEXT NOT NULL, -- JSON 数组格式的等待中任务 + snapshot_time DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + total_tasks INTEGER NOT NULL DEFAULT 0, + completed_tasks INTEGER NOT NULL DEFAULT 0 +); + +-- 队列历史表索引 +CREATE INDEX IF NOT EXISTS idx_comfyui_queue_history_snapshot_time ON comfyui_queue_history(snapshot_time); + +-- 配置表（用于存储 ComfyUI 服务配置） +CREATE TABLE IF NOT EXISTS comfyui_configs ( + id TEXT PRIMARY KEY DEFAULT 'default', + base_url TEXT NOT NULL, + timeout_seconds INTEGER NOT NULL DEFAULT 300, + retry_attempts INTEGER NOT NULL DEFAULT 3, + retry_delay_ms INTEGER NOT NULL DEFAULT 1000, + enable_websocket BOOLEAN NOT NULL DEFAULT 1, + enable_cache BOOLEAN NOT NULL DEFAULT 1, + max_concurrency INTEGER NOT NULL DEFAULT 4, + custom_headers TEXT, -- JSON 格式的自定义请求头 + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +-- 插入默认配置 +INSERT OR IGNORE INTO comfyui_configs (id, base_url) VALUES ('default', 'http://localhost:8188'); + +-- 创建触发器以自动更新 updated_at 字段 +CREATE TRIGGER IF NOT EXISTS update_comfyui_workflows_updated_at + AFTER UPDATE ON comfyui_workflows + FOR EACH ROW +BEGIN + UPDATE comfyui_workflows SET updated_at = CURRENT_TIMESTAMP WHERE id = NEW.id; +END; + +CREATE TRIGGER IF NOT EXISTS update_comfyui_templates_updated_at + AFTER UPDATE ON comfyui_templates + FOR EACH ROW +BEGIN + UPDATE comfyui_templates SET updated_at = CURRENT_TIMESTAMP WHERE id = NEW.id; +END; + +CREATE TRIGGER IF NOT EXISTS update_comfyui_configs_updated_at + AFTER UPDATE ON comfyui_configs + FOR EACH ROW +BEGIN + UPDATE comfyui_configs SET updated_at = CURRENT_TIMESTAMP WHERE id = NEW.id; +END; + +-- 创建视图以便于查询 +CREATE VIEW IF NOT EXISTS comfyui_execution_summary AS +SELECT + e.id, + e.prompt_id, + e.status, + e.created_at, + e.completed_at, + e.execution_time, + COALESCE(w.name, t.name) as workflow_name, + CASE + WHEN e.workflow_id IS NOT NULL THEN 'workflow' + WHEN e.template_id IS NOT NULL THEN 'template' + ELSE 'unknown' + END as execution_type +FROM comfyui_executions e +LEFT JOIN comfyui_workflows w ON e.workflow_id = w.id +LEFT JOIN comfyui_templates t ON e.template_id = t.id; + +-- 创建统计视图 +CREATE VIEW IF NOT EXISTS comfyui_execution_stats AS +SELECT + status, + COUNT(*) as count, + AVG(execution_time) as avg_execution_time, + MIN(execution_time) as min_execution_time, + MAX(execution_time) as max_execution_time +FROM comfyui_executions +WHERE execution_time IS NOT NULL +GROUP BY status; diff --git a/apps/desktop/src-tauri/src/data/models/comfyui.rs b/apps/desktop/src-tauri/src/data/models/comfyui.rs index bca01c3..e9c2c04 100644 --- a/apps/desktop/src-tauri/src/data/models/comfyui.rs +++ b/apps/desktop/src-tauri/src/data/models/comfyui.rs @@ -1,34 +1,521 @@ +//! ComfyUI 数据模型 +//! 基于 comfyui-sdk 重新设计的统一数据模型 + use serde::{Deserialize, Serialize}; use chrono::{DateTime, Utc}; use std::collections::HashMap; +use uuid::Uuid; -/// ComfyUI API 服务配置 +// 重新导出 SDK 类型 +pub use comfyui_sdk::types::{ + ComfyUIWorkflow, ParameterSchema, ParameterType, ParameterValues, + ValidationResult, ValidationError, TemplateMetadata, WorkflowTemplateData, + QueueStatus, SystemStats, ObjectInfo +}; + +/// 工作流数据模型 +/// 存储完整的工作流信息，包括元数据和工作流定义 #[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ComfyuiConfig { - /// API 基础 URL - pub base_url: String, - /// 请求超时时间（秒） - pub timeout: Option, - /// 重试次数 - pub retry_attempts: Option, - /// 是否启用缓存 - pub enable_cache: Option, - /// 最大并发数 - pub max_concurrency: Option, +pub struct WorkflowModel { + /// 工作流唯一标识符 + pub id: String, + /// 工作流名称 + pub name: String, + /// 工作流描述 + pub description: Option, + /// 工作流 JSON 数据 + pub workflow_data: ComfyUIWorkflow, + /// 工作流版本 + pub version: String, + /// 创建时间 + pub created_at: DateTime, + /// 更新时间 + pub updated_at: DateTime, + /// 是否启用 + pub enabled: bool, + /// 标签 + pub tags: Vec, + /// 分类 + pub category: Option, } -impl Default for ComfyuiConfig { +impl WorkflowModel { + /// 创建新的工作流模型 + pub fn new(name: String, workflow_data: ComfyUIWorkflow) -> Self { + let now = Utc::now(); + Self { + id: Uuid::new_v4().to_string(), + name, + description: None, + workflow_data, + version: "1.0".to_string(), + created_at: now, + updated_at: now, + enabled: true, + tags: Vec::new(), + category: None, + } + } + + /// 更新工作流数据 + pub fn update_workflow(&mut self, workflow_data: ComfyUIWorkflow) { + self.workflow_data = workflow_data; + self.updated_at = Utc::now(); + } + + /// 更新元数据 + pub fn update_metadata(&mut self, name: Option, description: Option, tags: Option>) { + if let Some(name) = name { + self.name = name; + } + if let Some(description) = description { + self.description = Some(description); + } + if let Some(tags) = tags { + self.tags = tags; + } + self.updated_at = Utc::now(); + } +} + +/// 模板数据模型 +/// 存储工作流模板信息，包括参数定义和模板数据 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TemplateModel { + /// 模板唯一标识符 + pub id: String, + /// 模板名称 + pub name: String, + /// 模板分类 + pub category: Option, + /// 模板描述 + pub description: Option, + /// 模板数据 + pub template_data: WorkflowTemplateData, + /// 参数定义 + pub parameter_schema: HashMap, + /// 创建时间 + pub created_at: DateTime, + /// 更新时间 + pub updated_at: DateTime, + /// 是否启用 + pub enabled: bool, + /// 标签 + pub tags: Vec, + /// 作者 + pub author: Option, + /// 版本 + pub version: String, +} + +impl TemplateModel { + /// 创建新的模板模型 + pub fn new(name: String, template_data: WorkflowTemplateData, parameter_schema: HashMap) -> Self { + let now = Utc::now(); + Self { + id: Uuid::new_v4().to_string(), + name, + category: None, + description: None, + template_data, + parameter_schema, + created_at: now, + updated_at: now, + enabled: true, + tags: Vec::new(), + author: None, + version: "1.0".to_string(), + } + } + + /// 验证参数 + pub fn validate_parameters(&self, params: &ParameterValues) -> ValidationResult { + let mut result = ValidationResult::success(); + + // 检查必需参数 + for (param_name, schema) in &self.parameter_schema { + if schema.required.unwrap_or(false) && !params.contains_key(param_name) { + result.add_error(ValidationError::new( + param_name, + format!("Required parameter '{}' is missing", param_name) + )); + } + } + + // 检查参数类型和值 + for (param_name, value) in params { + if let Some(schema) = self.parameter_schema.get(param_name) { + if let Err(error) = self.validate_parameter_value(param_name, value, schema) { + result.add_error(error); + } + } + } + + result + } + + /// 验证单个参数值 + fn validate_parameter_value(&self, param_name: &str, value: &serde_json::Value, schema: &ParameterSchema) -> Result<(), ValidationError> { + match (&schema.param_type, value) { + (ParameterType::String, serde_json::Value::String(_)) => Ok(()), + (ParameterType::Number, serde_json::Value::Number(_)) => Ok(()), + (ParameterType::Boolean, serde_json::Value::Bool(_)) => Ok(()), + (ParameterType::Array, serde_json::Value::Array(_)) => Ok(()), + (ParameterType::Object, serde_json::Value::Object(_)) => Ok(()), + _ => Err(ValidationError::with_value( + param_name, + format!("Parameter '{}' has invalid type, expected {:?}", param_name, schema.param_type), + value.clone() + )), + } + } + + /// 更新模板数据 + pub fn update_template(&mut self, template_data: WorkflowTemplateData, parameter_schema: HashMap) { + self.template_data = template_data; + self.parameter_schema = parameter_schema; + self.updated_at = Utc::now(); + } +} + +/// 执行状态枚举 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum ExecutionStatus { + /// 等待中 + Pending, + /// 运行中 + Running, + /// 已完成 + Completed, + /// 失败 + Failed, + /// 已取消 + Cancelled, +} + +impl std::fmt::Display for ExecutionStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ExecutionStatus::Pending => write!(f, "pending"), + ExecutionStatus::Running => write!(f, "running"), + ExecutionStatus::Completed => write!(f, "completed"), + ExecutionStatus::Failed => write!(f, "failed"), + ExecutionStatus::Cancelled => write!(f, "cancelled"), + } + } +} + +/// 执行记录模型 +/// 存储工作流执行的详细信息和结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecutionModel { + /// 执行唯一标识符 + pub id: String, + /// 关联的工作流 ID（可选） + pub workflow_id: Option, + /// 关联的模板 ID（可选） + pub template_id: Option, + /// ComfyUI 提示 ID + pub prompt_id: String, + /// 执行状态 + pub status: ExecutionStatus, + /// 执行参数 + pub parameters: Option, + /// 执行结果 + pub results: Option>, + /// 输出文件 URLs + pub output_urls: Vec, + /// 错误信息 + pub error_message: Option, + /// 执行时间（毫秒） + pub execution_time: Option, + /// 创建时间 + pub created_at: DateTime, + /// 完成时间 + pub completed_at: Option>, + /// 客户端 ID + pub client_id: Option, + /// 节点输出详情 + pub node_outputs: Option>, +} + +impl ExecutionModel { + /// 创建新的执行记录 + pub fn new(prompt_id: String) -> Self { + Self { + id: Uuid::new_v4().to_string(), + workflow_id: None, + template_id: None, + prompt_id, + status: ExecutionStatus::Pending, + parameters: None, + results: None, + output_urls: Vec::new(), + error_message: None, + execution_time: None, + created_at: Utc::now(), + completed_at: None, + client_id: None, + node_outputs: None, + } + } + + /// 创建工作流执行记录 + pub fn for_workflow(workflow_id: String, prompt_id: String, parameters: Option) -> Self { + let mut execution = Self::new(prompt_id); + execution.workflow_id = Some(workflow_id); + execution.parameters = parameters; + execution + } + + /// 创建模板执行记录 + pub fn for_template(template_id: String, prompt_id: String, parameters: ParameterValues) -> Self { + let mut execution = Self::new(prompt_id); + execution.template_id = Some(template_id); + execution.parameters = Some(parameters); + execution + } + + /// 更新执行状态 + pub fn update_status(&mut self, status: ExecutionStatus) { + self.status = status; + if matches!(status, ExecutionStatus::Completed | ExecutionStatus::Failed | ExecutionStatus::Cancelled) { + self.completed_at = Some(Utc::now()); + if let Some(created_at) = self.created_at.timestamp_millis().try_into().ok() { + if let Some(completed_at) = self.completed_at.and_then(|t| t.timestamp_millis().try_into().ok()) { + self.execution_time = Some(completed_at - created_at); + } + } + } + } + + /// 设置执行结果 + pub fn set_results(&mut self, results: HashMap, output_urls: Vec) { + self.results = Some(results); + self.output_urls = output_urls; + self.update_status(ExecutionStatus::Completed); + } + + /// 设置执行错误 + pub fn set_error(&mut self, error_message: String) { + self.error_message = Some(error_message); + self.update_status(ExecutionStatus::Failed); + } + + /// 检查是否已完成 + pub fn is_completed(&self) -> bool { + matches!(self.status, ExecutionStatus::Completed | ExecutionStatus::Failed | ExecutionStatus::Cancelled) + } + + /// 检查是否成功 + pub fn is_successful(&self) -> bool { + self.status == ExecutionStatus::Completed + } +} + +/// 队列状态模型 +/// 存储 ComfyUI 队列的状态信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueModel { + /// 队列快照 ID + pub id: String, + /// 正在运行的任务 + pub running_tasks: Vec, + /// 等待中的任务 + pub pending_tasks: Vec, + /// 快照时间 + pub snapshot_time: DateTime, + /// 总任务数 + pub total_tasks: u32, + /// 已完成任务数 + pub completed_tasks: u32, +} + +/// 队列任务信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueTaskInfo { + /// 提示 ID + pub prompt_id: String, + /// 任务编号 + pub number: u32, + /// 任务状态 + pub status: String, + /// 提交时间 + pub submitted_at: Option>, + /// 开始时间 + pub started_at: Option>, + /// 预估剩余时间（秒） + pub estimated_remaining: Option, +} + +impl QueueModel { + /// 从 SDK 队列状态创建模型 + pub fn from_queue_status(queue_status: &QueueStatus) -> Self { + let running_tasks: Vec = queue_status.queue_running + .iter() + .map(|item| QueueTaskInfo { + prompt_id: item.prompt_id.clone(), + number: item.number, + status: "running".to_string(), + submitted_at: None, + started_at: Some(Utc::now()), // 假设正在运行的任务刚开始 + estimated_remaining: None, + }) + .collect(); + + let pending_tasks: Vec = queue_status.queue_pending + .iter() + .map(|item| QueueTaskInfo { + prompt_id: item.prompt_id.clone(), + number: item.number, + status: "pending".to_string(), + submitted_at: Some(Utc::now()), // 假设等待中的任务刚提交 + started_at: None, + estimated_remaining: None, + }) + .collect(); + + let total_tasks = (running_tasks.len() + pending_tasks.len()) as u32; + + Self { + id: Uuid::new_v4().to_string(), + running_tasks, + pending_tasks, + snapshot_time: Utc::now(), + total_tasks, + completed_tasks: 0, // 这个需要从历史记录中计算 + } + } + + /// 获取队列统计信息 + pub fn get_statistics(&self) -> QueueStatistics { + QueueStatistics { + total_running: self.running_tasks.len() as u32, + total_pending: self.pending_tasks.len() as u32, + total_tasks: self.total_tasks, + completed_tasks: self.completed_tasks, + average_wait_time: self.calculate_average_wait_time(), + } + } + + /// 计算平均等待时间 + fn calculate_average_wait_time(&self) -> Option { + if self.pending_tasks.is_empty() { + return None; + } + + let now = Utc::now(); + let total_wait_time: i64 = self.pending_tasks + .iter() + .filter_map(|task| task.submitted_at) + .map(|submitted| (now - submitted).num_seconds()) + .sum(); + + if total_wait_time > 0 { + Some((total_wait_time / self.pending_tasks.len() as i64) as u64) + } else { + None + } + } +} + +/// 队列统计信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueStatistics { + /// 正在运行的任务数 + pub total_running: u32, + /// 等待中的任务数 + pub total_pending: u32, + /// 总任务数 + pub total_tasks: u32, + /// 已完成任务数 + pub completed_tasks: u32, + /// 平均等待时间（秒） + pub average_wait_time: Option, +} + +/// ComfyUI 服务配置 +/// 基于 SDK 的统一配置结构 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ComfyUIConfig { + /// 服务基础 URL + pub base_url: String, + /// 连接超时时间（秒） + pub timeout_seconds: u64, + /// 重试次数 + pub retry_attempts: u32, + /// 重试延迟（毫秒） + pub retry_delay_ms: u64, + /// 是否启用 WebSocket + pub enable_websocket: bool, + /// 是否启用缓存 + pub enable_cache: bool, + /// 最大并发数 + pub max_concurrency: u32, + /// 自定义请求头 + pub custom_headers: Option>, +} + +impl Default for ComfyUIConfig { fn default() -> Self { Self { - base_url: "https://bowongai-dev--waas-demo-fastapi-webapp.modal.run".to_string(), - timeout: Some(600), // 10分钟超时，适应 ComfyUI 工作流的长时间处理 - retry_attempts: Some(3), - enable_cache: Some(true), - max_concurrency: Some(8), + base_url: "http://localhost:8188".to_string(), + timeout_seconds: 300, // 5分钟默认超时 + retry_attempts: 3, + retry_delay_ms: 1000, + enable_websocket: true, + enable_cache: true, + max_concurrency: 4, + custom_headers: None, } } } +impl ComfyUIConfig { + /// 转换为 SDK 客户端配置 + pub fn to_sdk_config(&self) -> comfyui_sdk::types::ComfyUIClientConfig { + comfyui_sdk::types::ComfyUIClientConfig { + base_url: self.base_url.clone(), + timeout: Some(std::time::Duration::from_secs(self.timeout_seconds)), + retry_attempts: Some(self.retry_attempts), + retry_delay: Some(std::time::Duration::from_millis(self.retry_delay_ms)), + headers: self.custom_headers.clone(), + } + } + + /// 验证配置 + pub fn validate(&self) -> ValidationResult { + let mut result = ValidationResult::success(); + + // 验证 URL + if self.base_url.is_empty() { + result.add_error(ValidationError::new("base_url", "Base URL cannot be empty")); + } else if url::Url::parse(&self.base_url).is_err() { + result.add_error(ValidationError::new("base_url", "Invalid URL format")); + } + + // 验证超时时间 + if self.timeout_seconds == 0 { + result.add_error(ValidationError::new("timeout_seconds", "Timeout must be greater than 0")); + } + + // 验证重试次数 + if self.retry_attempts > 10 { + result.add_error(ValidationError::new("retry_attempts", "Retry attempts should not exceed 10")); + } + + // 验证并发数 + if self.max_concurrency == 0 { + result.add_error(ValidationError::new("max_concurrency", "Max concurrency must be greater than 0")); + } else if self.max_concurrency > 100 { + result.add_error(ValidationError::new("max_concurrency", "Max concurrency should not exceed 100")); + } + + result + } +} + /// 工作流对象 - 用于 GET /api/workflow 响应 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Workflow { diff --git a/apps/desktop/src-tauri/src/data/repositories/comfyui_repository.rs b/apps/desktop/src-tauri/src/data/repositories/comfyui_repository.rs new file mode 100644 index 0000000..059ee89 --- /dev/null +++ b/apps/desktop/src-tauri/src/data/repositories/comfyui_repository.rs @@ -0,0 +1,646 @@ +//! ComfyUI 数据访问层 +//! 提供对 ComfyUI 相关数据的 CRUD 操作 + +use anyhow::{Result, anyhow}; +use rusqlite::{Connection, params, Row}; +use serde_json; +use std::collections::HashMap; +use tracing::{info, warn, error, debug}; + +use crate::data::models::comfyui::{ + WorkflowModel, TemplateModel, ExecutionModel, QueueModel, ComfyUIConfig, + ExecutionStatus, QueueStatistics +}; + +/// ComfyUI 数据仓库 +pub struct ComfyUIRepository { + db_path: String, +} + +impl ComfyUIRepository { + /// 创建新的数据仓库实例 + pub fn new(db_path: String) -> Self { + Self { db_path } + } + + /// 获取数据库连接 + fn get_connection(&self) -> Result { + let conn = Connection::open(&self.db_path)?; + // 启用外键约束 + conn.execute("PRAGMA foreign_keys = ON", [])?; + Ok(conn) + } + + /// 初始化数据库表 + pub fn initialize(&self) -> Result<()> { + let conn = self.get_connection()?; + + // 读取并执行迁移脚本 + let migration_sql = include_str!("../migrations/comfyui_tables.sql"); + conn.execute_batch(migration_sql)?; + + info!("ComfyUI 数据库表初始化完成"); + Ok(()) + } + + // ==================== 工作流相关操作 ==================== + + /// 创建工作流 + pub fn create_workflow(&self, workflow: &WorkflowModel) -> Result<()> { + let conn = self.get_connection()?; + + let workflow_data_json = serde_json::to_string(&workflow.workflow_data)?; + let tags_json = serde_json::to_string(&workflow.tags)?; + + conn.execute( + "INSERT INTO comfyui_workflows (id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + params![ + workflow.id, + workflow.name, + workflow.description, + workflow_data_json, + workflow.version, + workflow.created_at, + workflow.updated_at, + workflow.enabled, + tags_json, + workflow.category + ], + )?; + + debug!("创建工作流: {}", workflow.id); + Ok(()) + } + + /// 获取工作流 + pub fn get_workflow(&self, id: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category + FROM comfyui_workflows WHERE id = ?1" + )?; + + let workflow_iter = stmt.query_map([id], |row| { + self.row_to_workflow_model(row) + })?; + + for workflow in workflow_iter { + return Ok(Some(workflow?)); + } + + Ok(None) + } + + /// 获取所有工作流 + pub fn list_workflows(&self, enabled_only: bool) -> Result> { + let conn = self.get_connection()?; + + let sql = if enabled_only { + "SELECT id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category + FROM comfyui_workflows WHERE enabled = 1 ORDER BY created_at DESC" + } else { + "SELECT id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category + FROM comfyui_workflows ORDER BY created_at DESC" + }; + + let mut stmt = conn.prepare(sql)?; + let workflow_iter = stmt.query_map([], |row| { + self.row_to_workflow_model(row) + })?; + + let mut workflows = Vec::new(); + for workflow in workflow_iter { + workflows.push(workflow?); + } + + Ok(workflows) + } + + /// 更新工作流 + pub fn update_workflow(&self, workflow: &WorkflowModel) -> Result<()> { + let conn = self.get_connection()?; + + let workflow_data_json = serde_json::to_string(&workflow.workflow_data)?; + let tags_json = serde_json::to_string(&workflow.tags)?; + + let rows_affected = conn.execute( + "UPDATE comfyui_workflows + SET name = ?2, description = ?3, workflow_data = ?4, version = ?5, updated_at = ?6, enabled = ?7, tags = ?8, category = ?9 + WHERE id = ?1", + params![ + workflow.id, + workflow.name, + workflow.description, + workflow_data_json, + workflow.version, + workflow.updated_at, + workflow.enabled, + tags_json, + workflow.category + ], + )?; + + if rows_affected == 0 { + return Err(anyhow!("工作流不存在: {}", workflow.id)); + } + + debug!("更新工作流: {}", workflow.id); + Ok(()) + } + + /// 删除工作流 + pub fn delete_workflow(&self, id: &str) -> Result<()> { + let conn = self.get_connection()?; + + let rows_affected = conn.execute("DELETE FROM comfyui_workflows WHERE id = ?1", [id])?; + + if rows_affected == 0 { + return Err(anyhow!("工作流不存在: {}", id)); + } + + debug!("删除工作流: {}", id); + Ok(()) + } + + /// 按分类获取工作流 + pub fn get_workflows_by_category(&self, category: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category + FROM comfyui_workflows WHERE category = ?1 AND enabled = 1 ORDER BY created_at DESC" + )?; + + let workflow_iter = stmt.query_map([category], |row| { + self.row_to_workflow_model(row) + })?; + + let mut workflows = Vec::new(); + for workflow in workflow_iter { + workflows.push(workflow?); + } + + Ok(workflows) + } + + /// 搜索工作流 + pub fn search_workflows(&self, query: &str) -> Result> { + let conn = self.get_connection()?; + + let search_pattern = format!("%{}%", query); + let mut stmt = conn.prepare( + "SELECT id, name, description, workflow_data, version, created_at, updated_at, enabled, tags, category + FROM comfyui_workflows + WHERE (name LIKE ?1 OR description LIKE ?1) AND enabled = 1 + ORDER BY created_at DESC" + )?; + + let workflow_iter = stmt.query_map([&search_pattern], |row| { + self.row_to_workflow_model(row) + })?; + + let mut workflows = Vec::new(); + for workflow in workflow_iter { + workflows.push(workflow?); + } + + Ok(workflows) + } + + // ==================== 辅助方法 ==================== + + /// 将数据库行转换为工作流模型 + fn row_to_workflow_model(&self, row: &Row) -> rusqlite::Result { + let workflow_data_json: String = row.get("workflow_data")?; + let tags_json: String = row.get("tags")?; + + let workflow_data = serde_json::from_str(&workflow_data_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "workflow_data".to_string(), rusqlite::types::Type::Text))?; + let tags = serde_json::from_str(&tags_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "tags".to_string(), rusqlite::types::Type::Text))?; + + Ok(WorkflowModel { + id: row.get("id")?, + name: row.get("name")?, + description: row.get("description")?, + workflow_data, + version: row.get("version")?, + created_at: row.get("created_at")?, + updated_at: row.get("updated_at")?, + enabled: row.get("enabled")?, + tags, + category: row.get("category")?, + }) + } + + // ==================== 模板相关操作 ==================== + + /// 创建模板 + pub fn create_template(&self, template: &TemplateModel) -> Result<()> { + let conn = self.get_connection()?; + + let template_data_json = serde_json::to_string(&template.template_data)?; + let parameter_schema_json = serde_json::to_string(&template.parameter_schema)?; + let tags_json = serde_json::to_string(&template.tags)?; + + conn.execute( + "INSERT INTO comfyui_templates (id, name, category, description, template_data, parameter_schema, created_at, updated_at, enabled, tags, author, version) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)", + params![ + template.id, + template.name, + template.category, + template.description, + template_data_json, + parameter_schema_json, + template.created_at, + template.updated_at, + template.enabled, + tags_json, + template.author, + template.version + ], + )?; + + debug!("创建模板: {}", template.id); + Ok(()) + } + + /// 获取模板 + pub fn get_template(&self, id: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, name, category, description, template_data, parameter_schema, created_at, updated_at, enabled, tags, author, version + FROM comfyui_templates WHERE id = ?1" + )?; + + let template_iter = stmt.query_map([id], |row| { + self.row_to_template_model(row) + })?; + + for template in template_iter { + return Ok(Some(template?)); + } + + Ok(None) + } + + /// 获取所有模板 + pub fn list_templates(&self, enabled_only: bool) -> Result> { + let conn = self.get_connection()?; + + let sql = if enabled_only { + "SELECT id, name, category, description, template_data, parameter_schema, created_at, updated_at, enabled, tags, author, version + FROM comfyui_templates WHERE enabled = 1 ORDER BY created_at DESC" + } else { + "SELECT id, name, category, description, template_data, parameter_schema, created_at, updated_at, enabled, tags, author, version + FROM comfyui_templates ORDER BY created_at DESC" + }; + + let mut stmt = conn.prepare(sql)?; + let template_iter = stmt.query_map([], |row| { + self.row_to_template_model(row) + })?; + + let mut templates = Vec::new(); + for template in template_iter { + templates.push(template?); + } + + Ok(templates) + } + + /// 更新模板 + pub fn update_template(&self, template: &TemplateModel) -> Result<()> { + let conn = self.get_connection()?; + + let template_data_json = serde_json::to_string(&template.template_data)?; + let parameter_schema_json = serde_json::to_string(&template.parameter_schema)?; + let tags_json = serde_json::to_string(&template.tags)?; + + let rows_affected = conn.execute( + "UPDATE comfyui_templates + SET name = ?2, category = ?3, description = ?4, template_data = ?5, parameter_schema = ?6, updated_at = ?7, enabled = ?8, tags = ?9, author = ?10, version = ?11 + WHERE id = ?1", + params![ + template.id, + template.name, + template.category, + template.description, + template_data_json, + parameter_schema_json, + template.updated_at, + template.enabled, + tags_json, + template.author, + template.version + ], + )?; + + if rows_affected == 0 { + return Err(anyhow!("模板不存在: {}", template.id)); + } + + debug!("更新模板: {}", template.id); + Ok(()) + } + + /// 删除模板 + pub fn delete_template(&self, id: &str) -> Result<()> { + let conn = self.get_connection()?; + + let rows_affected = conn.execute("DELETE FROM comfyui_templates WHERE id = ?1", [id])?; + + if rows_affected == 0 { + return Err(anyhow!("模板不存在: {}", id)); + } + + debug!("删除模板: {}", id); + Ok(()) + } + + /// 按分类获取模板 + pub fn get_templates_by_category(&self, category: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, name, category, description, template_data, parameter_schema, created_at, updated_at, enabled, tags, author, version + FROM comfyui_templates WHERE category = ?1 AND enabled = 1 ORDER BY created_at DESC" + )?; + + let template_iter = stmt.query_map([category], |row| { + self.row_to_template_model(row) + })?; + + let mut templates = Vec::new(); + for template in template_iter { + templates.push(template?); + } + + Ok(templates) + } + + /// 将数据库行转换为模板模型 + fn row_to_template_model(&self, row: &Row) -> rusqlite::Result { + let template_data_json: String = row.get("template_data")?; + let parameter_schema_json: String = row.get("parameter_schema")?; + let tags_json: String = row.get("tags")?; + + let template_data = serde_json::from_str(&template_data_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "template_data".to_string(), rusqlite::types::Type::Text))?; + let parameter_schema = serde_json::from_str(¶meter_schema_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "parameter_schema".to_string(), rusqlite::types::Type::Text))?; + let tags = serde_json::from_str(&tags_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "tags".to_string(), rusqlite::types::Type::Text))?; + + Ok(TemplateModel { + id: row.get("id")?, + name: row.get("name")?, + category: row.get("category")?, + description: row.get("description")?, + template_data, + parameter_schema, + created_at: row.get("created_at")?, + updated_at: row.get("updated_at")?, + enabled: row.get("enabled")?, + tags, + author: row.get("author")?, + version: row.get("version")?, + }) + } + + // ==================== 执行记录相关操作 ==================== + + /// 创建执行记录 + pub fn create_execution(&self, execution: &ExecutionModel) -> Result<()> { + let conn = self.get_connection()?; + + let parameters_json = execution.parameters.as_ref() + .map(|p| serde_json::to_string(p)) + .transpose()?; + let results_json = execution.results.as_ref() + .map(|r| serde_json::to_string(r)) + .transpose()?; + let output_urls_json = serde_json::to_string(&execution.output_urls)?; + let node_outputs_json = execution.node_outputs.as_ref() + .map(|n| serde_json::to_string(n)) + .transpose()?; + + conn.execute( + "INSERT INTO comfyui_executions (id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + params![ + execution.id, + execution.workflow_id, + execution.template_id, + execution.prompt_id, + execution.status.to_string(), + parameters_json, + results_json, + output_urls_json, + execution.error_message, + execution.execution_time, + execution.created_at, + execution.completed_at, + execution.client_id, + node_outputs_json + ], + )?; + + debug!("创建执行记录: {}", execution.id); + Ok(()) + } + + /// 获取执行记录 + pub fn get_execution(&self, id: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions WHERE id = ?1" + )?; + + let execution_iter = stmt.query_map([id], |row| { + self.row_to_execution_model(row) + })?; + + for execution in execution_iter { + return Ok(Some(execution?)); + } + + Ok(None) + } + + /// 通过 prompt_id 获取执行记录 + pub fn get_execution_by_prompt_id(&self, prompt_id: &str) -> Result> { + let conn = self.get_connection()?; + + let mut stmt = conn.prepare( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions WHERE prompt_id = ?1" + )?; + + let execution_iter = stmt.query_map([prompt_id], |row| { + self.row_to_execution_model(row) + })?; + + for execution in execution_iter { + return Ok(Some(execution?)); + } + + Ok(None) + } + + /// 获取执行记录列表 + pub fn list_executions(&self, limit: Option, status_filter: Option) -> Result> { + let conn = self.get_connection()?; + + let (sql, params): (String, Vec) = match (limit, status_filter) { + (Some(limit), Some(status)) => ( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions WHERE status = ?1 ORDER BY created_at DESC LIMIT ?2".to_string(), + vec![status.to_string(), limit.to_string()] + ), + (Some(limit), None) => ( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions ORDER BY created_at DESC LIMIT ?1".to_string(), + vec![limit.to_string()] + ), + (None, Some(status)) => ( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions WHERE status = ?1 ORDER BY created_at DESC".to_string(), + vec![status.to_string()] + ), + (None, None) => ( + "SELECT id, workflow_id, template_id, prompt_id, status, parameters, results, output_urls, error_message, execution_time, created_at, completed_at, client_id, node_outputs + FROM comfyui_executions ORDER BY created_at DESC".to_string(), + vec![] + ), + }; + + let mut stmt = conn.prepare(&sql)?; + let execution_iter = stmt.query_map(rusqlite::params_from_iter(params.iter()), |row| { + self.row_to_execution_model(row) + })?; + + let mut executions = Vec::new(); + for execution in execution_iter { + executions.push(execution?); + } + + Ok(executions) + } + + /// 更新执行记录 + pub fn update_execution(&self, execution: &ExecutionModel) -> Result<()> { + let conn = self.get_connection()?; + + let parameters_json = execution.parameters.as_ref() + .map(|p| serde_json::to_string(p)) + .transpose()?; + let results_json = execution.results.as_ref() + .map(|r| serde_json::to_string(r)) + .transpose()?; + let output_urls_json = serde_json::to_string(&execution.output_urls)?; + let node_outputs_json = execution.node_outputs.as_ref() + .map(|n| serde_json::to_string(n)) + .transpose()?; + + let rows_affected = conn.execute( + "UPDATE comfyui_executions + SET workflow_id = ?2, template_id = ?3, status = ?4, parameters = ?5, results = ?6, output_urls = ?7, error_message = ?8, execution_time = ?9, completed_at = ?10, client_id = ?11, node_outputs = ?12 + WHERE id = ?1", + params![ + execution.id, + execution.workflow_id, + execution.template_id, + execution.status.to_string(), + parameters_json, + results_json, + output_urls_json, + execution.error_message, + execution.execution_time, + execution.completed_at, + execution.client_id, + node_outputs_json + ], + )?; + + if rows_affected == 0 { + return Err(anyhow!("执行记录不存在: {}", execution.id)); + } + + debug!("更新执行记录: {}", execution.id); + Ok(()) + } + + /// 删除执行记录 + pub fn delete_execution(&self, id: &str) -> Result<()> { + let conn = self.get_connection()?; + + let rows_affected = conn.execute("DELETE FROM comfyui_executions WHERE id = ?1", [id])?; + + if rows_affected == 0 { + return Err(anyhow!("执行记录不存在: {}", id)); + } + + debug!("删除执行记录: {}", id); + Ok(()) + } + + /// 将数据库行转换为执行记录模型 + fn row_to_execution_model(&self, row: &Row) -> rusqlite::Result { + let parameters_json: Option = row.get("parameters")?; + let results_json: Option = row.get("results")?; + let output_urls_json: String = row.get("output_urls")?; + let node_outputs_json: Option = row.get("node_outputs")?; + let status_str: String = row.get("status")?; + + let parameters = parameters_json + .map(|json| serde_json::from_str(&json)) + .transpose() + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "parameters".to_string(), rusqlite::types::Type::Text))?; + + let results = results_json + .map(|json| serde_json::from_str(&json)) + .transpose() + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "results".to_string(), rusqlite::types::Type::Text))?; + + let output_urls = serde_json::from_str(&output_urls_json) + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "output_urls".to_string(), rusqlite::types::Type::Text))?; + + let node_outputs = node_outputs_json + .map(|json| serde_json::from_str(&json)) + .transpose() + .map_err(|e| rusqlite::Error::InvalidColumnType(0, "node_outputs".to_string(), rusqlite::types::Type::Text))?; + + let status = match status_str.as_str() { + "pending" => ExecutionStatus::Pending, + "running" => ExecutionStatus::Running, + "completed" => ExecutionStatus::Completed, + "failed" => ExecutionStatus::Failed, + "cancelled" => ExecutionStatus::Cancelled, + _ => ExecutionStatus::Failed, // 默认为失败状态 + }; + + Ok(ExecutionModel { + id: row.get("id")?, + workflow_id: row.get("workflow_id")?, + template_id: row.get("template_id")?, + prompt_id: row.get("prompt_id")?, + status, + parameters, + results, + output_urls, + error_message: row.get("error_message")?, + execution_time: row.get("execution_time")?, + created_at: row.get("created_at")?, + completed_at: row.get("completed_at")?, + client_id: row.get("client_id")?, + node_outputs, + }) + } +} diff --git a/apps/desktop/src-tauri/src/data/repositories/mod.rs b/apps/desktop/src-tauri/src/data/repositories/mod.rs index 6984b5e..8a6b427 100644 --- a/apps/desktop/src-tauri/src/data/repositories/mod.rs +++ b/apps/desktop/src-tauri/src/data/repositories/mod.rs @@ -20,6 +20,7 @@ pub mod outfit_image_repository; pub mod outfit_photo_generation_repository; pub mod video_generation_record_repository; pub mod hedra_lipsync_repository; +pub mod comfyui_repository; // Multi-workflow system repositories pub mod workflow_template_repository;