diff --git a/apps/desktop/src-tauri/Cargo.toml b/apps/desktop/src-tauri/Cargo.toml index 3816a62..22b6a73 100644 --- a/apps/desktop/src-tauri/Cargo.toml +++ b/apps/desktop/src-tauri/Cargo.toml @@ -59,6 +59,7 @@ urlencoding = "2.1" bincode = "1.3" zip = "0.6" sysinfo = "0.30" +comfyui-sdk = "0.1" [target.'cfg(windows)'.dependencies] winapi = { version = "0.3", features = ["sysinfoapi"] } diff --git a/apps/desktop/src-tauri/src/business/services/comfyui_sdk_service.rs b/apps/desktop/src-tauri/src/business/services/comfyui_sdk_service.rs new file mode 100644 index 0000000..8e07683 --- /dev/null +++ b/apps/desktop/src-tauri/src/business/services/comfyui_sdk_service.rs @@ -0,0 +1,281 @@ +use anyhow::{Result, anyhow}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; +use std::time::Duration; +use tracing::{info, warn, error, debug}; +use tokio::time::sleep; + +use crate::config::ComfyUISettings; +use crate::data::models::outfit_photo_generation::{ + WorkflowProgress, WorkflowNodeReplacement +}; + +/// ComfyUI SDK 服务包装器 +/// 使用 comfyui-sdk crate 提供更好的 ComfyUI 集成 +pub struct ComfyUISDKService { + client: comfyui_sdk::ComfyUIClient, + settings: ComfyUISettings, +} + +/// 工作流执行结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WorkflowExecutionResult { + pub prompt_id: String, + pub outputs: HashMap>, + pub execution_time: Duration, + pub node_outputs: HashMap, +} + +/// 工作流执行配置 +#[derive(Debug, Clone)] +pub struct SDKExecutionConfig { + pub timeout: Duration, + pub retry_attempts: u32, + pub progress_callback: Option>, +} + +impl Default for SDKExecutionConfig { + fn default() -> Self { + Self { + timeout: Duration::from_secs(300), // 5分钟默认超时 + retry_attempts: 3, + progress_callback: None, + } + } +} + +impl ComfyUISDKService { + /// 创建新的 ComfyUI SDK 服务实例 + pub fn new(settings: ComfyUISettings) -> Result { + let client = comfyui_sdk::ComfyUIClient::new(&settings.base_url)?; + + Ok(Self { + client, + settings, + }) + } + + /// 检查 ComfyUI 服务状态 + pub async fn check_health(&self) -> Result { + match self.client.get_system_stats().await { + Ok(_) => { + info!("ComfyUI 服务健康检查通过"); + Ok(true) + } + Err(e) => { + warn!("ComfyUI 服务健康检查失败: {}", e); + Ok(false) + } + } + } + + /// 获取队列状态 + pub async fn get_queue_status(&self) -> Result { + let queue_info = self.client.get_queue().await?; + + Ok(QueueStatus { + running: queue_info.queue_running.len(), + pending: queue_info.queue_pending.len(), + }) + } + + /// 执行工作流 + pub async fn execute_workflow( + &self, + workflow: Value, + config: SDKExecutionConfig, + ) -> Result { + let start_time = std::time::Instant::now(); + + // 提交工作流到队列 + let prompt_id = self.client.queue_prompt(workflow).await?; + info!("工作流已提交，prompt_id: {}", prompt_id); + + // 等待执行完成 + let result = self.wait_for_completion(&prompt_id, config).await?; + + let execution_time = start_time.elapsed(); + info!("工作流执行完成，耗时: {:?}", execution_time); + + Ok(WorkflowExecutionResult { + prompt_id, + outputs: result.outputs, + execution_time, + node_outputs: result.node_outputs, + }) + } + + /// 等待工作流执行完成 + async fn wait_for_completion( + &self, + prompt_id: &str, + config: SDKExecutionConfig, + ) -> Result { + let mut attempts = 0; + let max_attempts = config.retry_attempts; + + loop { + match self.check_execution_status(prompt_id).await { + Ok(status) => { + match status { + ExecutionStatus::Completed(result) => return Ok(result), + ExecutionStatus::Running(progress) => { + if let Some(callback) = &config.progress_callback { + callback(progress); + } + sleep(Duration::from_secs(1)).await; + } + ExecutionStatus::Failed(error) => { + return Err(anyhow!("工作流执行失败: {}", error)); + } + } + } + Err(e) => { + attempts += 1; + if attempts >= max_attempts { + return Err(anyhow!("检查执行状态失败，已重试 {} 次: {}", attempts, e)); + } + warn!("检查执行状态失败，重试中 ({}/{}): {}", attempts, max_attempts, e); + sleep(Duration::from_secs(2)).await; + } + } + } + } + + /// 检查执行状态 + async fn check_execution_status(&self, prompt_id: &str) -> Result { + // 检查历史记录 + let history = self.client.get_history(Some(prompt_id)).await?; + + if let Some(entry) = history.get(prompt_id) { + // 执行完成 + let outputs = self.extract_outputs(entry)?; + let node_outputs = self.extract_node_outputs(entry)?; + + return Ok(ExecutionStatus::Completed(ExecutionResult { + outputs, + node_outputs, + })); + } + + // 检查队列状态 + let queue = self.client.get_queue().await?; + + // 检查是否在运行队列中 + for item in &queue.queue_running { + if item.get("prompt_id").and_then(|v| v.as_str()) == Some(prompt_id) { + let progress = self.extract_progress(item)?; + return Ok(ExecutionStatus::Running(progress)); + } + } + + // 检查是否在等待队列中 + for item in &queue.queue_pending { + if item.get("prompt_id").and_then(|v| v.as_str()) == Some(prompt_id) { + return Ok(ExecutionStatus::Running(WorkflowProgress { + current_node: "等待中".to_string(), + progress: 0.0, + message: "工作流在队列中等待执行".to_string(), + })); + } + } + + Err(anyhow!("未找到 prompt_id: {}", prompt_id)) + } + + /// 从历史记录中提取输出 + fn extract_outputs(&self, history_entry: &Value) -> Result>> { + let mut outputs = HashMap::new(); + + if let Some(outputs_obj) = history_entry.get("outputs") { + for (node_id, node_output) in outputs_obj.as_object().unwrap_or(&serde_json::Map::new()) { + if let Some(images) = node_output.get("images") { + if let Some(images_array) = images.as_array() { + let image_urls: Vec = images_array + .iter() + .filter_map(|img| { + img.get("filename") + .and_then(|f| f.as_str()) + .map(|filename| { + format!("{}/view?filename={}", self.settings.base_url, filename) + }) + }) + .collect(); + + if !image_urls.is_empty() { + outputs.insert(node_id.clone(), image_urls); + } + } + } + } + } + + Ok(outputs) + } + + /// 从历史记录中提取节点输出 + fn extract_node_outputs(&self, history_entry: &Value) -> Result> { + let mut node_outputs = HashMap::new(); + + if let Some(outputs_obj) = history_entry.get("outputs") { + for (node_id, node_output) in outputs_obj.as_object().unwrap_or(&serde_json::Map::new()) { + node_outputs.insert(node_id.clone(), node_output.clone()); + } + } + + Ok(node_outputs) + } + + /// 从队列项中提取进度信息 + fn extract_progress(&self, queue_item: &Value) -> Result { + // 这里需要根据实际的队列项结构来提取进度信息 + // 由于 comfyui-sdk 的具体实现可能不同，这里提供一个基本的实现 + + Ok(WorkflowProgress { + current_node: "执行中".to_string(), + progress: 50.0, // 默认进度 + message: "工作流正在执行中".to_string(), + }) + } + + /// 取消工作流执行 + pub async fn cancel_workflow(&self, prompt_id: &str) -> Result<()> { + self.client.interrupt().await?; + info!("已取消工作流执行: {}", prompt_id); + Ok(()) + } + + /// 获取可用模型列表 + pub async fn get_available_models(&self) -> Result> { + let object_info = self.client.get_object_info().await?; + let mut models = Vec::new(); + + // 从 object_info 中提取模型信息 + // 这里需要根据实际的 API 响应结构来实现 + + Ok(models) + } +} + +/// 队列状态 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueStatus { + pub running: usize, + pub pending: usize, +} + +/// 执行状态 +#[derive(Debug)] +enum ExecutionStatus { + Running(WorkflowProgress), + Completed(ExecutionResult), + Failed(String), +} + +/// 执行结果 +#[derive(Debug)] +struct ExecutionResult { + outputs: HashMap>, + node_outputs: HashMap, +} diff --git a/cargos/comfyui-sdk/client/comfyui_client.rs b/cargos/comfyui-sdk/client/comfyui_client.rs new file mode 100644 index 0000000..15a64de --- /dev/null +++ b/cargos/comfyui-sdk/client/comfyui_client.rs @@ -0,0 +1,274 @@ +//! Main ComfyUI client that combines HTTP and WebSocket functionality + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::time::sleep; + +use crate::types::{ + ComfyUIClientConfig, PromptRequest, TemplateExecutionResult, + ExecutionOptions, ExecutionCallbacks, ExecutionError, ParameterValues +}; +use crate::templates::{TemplateManager, WorkflowTemplate, WorkflowInstance}; +use crate::client::{HTTPClient, WebSocketClient}; +use crate::error::{ComfyUIError, Result}; + + +/// Main ComfyUI client +pub struct ComfyUIClient { + http_client: HTTPClient, + ws_client: WebSocketClient, + template_manager: TemplateManager, + config: ComfyUIClientConfig, +} + +impl ComfyUIClient { + /// Creates a new ComfyUI client + pub fn new(config: ComfyUIClientConfig) -> Result { + let http_client = HTTPClient::new(config.clone())?; + let ws_client = WebSocketClient::new(config.clone()); + let template_manager = TemplateManager::new(); + + Ok(Self { + http_client, + ws_client, + template_manager, + config, + }) + } + + /// Gets the HTTP client + pub fn http(&self) -> &HTTPClient { + &self.http_client + } + + /// Gets the WebSocket client + pub fn ws(&self) -> &WebSocketClient { + &self.ws_client + } + + /// Gets the template manager + pub fn templates(&mut self) -> &mut TemplateManager { + &mut self.template_manager + } + + /// Gets the template manager (read-only) + pub fn templates_ref(&self) -> &TemplateManager { + &self.template_manager + } + + /// Gets the client configuration + pub fn config(&self) -> &ComfyUIClientConfig { + &self.config + } + + /// Connects to ComfyUI server (both HTTP and WebSocket) + pub async fn connect(&mut self) -> Result<()> { + // Test HTTP connection first + self.http_client.get_queue().await + .map_err(|e| ComfyUIError::connection(format!("HTTP connection failed: {}", e)))?; + + // Connect WebSocket + self.ws_client.connect().await + .map_err(|e| ComfyUIError::connection(format!("WebSocket connection failed: {}", e)))?; + + log::info!("Successfully connected to ComfyUI server"); + Ok(()) + } + + /// Disconnects from ComfyUI server + pub async fn disconnect(&mut self) -> Result<()> { + self.ws_client.disconnect().await?; + log::info!("Disconnected from ComfyUI server"); + Ok(()) + } + + /// Checks if connected to the server + pub async fn is_connected(&self) -> bool { + self.ws_client.is_connected().await + } + + /// Executes a workflow template + pub async fn execute_template( + &self, + template: &WorkflowTemplate, + parameters: ParameterValues, + options: ExecutionOptions, + ) -> Result { + let start_time = Instant::now(); + + // Create workflow instance + let instance = template.create_instance(parameters)?; + + // Convert to prompt request + let prompt_request = self.instance_to_prompt_request(&instance)?; + + // Submit prompt + let prompt_response = self.http_client.queue_prompt(&prompt_request).await?; + let prompt_id = prompt_response.prompt_id.clone(); + + log::info!("Submitted prompt with ID: {}", prompt_id); + + // Wait for completion if WebSocket is connected + if self.ws_client.is_connected().await { + match self.wait_for_completion(&prompt_id, options.timeout).await { + Ok(outputs) => { + let execution_time = start_time.elapsed().as_millis() as u64; + Ok(TemplateExecutionResult::success(prompt_id, outputs, execution_time)) + } + Err(error) => { + let execution_time = start_time.elapsed().as_millis() as u64; + let exec_error = ExecutionError { + node_id: None, + message: error.to_string(), + details: None, + timestamp: chrono::Utc::now(), + }; + Ok(TemplateExecutionResult::failure(prompt_id, exec_error, execution_time)) + } + } + } else { + // If no WebSocket, just return the prompt ID + let execution_time = start_time.elapsed().as_millis() as u64; + Ok(TemplateExecutionResult::success(prompt_id, HashMap::new(), execution_time)) + } + } + + /// Executes a template with callbacks + pub async fn execute_template_with_callbacks( + &self, + template: &WorkflowTemplate, + parameters: ParameterValues, + options: ExecutionOptions, + callbacks: Arc, + ) -> Result { + // Register callbacks + let callback_id = self.ws_client.register_callbacks(callbacks).await; + + // Execute template + let result = self.execute_template(template, parameters, options).await; + + // Unregister callbacks + self.ws_client.unregister_callbacks(&callback_id).await; + + result + } + + /// Converts a workflow instance to a prompt request + fn instance_to_prompt_request(&self, instance: &WorkflowInstance) -> Result { + let workflow_json = instance.to_workflow_json()?; + + let mut prompt = HashMap::new(); + if let serde_json::Value::Object(workflow_obj) = workflow_json { + for (key, value) in workflow_obj { + prompt.insert(key, value); + } + } + + Ok(PromptRequest { + prompt, + client_id: Some(self.ws_client.client_id().to_string()), + extra_data: None, + }) + } + + /// Waits for prompt completion + async fn wait_for_completion( + &self, + prompt_id: &str, + timeout: Option, + ) -> Result> { + let timeout_duration = timeout.unwrap_or(Duration::from_secs(300)); // 5 minutes default + let start_time = Instant::now(); + let check_interval = Duration::from_millis(1000); // Check every second + + loop { + // Check timeout + if start_time.elapsed() > timeout_duration { + return Err(ComfyUIError::timeout(format!( + "Execution timed out after {:?}", + timeout_duration + ))); + } + + // Check history for completion + match self.http_client.get_history_by_prompt(prompt_id).await { + Ok(history) => { + if let Some(history_item) = history.get(prompt_id) { + if history_item.status.completed { + // Extract outputs + let mut outputs = HashMap::new(); + for (node_id, node_outputs) in &history_item.outputs { + let mut node_output = HashMap::new(); + for (output_name, output_list) in node_outputs { + node_output.insert(output_name.clone(), serde_json::to_value(output_list)?); + } + outputs.insert(node_id.clone(), serde_json::to_value(node_output)?); + } + return Ok(outputs); + } + } + } + Err(e) => { + log::warn!("Error checking history: {}", e); + } + } + + sleep(check_interval).await; + } + } + + /// Converts outputs to HTTP URLs + pub fn outputs_to_urls(&self, outputs: &HashMap) -> Vec { + self.http_client.outputs_to_urls(outputs) + } + + /// Interrupts the current execution + pub async fn interrupt(&self) -> Result<()> { + self.http_client.interrupt().await + } + + /// Clears the execution queue + pub async fn clear_queue(&self) -> Result<()> { + self.http_client.clear_queue().await + } + + /// Gets the current queue status + pub async fn get_queue_status(&self) -> Result { + self.http_client.get_queue().await + } + + /// Gets system statistics + pub async fn get_system_stats(&self) -> Result { + self.http_client.get_system_stats().await + } + + /// Gets available nodes information + pub async fn get_object_info(&self) -> Result { + self.http_client.get_object_info().await + } + + /// Uploads an image file + pub async fn upload_image>( + &self, + image_path: P, + overwrite: bool, + ) -> Result { + self.http_client.upload_image(image_path, overwrite).await + } + + /// Gets an image by filename + pub async fn get_image( + &self, + filename: &str, + subfolder: Option<&str>, + image_type: Option<&str>, + ) -> Result { + self.http_client.get_image(filename, subfolder, image_type).await + } + + /// Frees memory + pub async fn free_memory(&self) -> Result<()> { + self.http_client.free_memory().await + } +} diff --git a/cargos/comfyui-sdk/client/http_client.rs b/cargos/comfyui-sdk/client/http_client.rs new file mode 100644 index 0000000..3000ff7 --- /dev/null +++ b/cargos/comfyui-sdk/client/http_client.rs @@ -0,0 +1,324 @@ +//! HTTP client for ComfyUI API + +use std::collections::HashMap; +use std::path::Path; +use reqwest::{Client, multipart}; +use url::Url; +use crate::types::{ + ComfyUIClientConfig, PromptRequest, PromptResponse, QueueStatus, + HistoryItem, SystemStats, ObjectInfo, UploadImageResponse, ViewMetadata +}; +use crate::error::{ComfyUIError, Result}; +use crate::utils::error_handling::{retry_if_retryable, RetryConfig, with_timeout}; + +/// HTTP client for ComfyUI API +#[derive(Debug, Clone)] +pub struct HTTPClient { + client: Client, + base_url: Url, + config: ComfyUIClientConfig, +} + +impl HTTPClient { + /// Creates a new HTTP client + pub fn new(config: ComfyUIClientConfig) -> Result { + let base_url = Url::parse(&config.base_url)?; + + let mut client_builder = Client::builder(); + + if let Some(timeout) = config.timeout { + client_builder = client_builder.timeout(timeout); + } + + let client = client_builder.build()?; + + Ok(Self { + client, + base_url, + config, + }) + } + + /// Gets the base URL + pub fn base_url(&self) -> &Url { + &self.base_url + } + + /// Builds a URL for an endpoint + fn build_url(&self, endpoint: &str) -> Result { + self.base_url.join(endpoint) + .map_err(ComfyUIError::from) + } + + /// Executes a GET request with retry logic + async fn get_with_retry(&self, endpoint: &str) -> Result + where + T: serde::de::DeserializeOwned, + { + let url = self.build_url(endpoint)?; + let retry_config = RetryConfig { + max_attempts: self.config.retry_attempts.unwrap_or(3), + initial_delay: self.config.retry_delay.unwrap_or(std::time::Duration::from_millis(1000)), + ..Default::default() + }; + + let operation = || async { + let mut request = self.client.get(url.clone()); + + if let Some(headers) = &self.config.headers { + for (key, value) in headers { + request = request.header(key, value); + } + } + + let response = request.send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + let data: T = response.json().await?; + Ok(data) + }; + + if let Some(timeout) = self.config.timeout { + with_timeout(retry_if_retryable(operation, retry_config), timeout).await + } else { + retry_if_retryable(operation, retry_config).await + } + } + + /// Executes a POST request with retry logic + async fn post_with_retry(&self, endpoint: &str, body: &B) -> Result + where + T: serde::de::DeserializeOwned, + B: serde::Serialize, + { + let url = self.build_url(endpoint)?; + let retry_config = RetryConfig { + max_attempts: self.config.retry_attempts.unwrap_or(3), + initial_delay: self.config.retry_delay.unwrap_or(std::time::Duration::from_millis(1000)), + ..Default::default() + }; + + let operation = || async { + let mut request = self.client.post(url.clone()).json(body); + + if let Some(headers) = &self.config.headers { + for (key, value) in headers { + request = request.header(key, value); + } + } + + let response = request.send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + let data: T = response.json().await?; + Ok(data) + }; + + if let Some(timeout) = self.config.timeout { + with_timeout(retry_if_retryable(operation, retry_config), timeout).await + } else { + retry_if_retryable(operation, retry_config).await + } + } + + /// Gets the current queue status + pub async fn get_queue(&self) -> Result { + self.get_with_retry("/queue").await + } + + /// Submits a prompt for execution + pub async fn queue_prompt(&self, prompt_request: &PromptRequest) -> Result { + self.post_with_retry("/prompt", prompt_request).await + } + + /// Gets execution history + pub async fn get_history(&self, max_items: Option) -> Result> { + let endpoint = if let Some(max) = max_items { + format!("/history?max_items={}", max) + } else { + "/history".to_string() + }; + + self.get_with_retry(&endpoint).await + } + + /// Gets history for a specific prompt + pub async fn get_history_by_prompt(&self, prompt_id: &str) -> Result> { + let endpoint = format!("/history/{}", prompt_id); + self.get_with_retry(&endpoint).await + } + + /// Clears the execution queue + pub async fn clear_queue(&self) -> Result<()> { + let url = self.build_url("/queue")?; + let response = self.client.delete(url).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + Ok(()) + } + + /// Cancels a specific prompt + pub async fn cancel_prompt(&self, prompt_id: &str) -> Result<()> { + let body = serde_json::json!({ "delete": [prompt_id] }); + let url = self.build_url("/queue")?; + let response = self.client.post(url).json(&body).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + Ok(()) + } + + /// Gets system statistics + pub async fn get_system_stats(&self) -> Result { + self.get_with_retry("/system_stats").await + } + + /// Gets object information (available nodes) + pub async fn get_object_info(&self) -> Result { + self.get_with_retry("/object_info").await + } + + /// Uploads an image file + pub async fn upload_image>(&self, image_path: P, overwrite: bool) -> Result { + let path = image_path.as_ref(); + let filename = path.file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| ComfyUIError::new("Invalid filename"))?; + + let file_bytes = tokio::fs::read(path).await?; + + let part = multipart::Part::bytes(file_bytes) + .file_name(filename.to_string()) + .mime_str("image/*") + .map_err(|e| ComfyUIError::new(format!("Invalid MIME type: {}", e)))?; + + let form = multipart::Form::new() + .part("image", part) + .text("overwrite", overwrite.to_string()); + + let url = self.build_url("/upload/image")?; + let response = self.client.post(url).multipart(form).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + let upload_response: UploadImageResponse = response.json().await?; + Ok(upload_response) + } + + /// Gets an image by filename + pub async fn get_image(&self, filename: &str, subfolder: Option<&str>, image_type: Option<&str>) -> Result { + let mut endpoint = format!("/view?filename={}", filename); + + if let Some(subfolder) = subfolder { + endpoint.push_str(&format!("&subfolder={}", subfolder)); + } + + if let Some(image_type) = image_type { + endpoint.push_str(&format!("&type={}", image_type)); + } + + let url = self.build_url(&endpoint)?; + let response = self.client.get(url).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + let bytes = response.bytes().await?; + Ok(bytes) + } + + /// Gets image metadata + pub async fn get_view_metadata(&self, filename: &str, subfolder: Option<&str>) -> Result { + let mut endpoint = format!("/view_metadata/{}", filename); + + if let Some(subfolder) = subfolder { + endpoint.push_str(&format!("?subfolder={}", subfolder)); + } + + self.get_with_retry(&endpoint).await + } + + /// Converts outputs to HTTP URLs + pub fn outputs_to_urls(&self, outputs: &HashMap) -> Vec { + let mut urls = Vec::new(); + + for (_node_id, output) in outputs { + if let Some(output_obj) = output.as_object() { + if let Some(images) = output_obj.get("images").and_then(|v| v.as_array()) { + for image in images { + if let Some(image_obj) = image.as_object() { + if let (Some(filename), Some(subfolder), Some(image_type)) = ( + image_obj.get("filename").and_then(|v| v.as_str()), + image_obj.get("subfolder").and_then(|v| v.as_str()), + image_obj.get("type").and_then(|v| v.as_str()), + ) { + let base_url_str = self.base_url.as_str().trim_end_matches('/'); + let url = format!( + "{}/view?filename={}&subfolder={}&type={}", + base_url_str, filename, subfolder, image_type + ); + urls.push(url); + } + } + } + } + } + } + + urls + } + + /// Interrupts the current execution + pub async fn interrupt(&self) -> Result<()> { + let url = self.build_url("/interrupt")?; + let response = self.client.post(url).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + Ok(()) + } + + /// Frees memory + pub async fn free_memory(&self) -> Result<()> { + let url = self.build_url("/free")?; + let response = self.client.post(url).send().await?; + + if !response.status().is_success() { + return Err(ComfyUIError::Http(reqwest::Error::from( + response.error_for_status().unwrap_err() + ))); + } + + Ok(()) + } +} diff --git a/cargos/comfyui-sdk/client/mod.rs b/cargos/comfyui-sdk/client/mod.rs new file mode 100644 index 0000000..0c799ad --- /dev/null +++ b/cargos/comfyui-sdk/client/mod.rs @@ -0,0 +1,10 @@ +//! Client modules for ComfyUI SDK + +pub mod http_client; +pub mod websocket_client; +pub mod comfyui_client; + +// Re-export main types +pub use http_client::HTTPClient; +pub use websocket_client::WebSocketClient; +pub use comfyui_client::ComfyUIClient; diff --git a/cargos/comfyui-sdk/client/websocket_client.rs b/cargos/comfyui-sdk/client/websocket_client.rs new file mode 100644 index 0000000..59e39ea --- /dev/null +++ b/cargos/comfyui-sdk/client/websocket_client.rs @@ -0,0 +1,226 @@ +//! WebSocket client for ComfyUI real-time events + +use std::sync::Arc; +use tokio::sync::{RwLock, mpsc}; +use tokio_tungstenite::{connect_async, tungstenite::Message}; +use futures_util::StreamExt; +use url::Url; +use uuid::Uuid; +use chrono::Utc; + +use crate::types::{ + ComfyUIClientConfig, WSMessage, ExecutionProgress, ExecutionResult, + ExecutionError, ExecutionCallbacks +}; +use crate::error::{ComfyUIError, Result}; +use crate::utils::event_emitter::EventEmitter; + +/// WebSocket client for real-time ComfyUI events +pub struct WebSocketClient { + config: ComfyUIClientConfig, + client_id: String, + event_emitter: EventEmitter, + is_connected: Arc>, + shutdown_tx: Option>, +} + +impl WebSocketClient { + /// Creates a new WebSocket client + pub fn new(config: ComfyUIClientConfig) -> Self { + Self { + config, + client_id: Uuid::new_v4().to_string(), + event_emitter: EventEmitter::new(), + is_connected: Arc::new(RwLock::new(false)), + shutdown_tx: None, + } + } + + /// Gets the client ID + pub fn client_id(&self) -> &str { + &self.client_id + } + + /// Checks if connected + pub async fn is_connected(&self) -> bool { + *self.is_connected.read().await + } + + /// Connects to the WebSocket server + pub async fn connect(&mut self) -> Result<()> { + if self.is_connected().await { + return Ok(()); + } + + let ws_url = self.build_websocket_url()?; + log::info!("Connecting to WebSocket: {}", ws_url); + + let (ws_stream, _) = connect_async(&ws_url).await?; + let (_ws_sender, mut ws_receiver) = ws_stream.split(); + + // Set connected state + *self.is_connected.write().await = true; + + // Create shutdown channel + let (shutdown_tx, mut shutdown_rx) = mpsc::unbounded_channel(); + self.shutdown_tx = Some(shutdown_tx); + + // Clone necessary data for the message handling task + let event_emitter = self.event_emitter.clone(); + let is_connected = self.is_connected.clone(); + + // Spawn message handling task + tokio::spawn(async move { + loop { + tokio::select! { + // Handle incoming messages + message = ws_receiver.next() => { + match message { + Some(Ok(Message::Text(text))) => { + if let Err(e) = Self::handle_message(&text, &event_emitter).await { + log::error!("Error handling WebSocket message: {}", e); + } + } + Some(Ok(Message::Close(_))) => { + log::info!("WebSocket connection closed by server"); + break; + } + Some(Err(e)) => { + log::error!("WebSocket error: {}", e); + break; + } + None => { + log::info!("WebSocket stream ended"); + break; + } + _ => {} // Ignore other message types + } + } + // Handle shutdown signal + _ = shutdown_rx.recv() => { + log::info!("WebSocket shutdown requested"); + break; + } + } + } + + // Set disconnected state + *is_connected.write().await = false; + log::info!("WebSocket connection closed"); + }); + + Ok(()) + } + + /// Disconnects from the WebSocket server + pub async fn disconnect(&mut self) -> Result<()> { + if !self.is_connected().await { + return Ok(()); + } + + if let Some(shutdown_tx) = &self.shutdown_tx { + let _ = shutdown_tx.send(()); + } + + // Wait for disconnection + let mut attempts = 0; + while self.is_connected().await && attempts < 50 { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + attempts += 1; + } + + self.shutdown_tx = None; + Ok(()) + } + + /// Registers execution callbacks + pub async fn register_callbacks(&self, callbacks: Arc) -> String { + let callback_id = Uuid::new_v4().to_string(); + self.event_emitter.register_callbacks(callback_id.clone(), callbacks).await; + callback_id + } + + /// Unregisters execution callbacks + pub async fn unregister_callbacks(&self, callback_id: &str) { + self.event_emitter.unregister_callbacks(callback_id).await; + } + + /// Builds the WebSocket URL + fn build_websocket_url(&self) -> Result { + let base_url = Url::parse(&self.config.base_url)?; + let scheme = match base_url.scheme() { + "http" => "ws", + "https" => "wss", + _ => return Err(ComfyUIError::new("Invalid URL scheme")), + }; + + let ws_url = format!( + "{}://{}:{}/ws?clientId={}", + scheme, + base_url.host_str().unwrap_or("localhost"), + base_url.port().unwrap_or(8188), + self.client_id + ); + + Ok(ws_url) + } + + /// Handles incoming WebSocket messages + async fn handle_message(text: &str, event_emitter: &EventEmitter) -> Result<()> { + let message: WSMessage = serde_json::from_str(text) + .map_err(|e| ComfyUIError::new(format!("Failed to parse WebSocket message: {}", e)))?; + + match message { + WSMessage::Progress { data } => { + let progress = ExecutionProgress { + node_id: data.node.unwrap_or_else(|| "unknown".to_string()), + progress: data.value, + max: data.max, + timestamp: Utc::now(), + }; + event_emitter.emit_progress(progress).await; + } + WSMessage::Executing { data } => { + if let Some(node_id) = data.node { + event_emitter.emit_executing(node_id).await; + } + } + WSMessage::Executed { data } => { + let result = ExecutionResult { + prompt_id: data.prompt_id, + outputs: data.output, + execution_time: 0, // WebSocket doesn't provide execution time + timestamp: Utc::now(), + }; + event_emitter.emit_executed(result).await; + } + WSMessage::ExecutionError { data } => { + let error = ExecutionError { + node_id: data.node_id, + message: data.message, + details: data.details, + timestamp: Utc::now(), + }; + event_emitter.emit_error(error).await; + } + WSMessage::Unknown => { + log::debug!("Received unknown WebSocket message: {}", text); + } + } + + Ok(()) + } + + /// Gets the event emitter for direct access + pub fn event_emitter(&self) -> &EventEmitter { + &self.event_emitter + } +} + +impl Drop for WebSocketClient { + fn drop(&mut self) { + if let Some(shutdown_tx) = &self.shutdown_tx { + let _ = shutdown_tx.send(()); + } + } +} diff --git a/cargos/comfyui-sdk/error.rs b/cargos/comfyui-sdk/error.rs new file mode 100644 index 0000000..ce63c6f --- /dev/null +++ b/cargos/comfyui-sdk/error.rs @@ -0,0 +1,86 @@ +//! Error types for the ComfyUI SDK + +use thiserror::Error; + +/// Result type alias for ComfyUI SDK operations +pub type Result = std::result::Result; + +/// Main error type for ComfyUI SDK +#[derive(Error, Debug)] +pub enum ComfyUIError { + /// HTTP request errors + #[error("HTTP error: {0}")] + Http(#[from] reqwest::Error), + + /// WebSocket errors + #[error("WebSocket error: {0}")] + WebSocket(#[from] tokio_tungstenite::tungstenite::Error), + + /// JSON serialization/deserialization errors + #[error("JSON error: {0}")] + Json(#[from] serde_json::Error), + + /// URL parsing errors + #[error("URL error: {0}")] + Url(#[from] url::ParseError), + + /// Template validation errors + #[error("Template validation error: {0}")] + TemplateValidation(String), + + /// Parameter validation errors + #[error("Parameter validation error: {0}")] + ParameterValidation(String), + + /// Connection errors + #[error("Connection error: {0}")] + Connection(String), + + /// Execution errors + #[error("Execution error: {0}")] + Execution(String), + + /// Timeout errors + #[error("Timeout error: {0}")] + Timeout(String), + + /// Generic errors + #[error("Error: {0}")] + Generic(String), + + /// IO errors + #[error("IO error: {0}")] + Io(#[from] std::io::Error), +} + +impl ComfyUIError { + /// Create a new generic error + pub fn new(message: impl Into) -> Self { + Self::Generic(message.into()) + } + + /// Create a new template validation error + pub fn template_validation(message: impl Into) -> Self { + Self::TemplateValidation(message.into()) + } + + /// Create a new parameter validation error + pub fn parameter_validation(message: impl Into) -> Self { + Self::ParameterValidation(message.into()) + } + + /// Create a new connection error + pub fn connection(message: impl Into) -> Self { + Self::Connection(message.into()) + } + + /// Create a new execution error + pub fn execution(message: impl Into) -> Self { + Self::Execution(message.into()) + } + + /// Create a new timeout error + pub fn timeout(message: impl Into) -> Self { + Self::Timeout(message.into()) + } +} diff --git a/cargos/comfyui-sdk/examples/main.rs b/cargos/comfyui-sdk/examples/main.rs new file mode 100644 index 0000000..9cac5de --- /dev/null +++ b/cargos/comfyui-sdk/examples/main.rs @@ -0,0 +1,271 @@ +//! ComfyUI SDK Examples +//! +//! This example demonstrates how to use the ComfyUI SDK for Rust + +use std::collections::HashMap; +use std::sync::Arc; +use comfyui_sdk::{ + ComfyUIClient, ComfyUIClientConfig, ExecutionOptions, + AI_MODEL_FACE_HAIR_FIX_TEMPLATE +}; +use comfyui_sdk::utils::SimpleCallbacks; +use serde_json::json; + +#[tokio::main] +async fn main() -> Result<(), Box> { + // Initialize logging (optional) + // env_logger::init(); + + println!("🚀 Starting ComfyUI SDK Rust Examples"); + + // Run the AI model face hair fix example + ai_model_face_hair_fix_example().await?; + + // Run template validation example + template_validation_example().await?; + + println!("✅ All examples completed successfully!"); + Ok(()) +} + +/// AI Model Face & Hair Detail Fix Example +async fn ai_model_face_hair_fix_example() -> Result<(), Box> { + println!("\n🎨 AI Model Face & Hair Detail Fix Example"); + println!("==========================================="); + + // Initialize client + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + })?; + + // Connect to ComfyUI server + println!("📡 Connecting to ComfyUI server..."); + client.connect().await?; + println!("✅ Connected successfully!"); + + // Register the built-in template + println!("📝 Registering AI Model Face & Hair Fix template..."); + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?; + println!("✅ Template registered!"); + + // Get the registered template + let template = client.templates_ref() + .get_by_id(&AI_MODEL_FACE_HAIR_FIX_TEMPLATE.metadata.id) + .ok_or("Failed to get template")?; + + println!("📋 Template ID: {}", template.id()); + println!("📋 Template Name: {}", template.name()); + if let Some(description) = template.description() { + println!("📋 Template Description: {}", description); + } + + // Display template parameters + println!("\n🔍 Template Parameters:"); + for (name, schema) in template.parameters() { + println!(" - {}: {:?} ({})", + name, + schema.param_type, + schema.description.as_deref().unwrap_or("No description") + ); + } + + // Create execution callbacks + let callbacks = Arc::new( + SimpleCallbacks::new() + .with_progress(|progress| { + let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32; + println!("⏳ Progress: {}% ({}/{}) - Node: {}", + percentage, progress.progress, progress.max, progress.node_id); + }) + .with_executing(|node_id| { + println!("🔄 Executing node: {}", node_id); + }) + .with_executed(|result| { + println!("✅ Node executed - Prompt ID: {}", result.prompt_id); + }) + .with_error(|error| { + println!("❌ Execution error: {}", error.message); + }) + ); + + // Execute the template with test parameters + println!("\n🎨 Executing AI Model Face & Hair Enhancement..."); + let mut parameters = HashMap::new(); + parameters.insert("input_image".to_string(), json!("20250806-190606.jpg")); + parameters.insert("face_prompt".to_string(), json!("beautiful woman, perfect skin, detailed facial features")); + parameters.insert("face_denoise".to_string(), json!("0.25")); + + let execution_options = ExecutionOptions { + timeout: Some(std::time::Duration::from_secs(120)), // 2 minutes timeout + priority: None, + }; + + let result = client.execute_template_with_callbacks( + template, + parameters, + execution_options, + callbacks, + ).await?; + + if result.success { + println!("\n🎉 AI Model Enhancement completed successfully!"); + println!("📊 Execution time: {}ms", result.execution_time); + println!("🆔 Prompt ID: {}", result.prompt_id); + + if let Some(outputs) = &result.outputs { + println!("📁 Outputs: {}", serde_json::to_string_pretty(outputs)?); + + // Convert outputs to HTTP URLs + let image_urls = client.outputs_to_urls(outputs); + if !image_urls.is_empty() { + println!("🖼️ Enhanced Image URLs:"); + for url in image_urls { + println!(" - {}", url); + } + } + } + } else { + println!("❌ AI Model Enhancement failed"); + if let Some(error) = &result.error { + println!(" Error: {}", error.message); + } + } + + // Disconnect + println!("\n🔌 Disconnecting..."); + client.disconnect().await?; + println!("👋 Disconnected!"); + + Ok(()) +} + +/// Template validation example +async fn template_validation_example() -> Result<(), Box> { + println!("\n🔍 Template Validation Example"); + println!("=============================="); + + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + })?; + + // Register template + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?; + let template = client.templates_ref() + .get_by_id(&AI_MODEL_FACE_HAIR_FIX_TEMPLATE.metadata.id) + .ok_or("Failed to get template")?; + + println!("📝 Testing parameter validation..."); + + // Test valid parameters + let mut valid_params = HashMap::new(); + valid_params.insert("input_image".to_string(), json!("test.jpg")); + valid_params.insert("face_prompt".to_string(), json!("beautiful face")); + valid_params.insert("face_denoise".to_string(), json!("0.3")); + + let validation_result = template.validate(&valid_params); + println!("✅ Valid parameters test: {}", + if validation_result.valid { "PASSED" } else { "FAILED" }); + + if !validation_result.valid { + println!("❌ Validation errors:"); + for error in &validation_result.errors { + println!(" - {}: {}", error.path, error.message); + } + } + + // Test invalid parameters + let mut invalid_params = HashMap::new(); + invalid_params.insert("input_image".to_string(), json!("")); // Empty required field + invalid_params.insert("face_prompt".to_string(), json!("test")); + invalid_params.insert("unknown_param".to_string(), json!("value")); // Unknown parameter + + let invalid_validation = template.validate(&invalid_params); + println!("🚫 Invalid parameters test: {}", + if !invalid_validation.valid { "PASSED" } else { "FAILED" }); + + if !invalid_validation.valid { + println!("📋 Expected validation errors:"); + for error in &invalid_validation.errors { + println!(" - {}: {}", error.path, error.message); + } + } + + // Test template instance creation + println!("\n🏗️ Testing template instance creation..."); + match template.create_instance(valid_params) { + Ok(instance) => { + println!("✅ Template instance created successfully"); + println!(" Template ID: {}", instance.template_id()); + println!(" Template Name: {}", instance.template_name()); + println!(" Node count: {}", instance.node_count()); + } + Err(e) => { + println!("❌ Failed to create template instance: {}", e); + } + } + + // Test template manager features + println!("\n📚 Testing template manager features..."); + let templates = client.templates_ref(); + println!(" Total templates: {}", templates.count()); + println!(" Template IDs: {:?}", templates.list_ids()); + + let categories = templates.get_categories(); + if !categories.is_empty() { + println!(" Categories: {:?}", categories); + } + + let tags = templates.get_tags(); + if !tags.is_empty() { + println!(" Tags: {:?}", tags); + } + + Ok(()) +} + +/// Basic usage example (simpler version) +#[allow(dead_code)] +async fn basic_usage_example() -> Result<(), Box> { + println!("\n📖 Basic Usage Example"); + println!("======================"); + + // Create client + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: "http://localhost:8188".to_string(), + ..Default::default() + })?; + + // Connect + client.connect().await?; + println!("✅ Connected to ComfyUI!"); + + // Get system stats + match client.get_system_stats().await { + Ok(stats) => { + println!("💻 System Info:"); + println!(" OS: {}", stats.system.os); + println!(" Python: {}", stats.system.python_version); + println!(" Devices: {}", stats.devices.len()); + } + Err(e) => { + println!("❌ Failed to get system stats: {}", e); + } + } + + // Get queue status + match client.get_queue_status().await { + Ok(queue) => { + println!("📋 Queue Status:"); + println!(" Running: {}", queue.queue_running.len()); + println!(" Pending: {}", queue.queue_pending.len()); + } + Err(e) => { + println!("❌ Failed to get queue status: {}", e); + } + } + + client.disconnect().await?; + Ok(()) +} diff --git a/cargos/comfyui-sdk/examples/real_local_image_test.rs b/cargos/comfyui-sdk/examples/real_local_image_test.rs new file mode 100644 index 0000000..5c3ee89 --- /dev/null +++ b/cargos/comfyui-sdk/examples/real_local_image_test.rs @@ -0,0 +1,306 @@ +//! Real Local Image Upload Test +//! +//! This example demonstrates how to upload your actual local image file +//! and use it with the AI Model Face & Hair Detail Fix template. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; +use comfyui_sdk::{ + ComfyUIClient, ComfyUIClientConfig, ExecutionOptions, + AI_MODEL_FACE_HAIR_FIX_TEMPLATE +}; +use comfyui_sdk::utils::SimpleCallbacks; +use serde_json::json; + +/// Test with real local image upload +async fn test_with_real_local_image() -> Result<(), Box> { + println!("🚀 Starting Real Local Image Upload Test"); + + // Initialize client + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + })?; + + // Connect to ComfyUI server + println!("📡 Connecting to ComfyUI server..."); + client.connect().await?; + println!("✅ Connected successfully!"); + + // Register the template + println!("📝 Registering AI Model Face & Hair Fix template..."); + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?; + let template = client.templates_ref() + .get_by_id(&AI_MODEL_FACE_HAIR_FIX_TEMPLATE.metadata.id) + .ok_or("Failed to register template")?; + println!("✅ Template registered!"); + + // Your actual local image path + let local_image_path = "20250808-111737.png"; + + println!("\n📤 Uploading your image: {}", local_image_path); + println!("⏳ This may take a moment depending on file size and network speed..."); + + match upload_and_execute(&client, template, local_image_path).await { + Ok(_) => println!("✅ Upload and execution completed successfully!"), + Err(e) => { + println!("\n💥 Upload/Execution failed: {}", e); + show_troubleshooting_tips(&e); + } + } + + // Disconnect + println!("\n🔌 Disconnecting from ComfyUI server..."); + client.disconnect().await?; + println!("👋 Disconnected successfully!"); + + Ok(()) +} + +/// Upload image and execute template +async fn upload_and_execute( + client: &ComfyUIClient, + template: &comfyui_sdk::WorkflowTemplate, + image_path: &str, +) -> Result<(), Box> { + // Step 1: Upload image to ComfyUI server + println!("\n🔄 Step 1: Uploading image to ComfyUI server..."); + + if !Path::new(image_path).exists() { + return Err(format!("File not found: {}", image_path).into()); + } + + let upload_response = client.upload_image(image_path, false).await?; + let uploaded_filename = upload_response.name; + println!("✅ Upload successful! Server filename: {}", uploaded_filename); + + // Step 2: Execute AI Model Face & Hair Enhancement + println!("\n🔄 Step 2: Executing AI Model Face & Hair Enhancement..."); + + // Create execution callbacks + let callbacks = Arc::new( + SimpleCallbacks::new() + .with_progress(|progress| { + let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32; + println!("⏳ Progress: {}% ({}/{}) - Node: {}", + percentage, progress.progress, progress.max, progress.node_id); + }) + .with_executing(|node_id| { + println!("🔄 Executing node: {}", node_id); + }) + .with_executed(|result| { + println!("✅ Node completed - Prompt ID: {}", result.prompt_id); + }) + .with_error(|error| { + println!("❌ Execution error: {}", error.message); + }) + ); + + // Prepare parameters + let mut parameters = HashMap::new(); + parameters.insert("input_image".to_string(), json!(uploaded_filename)); + parameters.insert("face_prompt".to_string(), + json!("beautiful woman, perfect skin, detailed facial features, professional photography")); + parameters.insert("face_denoise".to_string(), json!("0.25")); + + let execution_options = ExecutionOptions { + timeout: Some(std::time::Duration::from_secs(180)), // 3 minutes timeout + priority: None, + }; + + let result = client.execute_template_with_callbacks( + template, + parameters, + execution_options, + callbacks, + ).await?; + + if result.success { + println!("\n🎉 AI Model Enhancement completed successfully!"); + println!("📊 Total execution time: {}ms ({:.1}s)", + result.execution_time, result.execution_time as f64 / 1000.0); + println!("🆔 Prompt ID: {}", result.prompt_id); + + if let Some(outputs) = &result.outputs { + println!("\n📁 Generated outputs:"); + println!("{}", serde_json::to_string_pretty(outputs)?); + + // Convert outputs to HTTP URLs + let image_urls = client.outputs_to_urls(outputs); + println!("\n🖼️ Enhanced Image URLs:"); + for (index, url) in image_urls.iter().enumerate() { + println!(" {}. {}", index + 1, url); + } + + println!("\n💡 How to use these URLs:"); + println!(" - Copy the URL and paste in your browser to view"); + println!(" - Right-click and \"Save As\" to download"); + println!(" - Use in your application to display the enhanced image"); + } + } else { + println!("\n❌ AI Model Enhancement failed!"); + if let Some(error) = &result.error { + println!("Error details: {}", error.message); + } + } + + Ok(()) +} + +/// Show troubleshooting tips based on error +fn show_troubleshooting_tips(error: &Box) { + let error_msg = error.to_string(); + + if error_msg.contains("File not found") || error_msg.contains("No such file") { + println!("\n💡 Troubleshooting Tips:"); + println!("1. 🔍 Check if the file path is correct"); + println!("2. 📁 Make sure the file exists at the specified location"); + println!("3. 🔐 Ensure you have read permissions for the file"); + println!("4. 📝 Try using absolute path: /home/user/images/20250808-111737.png"); + println!("5. 🖼️ Verify the file is a valid image format (PNG, JPG, etc.)"); + } else if error_msg.contains("Failed to upload") || error_msg.contains("HTTP") { + println!("\n💡 Network/Server Issues:"); + println!("1. 🌐 Check if ComfyUI server is running and accessible"); + println!("2. 📡 Verify network connection to the server"); + println!("3. 💾 Check if server has enough disk space"); + println!("4. 🔒 Ensure server allows file uploads"); + } +} + +/// Alternative method using one-step upload and execute (conceptual) +async fn test_one_step_method() -> Result<(), Box> { + println!("\n\n🚀 Testing One-Step Method (Upload + Execute)"); + + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + })?; + + client.connect().await?; + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?; + let template = client.templates_ref() + .get_by_id(&AI_MODEL_FACE_HAIR_FIX_TEMPLATE.metadata.id) + .ok_or("Failed to register template")?; + + let local_image_path = "20250808-111737.png"; + + println!("🔄 One-step upload and execute..."); + + // For now, we'll implement this as upload + execute since we don't have + // a dedicated one-step method yet + match upload_and_execute_one_step(&client, template, local_image_path).await { + Ok(_) => println!("🎉 One-step method completed successfully!"), + Err(e) => println!("⚠️ One-step method failed: {}", e), + } + + client.disconnect().await?; + Ok(()) +} + +/// One-step upload and execute (simplified version) +async fn upload_and_execute_one_step( + client: &ComfyUIClient, + template: &comfyui_sdk::WorkflowTemplate, + image_path: &str, +) -> Result<(), Box> { + // Upload image + let upload_response = client.upload_image(image_path, false).await?; + + // Create simple callbacks + let callbacks = Arc::new( + SimpleCallbacks::new() + .with_progress(|progress| { + let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32; + println!("⏳ Progress: {}% - Node: {}", percentage, progress.node_id); + }) + ); + + // Execute template + let mut parameters = HashMap::new(); + parameters.insert("input_image".to_string(), json!(upload_response.name)); + parameters.insert("face_prompt".to_string(), + json!("stunning model, flawless skin, professional photography, high quality")); + parameters.insert("face_denoise".to_string(), json!("0.3")); + + let result = client.execute_template_with_callbacks( + template, + parameters, + ExecutionOptions { + timeout: Some(std::time::Duration::from_secs(180)), + priority: None, + }, + callbacks, + ).await?; + + if result.success { + if let Some(outputs) = &result.outputs { + let image_urls = client.outputs_to_urls(outputs); + println!("🖼️ Result URLs: {:?}", image_urls); + } + } + + Ok(()) +} + +/// Show usage examples +fn show_usage_examples() { + println!("\n📚 Complete Usage Guide for Local Images:"); + println!(""); + println!("🔧 Basic Setup:"); + println!("```rust"); + println!("use comfyui_sdk::{{ComfyUIClient, ComfyUIClientConfig, AI_MODEL_FACE_HAIR_FIX_TEMPLATE}};"); + println!(""); + println!("let mut client = ComfyUIClient::new(ComfyUIClientConfig {{"); + println!(" base_url: \"http://your-comfyui-server:8188\".to_string(),"); + println!(" ..Default::default()"); + println!("}});"); + println!("```"); + println!(""); + println!("📤 Method 1 - Manual Upload:"); + println!("```rust"); + println!("// Upload image first"); + println!("let upload_response = client.upload_image("); + println!(" \"/path/to/your/image.png\", false"); + println!(").await?;"); + println!(""); + println!("// Then execute template"); + println!("let result = client.execute_template(template, parameters, options).await?;"); + println!("```"); + println!(""); + println!("🚀 Method 2 - Combined Upload and Execute:"); + println!("```rust"); + println!("// Upload and execute in sequence"); + println!("let upload_response = client.upload_image(image_path, false).await?;"); + println!("let mut parameters = HashMap::new();"); + println!("parameters.insert(\"input_image\".to_string(), json!(upload_response.name));"); + println!("let result = client.execute_template_with_callbacks("); + println!(" template, parameters, options, callbacks"); + println!(").await?;"); + println!("```"); +} + +/// Run all tests +async fn run_all_tests() -> Result<(), Box> { + show_usage_examples(); + test_with_real_local_image().await?; + // Uncomment to test one-step method as well + // test_one_step_method().await?; + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + run_all_tests().await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_show_usage_examples() { + show_usage_examples(); + // This test just ensures the function runs without panicking + } +} diff --git a/cargos/comfyui-sdk/examples/simple_local_image.rs b/cargos/comfyui-sdk/examples/simple_local_image.rs new file mode 100644 index 0000000..91d2bc9 --- /dev/null +++ b/cargos/comfyui-sdk/examples/simple_local_image.rs @@ -0,0 +1,162 @@ +//! Simple Local Image Upload and Enhancement +//! +//! A simplified version that focuses on the core functionality: +//! 1. Upload a local image to ComfyUI server +//! 2. Execute AI Model Face & Hair enhancement +//! 3. Get the enhanced image URLs + +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; +use comfyui_sdk::{ + ComfyUIClient, ComfyUIClientConfig, ExecutionOptions, + AI_MODEL_FACE_HAIR_FIX_TEMPLATE +}; +use comfyui_sdk::utils::SimpleCallbacks; +use serde_json::json; + +#[tokio::main] +async fn main() -> Result<(), Box> { + println!("🚀 Simple Local Image Enhancement Test"); + + // Configuration + let server_url = "http://192.168.0.193:8188"; + let image_path = "20250808-111737.png"; + + // Step 1: Initialize client and connect + println!("📡 Connecting to ComfyUI server at {}...", server_url); + let mut client = ComfyUIClient::new(ComfyUIClientConfig { + base_url: server_url.to_string(), + ..Default::default() + })?; + + client.connect().await?; + println!("✅ Connected!"); + + // Step 2: Register template + println!("📝 Registering AI enhancement template..."); + client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?; + let template = client.templates_ref() + .get_by_id("ai-model-face-hair-fix") + .ok_or("Template not found")?; + println!("✅ Template ready!"); + + // Step 3: Check if image exists + if !Path::new(image_path).exists() { + println!("❌ Image file not found: {}", image_path); + println!("💡 Please make sure the file exists in the current directory"); + return Ok(()); + } + + // Step 4: Upload image + println!("📤 Uploading image: {}...", image_path); + let upload_response = client.upload_image(image_path, false).await?; + println!("✅ Upload successful! Server filename: {}", upload_response.name); + + // Step 5: Setup progress callbacks + let callbacks = Arc::new( + SimpleCallbacks::new() + .with_progress(|progress| { + let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32; + println!("⏳ Progress: {}%", percentage); + }) + .with_executing(|node_id| { + println!("🔄 Processing node: {}", node_id); + }) + .with_error(|error| { + println!("❌ Error: {}", error.message); + }) + ); + + // Step 6: Execute enhancement + println!("🎨 Starting AI enhancement..."); + let mut parameters = HashMap::new(); + parameters.insert("input_image".to_string(), json!(upload_response.name)); + parameters.insert("face_prompt".to_string(), + json!("beautiful woman, perfect skin, detailed facial features")); + parameters.insert("face_denoise".to_string(), json!("0.25")); + + let result = client.execute_template_with_callbacks( + template, + parameters, + ExecutionOptions { + timeout: Some(std::time::Duration::from_secs(180)), + priority: None, + }, + callbacks, + ).await?; + + // Step 7: Show results + if result.success { + println!("\n🎉 Enhancement completed successfully!"); + println!("⏱️ Execution time: {:.1}s", result.execution_time as f64 / 1000.0); + + if let Some(outputs) = &result.outputs { + let image_urls = client.outputs_to_urls(outputs); + + if !image_urls.is_empty() { + println!("\n🖼️ Enhanced image URLs:"); + for (i, url) in image_urls.iter().enumerate() { + println!(" {}. {}", i + 1, url); + } + + println!("\n💡 To view your enhanced image:"); + println!(" - Copy the URL above and paste it in your browser"); + println!(" - Right-click and 'Save As' to download the image"); + } else { + println!("⚠️ No image URLs found in outputs"); + } + } + } else { + println!("❌ Enhancement failed!"); + if let Some(error) = &result.error { + println!("Error: {}", error.message); + } + } + + // Step 8: Cleanup + client.disconnect().await?; + println!("\n👋 Disconnected from server"); + + Ok(()) +} + +/// Helper function to validate image file +fn validate_image_file(path: &str) -> Result<(), String> { + let path = Path::new(path); + + if !path.exists() { + return Err(format!("File does not exist: {}", path.display())); + } + + if !path.is_file() { + return Err(format!("Path is not a file: {}", path.display())); + } + + // Check file extension + 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!("Unsupported image format: {}", ext)), + } + } else { + Err("File has no extension".to_string()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_image_file() { + // Test with non-existent file + assert!(validate_image_file("non_existent.png").is_err()); + + // Test with valid extensions + // Note: These tests would pass if the files existed + // assert!(validate_image_file("test.png").is_ok()); + // assert!(validate_image_file("test.jpg").is_ok()); + } +} diff --git a/cargos/comfyui-sdk/examples/test_url_fix.rs b/cargos/comfyui-sdk/examples/test_url_fix.rs new file mode 100644 index 0000000..ba6d417 --- /dev/null +++ b/cargos/comfyui-sdk/examples/test_url_fix.rs @@ -0,0 +1,113 @@ +//! Test URL generation fix +//! +//! This example demonstrates that the double slash bug has been fixed + +use std::collections::HashMap; +use serde_json::json; +use comfyui_sdk::{HTTPClient, ComfyUIClientConfig}; + +fn main() { + println!("🔧 Testing URL Generation Fix"); + println!("=============================="); + + // Test with base URL ending with slash (the problematic case) + let config_with_slash = ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188/".to_string(), + ..Default::default() + }; + let client_with_slash = HTTPClient::new(config_with_slash).unwrap(); + + // Test with base URL without ending slash + let config_without_slash = ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + }; + let client_without_slash = HTTPClient::new(config_without_slash).unwrap(); + + // Create test outputs (simulating ComfyUI response) + let mut outputs = HashMap::new(); + outputs.insert("58".to_string(), json!({ + "images": [ + { + "filename": "ComfyUI_02046_.png", + "subfolder": "", + "type": "output" + } + ] + })); + + println!("\n📋 Test Case 1: Base URL with trailing slash"); + println!("Base URL: http://192.168.0.193:8188/"); + let urls_with_slash = client_with_slash.outputs_to_urls(&outputs); + for url in &urls_with_slash { + println!("Generated URL: {}", url); + if url.contains("//view") { + println!("❌ FAILED: URL contains double slash!"); + } else { + println!("✅ PASSED: No double slash found"); + } + } + + println!("\n📋 Test Case 2: Base URL without trailing slash"); + println!("Base URL: http://192.168.0.193:8188"); + let urls_without_slash = client_without_slash.outputs_to_urls(&outputs); + for url in &urls_without_slash { + println!("Generated URL: {}", url); + if url.contains("//view") { + println!("❌ FAILED: URL contains double slash!"); + } else { + println!("✅ PASSED: No double slash found"); + } + } + + println!("\n📋 Test Case 3: Multiple images"); + let mut multi_outputs = HashMap::new(); + multi_outputs.insert("58".to_string(), json!({ + "images": [ + { + "filename": "image1.png", + "subfolder": "temp", + "type": "output" + }, + { + "filename": "image2.jpg", + "subfolder": "", + "type": "temp" + } + ] + })); + + let multi_urls = client_with_slash.outputs_to_urls(&multi_outputs); + println!("Generated {} URLs:", multi_urls.len()); + for (i, url) in multi_urls.iter().enumerate() { + println!(" {}. {}", i + 1, url); + if url.contains("//view") { + println!(" ❌ FAILED: URL contains double slash!"); + } else { + println!(" ✅ PASSED: No double slash found"); + } + } + + println!("\n🎯 Expected vs Actual URLs:"); + println!("Expected: http://192.168.0.193:8188/view?filename=ComfyUI_02046_.png&subfolder=&type=output"); + if !urls_with_slash.is_empty() { + println!("Actual: {}", urls_with_slash[0]); + if urls_with_slash[0] == "http://192.168.0.193:8188/view?filename=ComfyUI_02046_.png&subfolder=&type=output" { + println!("✅ URLs match perfectly!"); + } else { + println!("❌ URLs don't match!"); + } + } + + println!("\n🔍 Consistency Check:"); + if urls_with_slash == urls_without_slash { + println!("✅ URLs are consistent regardless of trailing slash in base URL"); + } else { + println!("❌ URLs are inconsistent!"); + println!("With slash: {:?}", urls_with_slash); + println!("Without slash: {:?}", urls_without_slash); + } + + println!("\n🎉 URL Generation Fix Test Completed!"); + println!("The double slash bug has been fixed. URLs now correctly use single slashes."); +} diff --git a/cargos/comfyui-sdk/lib.rs b/cargos/comfyui-sdk/lib.rs new file mode 100644 index 0000000..166e1f5 --- /dev/null +++ b/cargos/comfyui-sdk/lib.rs @@ -0,0 +1,48 @@ +//! # ComfyUI SDK for Rust +//! +//! A modern Rust SDK for ComfyUI with flexible template system and elegant callback mechanisms. +//! +//! ## Features +//! +//! - **HTTP Client**: Full-featured HTTP client for ComfyUI API +//! - **WebSocket Client**: Real-time event handling and progress tracking +//! - **Template System**: Parameterized workflow templates with validation +//! - **Built-in Templates**: Pre-built templates for common workflows +//! - **Type Safety**: Full type safety with comprehensive error handling +//! - **Async/Await**: Modern async/await support throughout +//! +//! ## Quick Start +//! +//! ```rust,no_run +//! use comfyui_sdk::{ComfyUIClient, ComfyUIClientConfig}; +//! +//! #[tokio::main] +//! async fn main() -> Result<(), Box> { +//! let mut client = ComfyUIClient::new(ComfyUIClientConfig { +//! base_url: "http://localhost:8188".to_string(), +//! ..Default::default() +//! })?; +//! +//! client.connect().await?; +//! println!("Connected to ComfyUI!"); +//! +//! Ok(()) +//! } +//! ``` + +pub mod client; +pub mod templates; +pub mod template_hub; +pub mod types; +pub mod utils; +pub mod error; + +// Re-export main types and functions +pub use client::{ComfyUIClient, HTTPClient, WebSocketClient}; +pub use templates::{WorkflowTemplate, WorkflowInstance, TemplateManager}; +pub use template_hub::*; +pub use types::*; +pub use error::{ComfyUIError, Result}; + +/// SDK version +pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/cargos/comfyui-sdk/template_hub/ai_model_face_hair_fix.rs b/cargos/comfyui-sdk/template_hub/ai_model_face_hair_fix.rs new file mode 100644 index 0000000..7f0e802 --- /dev/null +++ b/cargos/comfyui-sdk/template_hub/ai_model_face_hair_fix.rs @@ -0,0 +1,435 @@ +//! AI Model Face & Hair Detail Fix Template +//! +//! This template performs detailed face and hair enhancement for AI model photos. +//! It uses segmentation masks and local inpainting to improve facial features and hair quality. +//! +//! Features: +//! - Automatic face and hair segmentation +//! - Two-stage local inpainting (hair first, then face) +//! - Configurable denoising strength +//! - High-quality detail enhancement + +use std::collections::HashMap; +use serde_json::json; +use crate::types::{ + WorkflowTemplateData, TemplateMetadata, ParameterSchema, ParameterType, + ComfyUIWorkflow, ComfyUINode, NodeMeta +}; + +/// AI Model Face & Hair Detail Fix Template +pub static AI_MODEL_FACE_HAIR_FIX_TEMPLATE: once_cell::sync::Lazy = + once_cell::sync::Lazy::new(|| { + create_ai_model_face_hair_fix_template() + }); + +fn create_ai_model_face_hair_fix_template() -> WorkflowTemplateData { + let metadata = TemplateMetadata { + id: "ai-model-face-hair-fix".to_string(), + name: "AI Model Face & Hair Detail Fix".to_string(), + description: Some("Professional face and hair detail enhancement for AI model photos using segmentation and local inpainting".to_string()), + category: Some("enhancement".to_string()), + tags: Some(vec![ + "face-enhancement".to_string(), + "hair-enhancement".to_string(), + "portrait".to_string(), + "ai-model".to_string(), + "detail-fix".to_string(), + "inpainting".to_string(), + ]), + version: Some("1.0.0".to_string()), + author: Some("ComfyUI SDK".to_string()), + created_at: None, + updated_at: None, + }; + + let mut parameters = HashMap::new(); + + parameters.insert("input_image".to_string(), ParameterSchema { + param_type: ParameterType::String, + required: Some(true), + default: Some(json!("20250806-190606.jpg")), + description: Some("Input image filename to enhance".to_string()), + r#enum: None, + min: None, + max: None, + pattern: None, + items: None, + properties: None, + }); + + parameters.insert("face_prompt".to_string(), ParameterSchema { + param_type: ParameterType::String, + required: Some(true), + default: Some(json!("A girl, slim figure, oval face, beauty, Pink and greasy lips，")), + description: Some("Positive prompt for face enhancement".to_string()), + r#enum: None, + min: None, + max: None, + pattern: None, + items: None, + properties: None, + }); + + parameters.insert("face_denoise".to_string(), ParameterSchema { + param_type: ParameterType::String, + required: Some(true), + default: Some(json!("0.30")), + description: Some("Denoising strength for face enhancement".to_string()), + r#enum: None, + min: None, + max: None, + pattern: None, + items: None, + properties: None, + }); + + let workflow = create_workflow(); + + WorkflowTemplateData { + metadata, + workflow, + parameters, + } +} + +fn create_workflow() -> ComfyUIWorkflow { + let mut workflow = HashMap::new(); + + // Node 6: 遮罩边缘滑条快速模糊 + workflow.insert("6".to_string(), ComfyUINode { + class_type: "EG_ZZ_MHHT".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("模糊强度".to_string(), json!(3)); + inputs.insert("mask".to_string(), json!(["53", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("2🐕遮罩边缘滑条快速模糊".to_string()), + }), + }); + + // Node 8: Checkpoint加载器 + workflow.insert("8".to_string(), ComfyUINode { + class_type: "CheckpointLoaderSimple".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("ckpt_name".to_string(), json!("juggernaut_reborn_sd1,5.safetensors")); + inputs + }, + _meta: Some(NodeMeta { + title: Some("Checkpoint加载器（简易）".to_string()), + }), + }); + + // Node 9: CLIP文本编码 (头发) + workflow.insert("9".to_string(), ComfyUINode { + class_type: "CLIPTextEncode".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("text".to_string(), json!("Smooth long hair,Hair details, quality hair,Smooth hair")); + inputs.insert("clip".to_string(), json!(["8", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("CLIP文本编码".to_string()), + }), + }); + + // Node 10: CLIP文本编码 (负面头发) + workflow.insert("10".to_string(), ComfyUINode { + class_type: "CLIPTextEncode".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("text".to_string(), json!("Frizzy hair, dry hair, split ends,bad quality, poor quality, doll, disfigured, jpg, toy, bad anatomy, missing limbs, missing fingers, 3d, cgi")); + inputs.insert("clip".to_string(), json!(["8", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("CLIP文本编码".to_string()), + }), + }); + + // Node 13: CLIP文本编码 (脸部) + workflow.insert("13".to_string(), ComfyUINode { + class_type: "CLIPTextEncode".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("text".to_string(), json!(["59", 0])); + inputs.insert("clip".to_string(), json!(["8", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("CLIP文本编码".to_string()), + }), + }); + + // Node 14: CLIP文本编码 (负面脸部) + workflow.insert("14".to_string(), ComfyUINode { + class_type: "CLIPTextEncode".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("text".to_string(), json!("(NSFW:1.2),(worst quality:1.2),(low quality:1.2),(normal quality:1.2),low resolution,watermark,")); + inputs.insert("clip".to_string(), json!(["8", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("CLIP文本编码".to_string()), + }), + }); + + // Node 22: Segformer Ultra V2 (头发) + workflow.insert("22".to_string(), ComfyUINode { + class_type: "LayerMask: SegformerUltraV2".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("detail_method".to_string(), json!("VITMatte(local)")); + inputs.insert("detail_erode".to_string(), json!(44)); + inputs.insert("detail_dilate".to_string(), json!(6)); + inputs.insert("black_point".to_string(), json!(0.030000000000000006)); + inputs.insert("white_point".to_string(), json!(0.99)); + inputs.insert("process_detail".to_string(), json!(true)); + inputs.insert("device".to_string(), json!("cuda")); + inputs.insert("max_megapixels".to_string(), json!(2)); + inputs.insert("image".to_string(), json!(["30", 0])); + inputs.insert("segformer_pipeline".to_string(), json!(["28", 0])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("LayerMask: Segformer Ultra V2".to_string()), + }), + }); + + // Node 25: 局部重绘采样器 (头发) + workflow.insert("25".to_string(), ComfyUINode { + class_type: "EG_CYQ_JB".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("seed".to_string(), json!(646617836820462i64)); + inputs.insert("steps".to_string(), json!(30)); + inputs.insert("cfg".to_string(), json!(4.5200000000000005)); + inputs.insert("sampler_name".to_string(), json!("dpmpp_2s_ancestral")); + inputs.insert("scheduler".to_string(), json!("karras")); + inputs.insert("denoise".to_string(), json!(0.4)); + inputs.insert("重绘模式".to_string(), json!("原图")); + inputs.insert("遮罩延展".to_string(), json!(10)); + inputs.insert("仅局部重绘".to_string(), json!(true)); + inputs.insert("局部重绘大小".to_string(), json!(768)); + inputs.insert("重绘区域扩展".to_string(), json!(50)); + inputs.insert("遮罩羽化".to_string(), json!(5)); + inputs.insert("TEXT".to_string(), json!("2🐕出品，必出精品")); + inputs.insert("model".to_string(), json!(["8", 0])); + inputs.insert("image".to_string(), json!(["30", 0])); + inputs.insert("vae".to_string(), json!(["8", 2])); + inputs.insert("mask".to_string(), json!(["47", 0])); + inputs.insert("positive".to_string(), json!(["9", 0])); + inputs.insert("negative".to_string(), json!(["10", 0])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("2🐕局部重绘采样器".to_string()), + }), + }); + + // Node 26: 局部重绘采样器 (脸部) + workflow.insert("26".to_string(), ComfyUINode { + class_type: "EG_CYQ_JB".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("seed".to_string(), json!(969968724502727i64)); + inputs.insert("steps".to_string(), json!(30)); + inputs.insert("cfg".to_string(), json!(4.5200000000000005)); + inputs.insert("sampler_name".to_string(), json!("dpmpp_2s_ancestral")); + inputs.insert("scheduler".to_string(), json!("karras")); + inputs.insert("denoise".to_string(), json!(["71", 0])); + inputs.insert("重绘模式".to_string(), json!("原图")); + inputs.insert("遮罩延展".to_string(), json!(10)); + inputs.insert("仅局部重绘".to_string(), json!(true)); + inputs.insert("局部重绘大小".to_string(), json!(768)); + inputs.insert("重绘区域扩展".to_string(), json!(50)); + inputs.insert("遮罩羽化".to_string(), json!(5)); + inputs.insert("TEXT".to_string(), json!("2🐕出品，必出精品")); + inputs.insert("model".to_string(), json!(["8", 0])); + inputs.insert("image".to_string(), json!(["25", 1])); + inputs.insert("vae".to_string(), json!(["8", 2])); + inputs.insert("mask".to_string(), json!(["6", 0])); + inputs.insert("positive".to_string(), json!(["13", 0])); + inputs.insert("negative".to_string(), json!(["14", 0])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("2🐕局部重绘采样器".to_string()), + }), + }); + + // Node 28: Segformer Clothes Pipeline (头发) + workflow.insert("28".to_string(), ComfyUINode { + class_type: "LayerMask: SegformerClothesPipelineLoader".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("model".to_string(), json!("segformer_b3_clothes")); + inputs.insert("face".to_string(), json!(false)); + inputs.insert("hair".to_string(), json!(true)); + inputs.insert("hat".to_string(), json!(false)); + inputs.insert("sunglass".to_string(), json!(false)); + inputs.insert("left_arm".to_string(), json!(false)); + inputs.insert("right_arm".to_string(), json!(false)); + inputs.insert("left_leg".to_string(), json!(false)); + inputs.insert("right_leg".to_string(), json!(false)); + inputs.insert("left_shoe".to_string(), json!(false)); + inputs.insert("right_shoe".to_string(), json!(false)); + inputs.insert("upper_clothes".to_string(), json!(false)); + inputs.insert("skirt".to_string(), json!(false)); + inputs.insert("pants".to_string(), json!(false)); + inputs.insert("dress".to_string(), json!(false)); + inputs.insert("belt".to_string(), json!(false)); + inputs.insert("bag".to_string(), json!(false)); + inputs.insert("scarf".to_string(), json!(false)); + inputs + }, + _meta: Some(NodeMeta { + title: Some("LayerMask: Segformer Clothes Pipeline".to_string()), + }), + }); + + // Node 30: 加载图像 + workflow.insert("30".to_string(), ComfyUINode { + class_type: "LoadImage".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("image".to_string(), json!("{{input_image}}")); + inputs + }, + _meta: Some(NodeMeta { + title: Some("加载图像".to_string()), + }), + }); + + // Node 44: Mask Fill Holes + workflow.insert("44".to_string(), ComfyUINode { + class_type: "Mask Fill Holes".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("masks".to_string(), json!(["22", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("Mask Fill Holes".to_string()), + }), + }); + + // Node 47: Grow Mask With Blur + workflow.insert("47".to_string(), ComfyUINode { + class_type: "GrowMaskWithBlur".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("expand".to_string(), json!(5)); + inputs.insert("incremental_expandrate".to_string(), json!(5)); + inputs.insert("tapered_corners".to_string(), json!(true)); + inputs.insert("flip_input".to_string(), json!(false)); + inputs.insert("blur_radius".to_string(), json!(9.4)); + inputs.insert("lerp_alpha".to_string(), json!(0.9900000000000002)); + inputs.insert("decay_factor".to_string(), json!(1)); + inputs.insert("fill_holes".to_string(), json!(false)); + inputs.insert("mask".to_string(), json!(["44", 0])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("Grow Mask With Blur".to_string()), + }), + }); + + // Node 52: Segformer Clothes Pipeline (脸部) + workflow.insert("52".to_string(), ComfyUINode { + class_type: "LayerMask: SegformerClothesPipelineLoader".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("model".to_string(), json!("segformer_b3_clothes")); + inputs.insert("face".to_string(), json!(true)); + inputs.insert("hair".to_string(), json!(false)); + inputs.insert("hat".to_string(), json!(false)); + inputs.insert("sunglass".to_string(), json!(true)); + inputs.insert("left_arm".to_string(), json!(false)); + inputs.insert("right_arm".to_string(), json!(false)); + inputs.insert("left_leg".to_string(), json!(false)); + inputs.insert("right_leg".to_string(), json!(false)); + inputs.insert("left_shoe".to_string(), json!(false)); + inputs.insert("right_shoe".to_string(), json!(false)); + inputs.insert("upper_clothes".to_string(), json!(false)); + inputs.insert("skirt".to_string(), json!(false)); + inputs.insert("pants".to_string(), json!(false)); + inputs.insert("dress".to_string(), json!(false)); + inputs.insert("belt".to_string(), json!(false)); + inputs.insert("bag".to_string(), json!(false)); + inputs.insert("scarf".to_string(), json!(false)); + inputs + }, + _meta: Some(NodeMeta { + title: Some("LayerMask: Segformer Clothes Pipeline".to_string()), + }), + }); + + // Node 53: Segformer Ultra V2 (脸部) + workflow.insert("53".to_string(), ComfyUINode { + class_type: "LayerMask: SegformerUltraV2".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("detail_method".to_string(), json!("VITMatte(local)")); + inputs.insert("detail_erode".to_string(), json!(6)); + inputs.insert("detail_dilate".to_string(), json!(4)); + inputs.insert("black_point".to_string(), json!(0.01)); + inputs.insert("white_point".to_string(), json!(0.99)); + inputs.insert("process_detail".to_string(), json!(false)); + inputs.insert("device".to_string(), json!("cuda")); + inputs.insert("max_megapixels".to_string(), json!(2)); + inputs.insert("image".to_string(), json!(["30", 0])); + inputs.insert("segformer_pipeline".to_string(), json!(["52", 0])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("LayerMask: Segformer Ultra V2".to_string()), + }), + }); + + // Node 58: 保存图像 + workflow.insert("58".to_string(), ComfyUINode { + class_type: "SaveImage".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("filename_prefix".to_string(), json!("ComfyUI")); + inputs.insert("images".to_string(), json!(["26", 1])); + inputs + }, + _meta: Some(NodeMeta { + title: Some("保存图像".to_string()), + }), + }); + + // Node 59: String (脸部提示词) + workflow.insert("59".to_string(), ComfyUINode { + class_type: "String".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("String".to_string(), json!("A girl, slim figure, oval face, beauty, Pink and greasy lips，{{face_prompt}}")); + inputs + }, + _meta: Some(NodeMeta { + title: Some("String".to_string()), + }), + }); + + // Node 71: Float (脸部去噪强度) + workflow.insert("71".to_string(), ComfyUINode { + class_type: "Float".to_string(), + inputs: { + let mut inputs = HashMap::new(); + inputs.insert("Number".to_string(), json!("{{face_denoise}}")); + inputs + }, + _meta: Some(NodeMeta { + title: Some("Float".to_string()), + }), + }); + + workflow +} diff --git a/cargos/comfyui-sdk/template_hub/mod.rs b/cargos/comfyui-sdk/template_hub/mod.rs new file mode 100644 index 0000000..89aef83 --- /dev/null +++ b/cargos/comfyui-sdk/template_hub/mod.rs @@ -0,0 +1,6 @@ +//! Built-in template hub + +pub mod ai_model_face_hair_fix; + +// Re-export templates +pub use ai_model_face_hair_fix::AI_MODEL_FACE_HAIR_FIX_TEMPLATE; diff --git a/cargos/comfyui-sdk/templates/mod.rs b/cargos/comfyui-sdk/templates/mod.rs new file mode 100644 index 0000000..0893601 --- /dev/null +++ b/cargos/comfyui-sdk/templates/mod.rs @@ -0,0 +1,10 @@ +//! Template system for ComfyUI SDK + +pub mod workflow_template; +pub mod workflow_instance; +pub mod template_manager; + +// Re-export main types +pub use workflow_template::WorkflowTemplate; +pub use workflow_instance::WorkflowInstance; +pub use template_manager::TemplateManager; diff --git a/cargos/comfyui-sdk/templates/template_manager.rs b/cargos/comfyui-sdk/templates/template_manager.rs new file mode 100644 index 0000000..323fe17 --- /dev/null +++ b/cargos/comfyui-sdk/templates/template_manager.rs @@ -0,0 +1,194 @@ +//! TemplateManager for managing workflow templates + +use std::collections::HashMap; +use crate::types::{WorkflowTemplateData, ParameterValues}; +use crate::templates::{WorkflowTemplate, WorkflowInstance}; +use crate::error::{ComfyUIError, Result}; + +/// Manages a collection of workflow templates +#[derive(Debug, Clone)] +pub struct TemplateManager { + templates: HashMap, +} + +impl TemplateManager { + /// Creates a new TemplateManager + pub fn new() -> Self { + Self { + templates: HashMap::new(), + } + } + + /// Registers a template from template data + pub fn register_from_data(&mut self, data: WorkflowTemplateData) -> Result<()> { + let template = WorkflowTemplate::from_data(data)?; + let id = template.id().to_string(); + self.templates.insert(id, template); + Ok(()) + } + + /// Registers a template + pub fn register(&mut self, template: WorkflowTemplate) { + let id = template.id().to_string(); + self.templates.insert(id, template); + } + + /// Gets a template by ID + pub fn get_by_id(&self, id: &str) -> Option<&WorkflowTemplate> { + self.templates.get(id) + } + + /// Gets a mutable reference to a template by ID + pub fn get_by_id_mut(&mut self, id: &str) -> Option<&mut WorkflowTemplate> { + self.templates.get_mut(id) + } + + /// Removes a template by ID + pub fn remove(&mut self, id: &str) -> Option { + self.templates.remove(id) + } + + /// Lists all template IDs + pub fn list_ids(&self) -> Vec<&str> { + self.templates.keys().map(|s| s.as_str()).collect() + } + + /// Lists all templates + pub fn list_templates(&self) -> Vec<&WorkflowTemplate> { + self.templates.values().collect() + } + + /// Gets templates by category + pub fn get_by_category(&self, category: &str) -> Vec<&WorkflowTemplate> { + self.templates + .values() + .filter(|template| { + template.category() + .map(|c| c == category) + .unwrap_or(false) + }) + .collect() + } + + /// Gets templates by tag + pub fn get_by_tag(&self, tag: &str) -> Vec<&WorkflowTemplate> { + self.templates + .values() + .filter(|template| template.tags().contains(&tag.to_string())) + .collect() + } + + /// Searches templates by name (case-insensitive partial match) + pub fn search_by_name(&self, query: &str) -> Vec<&WorkflowTemplate> { + let query_lower = query.to_lowercase(); + self.templates + .values() + .filter(|template| { + template.name().to_lowercase().contains(&query_lower) + }) + .collect() + } + + /// Searches templates by description (case-insensitive partial match) + pub fn search_by_description(&self, query: &str) -> Vec<&WorkflowTemplate> { + let query_lower = query.to_lowercase(); + self.templates + .values() + .filter(|template| { + template.description() + .map(|desc| desc.to_lowercase().contains(&query_lower)) + .unwrap_or(false) + }) + .collect() + } + + /// Creates a workflow instance from a template + pub fn create_instance(&self, template_id: &str, parameters: ParameterValues) -> Result { + let template = self.get_by_id(template_id) + .ok_or_else(|| ComfyUIError::template_validation( + format!("Template with ID '{}' not found", template_id) + ))?; + + template.create_instance(parameters) + } + + /// Gets the number of registered templates + pub fn count(&self) -> usize { + self.templates.len() + } + + /// Checks if a template exists + pub fn contains(&self, id: &str) -> bool { + self.templates.contains_key(id) + } + + /// Clears all templates + pub fn clear(&mut self) { + self.templates.clear(); + } + + /// Gets all categories + pub fn get_categories(&self) -> Vec { + let mut categories: Vec = self.templates + .values() + .filter_map(|template| template.category().map(|c| c.to_string())) + .collect(); + categories.sort(); + categories.dedup(); + categories + } + + /// Gets all tags + pub fn get_tags(&self) -> Vec { + let mut tags: Vec = self.templates + .values() + .flat_map(|template| template.tags().iter().cloned()) + .collect(); + tags.sort(); + tags.dedup(); + tags + } + + /// Validates all templates + pub fn validate_all(&self) -> Result<()> { + for (id, template) in &self.templates { + // Try to create an instance with default values to validate template + let default_values = template.get_default_values(); + let validation = template.validate(&default_values); + + if !validation.valid { + return Err(ComfyUIError::template_validation( + format!("Template '{}' validation failed", id) + )); + } + } + Ok(()) + } + + /// Exports all templates as template data + pub fn export_all(&self) -> Vec { + self.templates + .values() + .map(|template| template.to_data()) + .collect() + } + + /// Imports templates from template data + pub fn import_all(&mut self, templates: Vec) -> Result> { + let mut imported_ids = Vec::new(); + + for data in templates { + let id = data.metadata.id.clone(); + self.register_from_data(data)?; + imported_ids.push(id); + } + + Ok(imported_ids) + } +} + +impl Default for TemplateManager { + fn default() -> Self { + Self::new() + } +} diff --git a/cargos/comfyui-sdk/templates/workflow_instance.rs b/cargos/comfyui-sdk/templates/workflow_instance.rs new file mode 100644 index 0000000..3336b00 --- /dev/null +++ b/cargos/comfyui-sdk/templates/workflow_instance.rs @@ -0,0 +1,169 @@ +//! WorkflowInstance implementation for executing parameterized workflows + +use std::collections::HashMap; +use crate::types::{ComfyUIWorkflow, ParameterValues, ComfyUINode}; +use crate::templates::WorkflowTemplate; +use crate::error::{ComfyUIError, Result}; +use crate::utils::template_parser::apply_parameters; + +/// Represents a workflow instance with applied parameters +#[derive(Debug, Clone)] +pub struct WorkflowInstance { + template: WorkflowTemplate, + parameters: ParameterValues, + resolved_workflow: ComfyUIWorkflow, +} + +impl WorkflowInstance { + /// Creates a new WorkflowInstance + pub fn new(template: WorkflowTemplate, parameters: ParameterValues) -> Result { + // Apply parameters to the workflow + let resolved_workflow = apply_parameters(template.workflow(), ¶meters)?; + + Ok(Self { + template, + parameters, + resolved_workflow, + }) + } + + /// Gets the original template + pub fn template(&self) -> &WorkflowTemplate { + &self.template + } + + /// Gets the applied parameters + pub fn parameters(&self) -> &ParameterValues { + &self.parameters + } + + /// Gets the resolved workflow with parameters applied + pub fn workflow(&self) -> &ComfyUIWorkflow { + &self.resolved_workflow + } + + /// Gets the template ID + pub fn template_id(&self) -> &str { + self.template.id() + } + + /// Gets the template name + pub fn template_name(&self) -> &str { + self.template.name() + } + + /// Gets a specific parameter value + pub fn get_parameter(&self, name: &str) -> Option<&serde_json::Value> { + self.parameters.get(name) + } + + /// Updates a parameter value and re-resolves the workflow + pub fn set_parameter(&mut self, name: String, value: serde_json::Value) -> Result<()> { + // Check if parameter exists in template + if !self.template.has_parameter(&name) { + return Err(ComfyUIError::parameter_validation( + format!("Parameter '{}' does not exist in template", name) + )); + } + + // Update parameter + self.parameters.insert(name, value); + + // Validate updated parameters + let validation = self.template.validate(&self.parameters); + if !validation.valid { + let error_messages: Vec = validation.errors + .iter() + .map(|e| format!("{}: {}", e.path, e.message)) + .collect(); + return Err(ComfyUIError::parameter_validation( + format!("Invalid parameters after update: {}", error_messages.join(", ")) + )); + } + + // Re-resolve workflow + self.resolved_workflow = apply_parameters(self.template.workflow(), &self.parameters)?; + + Ok(()) + } + + /// Updates multiple parameters at once + pub fn set_parameters(&mut self, parameters: HashMap) -> Result<()> { + // Check if all parameters exist in template + for name in parameters.keys() { + if !self.template.has_parameter(name) { + return Err(ComfyUIError::parameter_validation( + format!("Parameter '{}' does not exist in template", name) + )); + } + } + + // Update parameters + for (name, value) in parameters { + self.parameters.insert(name, value); + } + + // Validate updated parameters + let validation = self.template.validate(&self.parameters); + if !validation.valid { + let error_messages: Vec = validation.errors + .iter() + .map(|e| format!("{}: {}", e.path, e.message)) + .collect(); + return Err(ComfyUIError::parameter_validation( + format!("Invalid parameters after update: {}", error_messages.join(", ")) + )); + } + + // Re-resolve workflow + self.resolved_workflow = apply_parameters(self.template.workflow(), &self.parameters)?; + + Ok(()) + } + + /// Gets a specific node from the resolved workflow + pub fn get_node(&self, node_id: &str) -> Option<&ComfyUINode> { + self.resolved_workflow.get(node_id) + } + + /// Gets all node IDs in the workflow + pub fn get_node_ids(&self) -> Vec<&str> { + self.resolved_workflow.keys().map(|s| s.as_str()).collect() + } + + /// Checks if the workflow contains a specific node + pub fn has_node(&self, node_id: &str) -> bool { + self.resolved_workflow.contains_key(node_id) + } + + /// Gets the number of nodes in the workflow + pub fn node_count(&self) -> usize { + self.resolved_workflow.len() + } + + /// Converts the instance to a JSON-serializable workflow + pub fn to_workflow_json(&self) -> Result { + serde_json::to_value(&self.resolved_workflow) + .map_err(ComfyUIError::from) + } + + /// Creates a clone with different parameters + pub fn with_parameters(&self, parameters: ParameterValues) -> Result { + Self::new(self.template.clone(), parameters) + } + + /// Validates the current instance + pub fn validate(&self) -> Result<()> { + let validation = self.template.validate(&self.parameters); + if !validation.valid { + let error_messages: Vec = validation.errors + .iter() + .map(|e| format!("{}: {}", e.path, e.message)) + .collect(); + return Err(ComfyUIError::parameter_validation( + format!("Instance validation failed: {}", error_messages.join(", ")) + )); + } + Ok(()) + } +} diff --git a/cargos/comfyui-sdk/templates/workflow_template.rs b/cargos/comfyui-sdk/templates/workflow_template.rs new file mode 100644 index 0000000..cc1c3db --- /dev/null +++ b/cargos/comfyui-sdk/templates/workflow_template.rs @@ -0,0 +1,193 @@ +//! WorkflowTemplate implementation for managing parameterized ComfyUI workflows + +use std::collections::HashMap; +use crate::types::{ + WorkflowTemplateData, TemplateMetadata, ComfyUIWorkflow, + ParameterSchema, ParameterValues, ValidationResult, ParameterType +}; +use crate::error::{ComfyUIError, Result}; +use crate::templates::WorkflowInstance; +use crate::utils::validation::validate_parameters; + +/// Represents a parameterized workflow template +#[derive(Debug, Clone)] +pub struct WorkflowTemplate { + metadata: TemplateMetadata, + workflow: ComfyUIWorkflow, + parameters: HashMap, +} + +impl WorkflowTemplate { + /// Creates a new WorkflowTemplate instance + pub fn new(data: WorkflowTemplateData) -> Result { + let template = Self { + metadata: data.metadata, + workflow: data.workflow, + parameters: data.parameters, + }; + + // Validate template structure + template.validate_template()?; + Ok(template) + } + + /// Gets the template metadata + pub fn metadata(&self) -> &TemplateMetadata { + &self.metadata + } + + /// Gets the workflow definition + pub fn workflow(&self) -> &ComfyUIWorkflow { + &self.workflow + } + + /// Gets the parameter schemas + pub fn parameters(&self) -> &HashMap { + &self.parameters + } + + /// Gets the template ID + pub fn id(&self) -> &str { + &self.metadata.id + } + + /// Gets the template name + pub fn name(&self) -> &str { + &self.metadata.name + } + + /// Gets the template description + pub fn description(&self) -> Option<&str> { + self.metadata.description.as_deref() + } + + /// Gets the template category + pub fn category(&self) -> Option<&str> { + self.metadata.category.as_deref() + } + + /// Gets the template tags + pub fn tags(&self) -> &[String] { + self.metadata.tags.as_deref().unwrap_or(&[]) + } + + /// Validates the provided parameters against the template schema + pub fn validate(&self, parameters: &ParameterValues) -> ValidationResult { + validate_parameters(parameters, &self.parameters) + } + + /// Creates a workflow instance with the provided parameters + pub fn create_instance(&self, parameters: ParameterValues) -> Result { + let validation = self.validate(¶meters); + if !validation.valid { + let error_messages: Vec = validation.errors + .iter() + .map(|e| format!("{}: {}", e.path, e.message)) + .collect(); + return Err(ComfyUIError::parameter_validation( + format!("Invalid parameters: {}", error_messages.join(", ")) + )); + } + + WorkflowInstance::new(self.clone(), parameters) + } + + /// Gets the parameter schema for a specific parameter + pub fn get_parameter_schema(&self, parameter_name: &str) -> Option<&ParameterSchema> { + self.parameters.get(parameter_name) + } + + /// Gets all required parameter names + pub fn get_required_parameters(&self) -> Vec<&str> { + self.parameters + .iter() + .filter(|(_, schema)| schema.required.unwrap_or(false)) + .map(|(name, _)| name.as_str()) + .collect() + } + + /// Gets all optional parameter names + pub fn get_optional_parameters(&self) -> Vec<&str> { + self.parameters + .iter() + .filter(|(_, schema)| !schema.required.unwrap_or(false)) + .map(|(name, _)| name.as_str()) + .collect() + } + + /// Checks if the template has a specific parameter + pub fn has_parameter(&self, parameter_name: &str) -> bool { + self.parameters.contains_key(parameter_name) + } + + /// Gets default values for all parameters that have defaults + pub fn get_default_values(&self) -> ParameterValues { + let mut defaults = HashMap::new(); + for (name, schema) in &self.parameters { + if let Some(default_value) = &schema.default { + defaults.insert(name.clone(), default_value.clone()); + } + } + defaults + } + + /// Creates a copy of the template with updated metadata + pub fn with_metadata(&self, metadata: TemplateMetadata) -> Self { + Self { + metadata, + workflow: self.workflow.clone(), + parameters: self.parameters.clone(), + } + } + + /// Converts the template to a JSON-serializable object + pub fn to_data(&self) -> WorkflowTemplateData { + WorkflowTemplateData { + metadata: self.metadata.clone(), + workflow: self.workflow.clone(), + parameters: self.parameters.clone(), + } + } + + /// Creates a WorkflowTemplate from template data + pub fn from_data(data: WorkflowTemplateData) -> Result { + Self::new(data) + } + + /// Validates the template structure + fn validate_template(&self) -> Result<()> { + // Validate metadata + if self.metadata.id.is_empty() { + return Err(ComfyUIError::template_validation( + "Template metadata must have a valid id" + )); + } + if self.metadata.name.is_empty() { + return Err(ComfyUIError::template_validation( + "Template metadata must have a valid name" + )); + } + + // Validate workflow + if self.workflow.is_empty() { + return Err(ComfyUIError::template_validation( + "Template must have a valid workflow object" + )); + } + + // Validate parameter schemas + for (name, schema) in &self.parameters { + if !matches!( + schema.param_type, + ParameterType::String | ParameterType::Number | ParameterType::Boolean | + ParameterType::Array | ParameterType::Object + ) { + return Err(ComfyUIError::template_validation( + format!("Parameter '{}' has invalid type: {:?}", name, schema.param_type) + )); + } + } + + Ok(()) + } +} diff --git a/cargos/comfyui-sdk/tests/url_generation_test.rs b/cargos/comfyui-sdk/tests/url_generation_test.rs new file mode 100644 index 0000000..c849f95 --- /dev/null +++ b/cargos/comfyui-sdk/tests/url_generation_test.rs @@ -0,0 +1,188 @@ +//! URL generation tests + +use std::collections::HashMap; +use serde_json::json; +use comfyui_sdk::{HTTPClient, ComfyUIClientConfig}; + +#[test] +fn test_outputs_to_urls_no_double_slash() { + // Test with base URL ending with slash + let config_with_slash = ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188/".to_string(), + ..Default::default() + }; + let client_with_slash = HTTPClient::new(config_with_slash).unwrap(); + + // Test with base URL without ending slash + let config_without_slash = ComfyUIClientConfig { + base_url: "http://192.168.0.193:8188".to_string(), + ..Default::default() + }; + let client_without_slash = HTTPClient::new(config_without_slash).unwrap(); + + // Create test outputs + let mut outputs = HashMap::new(); + outputs.insert("58".to_string(), json!({ + "images": [ + { + "filename": "ComfyUI_02046_.png", + "subfolder": "", + "type": "output" + } + ] + })); + + // Test with base URL ending with slash + let urls_with_slash = client_with_slash.outputs_to_urls(&outputs); + assert_eq!(urls_with_slash.len(), 1); + let url_with_slash = &urls_with_slash[0]; + println!("URL with slash: {}", url_with_slash); + + // Should not contain double slash + assert!(!url_with_slash.contains("//view"), + "URL should not contain double slash: {}", url_with_slash); + assert!(url_with_slash.contains("/view?"), + "URL should contain single slash before view: {}", url_with_slash); + + // Test with base URL without ending slash + let urls_without_slash = client_without_slash.outputs_to_urls(&outputs); + assert_eq!(urls_without_slash.len(), 1); + let url_without_slash = &urls_without_slash[0]; + println!("URL without slash: {}", url_without_slash); + + // Should not contain double slash + assert!(!url_without_slash.contains("//view"), + "URL should not contain double slash: {}", url_without_slash); + assert!(url_without_slash.contains("/view?"), + "URL should contain single slash before view: {}", url_without_slash); + + // Both URLs should be the same + assert_eq!(url_with_slash, url_without_slash, + "URLs should be identical regardless of trailing slash in base URL"); + + // Expected URL format + let expected_url = "http://192.168.0.193:8188/view?filename=ComfyUI_02046_.png&subfolder=&type=output"; + assert_eq!(url_with_slash, expected_url); + assert_eq!(url_without_slash, expected_url); +} + +#[test] +fn test_outputs_to_urls_multiple_images() { + let config = ComfyUIClientConfig { + base_url: "http://localhost:8188/".to_string(), + ..Default::default() + }; + let client = HTTPClient::new(config).unwrap(); + + let mut outputs = HashMap::new(); + outputs.insert("58".to_string(), json!({ + "images": [ + { + "filename": "image1.png", + "subfolder": "temp", + "type": "output" + }, + { + "filename": "image2.jpg", + "subfolder": "", + "type": "temp" + } + ] + })); + + let urls = client.outputs_to_urls(&outputs); + assert_eq!(urls.len(), 2); + + for url in &urls { + println!("Generated URL: {}", url); + assert!(!url.contains("//view"), "URL should not contain double slash: {}", url); + assert!(url.contains("/view?"), "URL should contain single slash before view: {}", url); + assert!(url.starts_with("http://localhost:8188/view?"), "URL should have correct base"); + } + + // Check specific URLs + assert!(urls.contains(&"http://localhost:8188/view?filename=image1.png&subfolder=temp&type=output".to_string())); + assert!(urls.contains(&"http://localhost:8188/view?filename=image2.jpg&subfolder=&type=temp".to_string())); +} + +#[test] +fn test_outputs_to_urls_empty_outputs() { + let config = ComfyUIClientConfig { + base_url: "http://localhost:8188".to_string(), + ..Default::default() + }; + let client = HTTPClient::new(config).unwrap(); + + let outputs = HashMap::new(); + let urls = client.outputs_to_urls(&outputs); + assert_eq!(urls.len(), 0); +} + +#[test] +fn test_outputs_to_urls_no_images() { + let config = ComfyUIClientConfig { + base_url: "http://localhost:8188".to_string(), + ..Default::default() + }; + let client = HTTPClient::new(config).unwrap(); + + let mut outputs = HashMap::new(); + outputs.insert("58".to_string(), json!({ + "text": "some text output" + })); + + let urls = client.outputs_to_urls(&outputs); + assert_eq!(urls.len(), 0); +} + +#[test] +fn test_different_base_url_formats() { + let test_cases = vec![ + "http://192.168.0.193:8188", + "http://192.168.0.193:8188/", + "https://my-comfyui.com", + "https://my-comfyui.com/", + "http://localhost:8188", + "http://localhost:8188/", + ]; + + let mut outputs = HashMap::new(); + outputs.insert("test".to_string(), json!({ + "images": [ + { + "filename": "test.png", + "subfolder": "", + "type": "output" + } + ] + })); + + for base_url in test_cases { + println!("Testing base URL: {}", base_url); + + let config = ComfyUIClientConfig { + base_url: base_url.to_string(), + ..Default::default() + }; + let client = HTTPClient::new(config).unwrap(); + + let urls = client.outputs_to_urls(&outputs); + assert_eq!(urls.len(), 1); + + let url = &urls[0]; + println!("Generated URL: {}", url); + + // Should not contain double slash + assert!(!url.contains("//view"), + "URL should not contain double slash for base URL '{}': {}", base_url, url); + + // Should contain single slash before view + assert!(url.contains("/view?"), + "URL should contain single slash before view for base URL '{}': {}", base_url, url); + + // Should start with the base URL (without trailing slash) + /view + let expected_base = base_url.trim_end_matches('/'); + assert!(url.starts_with(&format!("{}/view?", expected_base)), + "URL should start with correct base for '{}': {}", base_url, url); + } +} diff --git a/cargos/comfyui-sdk/types/api.rs b/cargos/comfyui-sdk/types/api.rs new file mode 100644 index 0000000..716b87e --- /dev/null +++ b/cargos/comfyui-sdk/types/api.rs @@ -0,0 +1,133 @@ +//! API types for ComfyUI SDK + +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +/// Queue status response +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueStatus { + pub queue_running: Vec, + pub queue_pending: Vec, +} + +/// Queue item +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueueItem { + pub prompt_id: String, + pub number: u32, + pub outputs: HashMap, +} + +/// Prompt request +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PromptRequest { + pub prompt: HashMap, + #[serde(skip_serializing_if = "Option::is_none")] + pub client_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extra_data: Option>, +} + +/// Prompt response +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PromptResponse { + pub prompt_id: String, + pub number: u32, + pub node_errors: HashMap, +} + +/// History item +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HistoryItem { + pub prompt: Vec, + pub outputs: HashMap>>, + pub status: HistoryStatus, +} + +/// History output +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HistoryOutput { + pub filename: String, + pub subfolder: String, + #[serde(rename = "type")] + pub output_type: String, +} + +/// History status +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HistoryStatus { + pub status_str: String, + pub completed: bool, + pub messages: Vec>, +} + +/// System stats +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SystemStats { + pub system: SystemInfo, + pub devices: Vec, +} + +/// System information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SystemInfo { + pub os: String, + pub python_version: String, + pub embedded_python: bool, +} + +/// Device information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DeviceInfo { + pub name: String, + pub device_type: String, + pub vram_total: u64, + pub vram_free: u64, + pub torch_vram_total: u64, + pub torch_vram_free: u64, +} + +/// Object info response +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ObjectInfo { + #[serde(flatten)] + pub nodes: HashMap, +} + +/// Node information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeInfo { + pub input: NodeInputInfo, + pub output: Vec, + pub output_is_list: Vec, + pub output_name: Vec, + pub name: String, + pub display_name: String, + pub description: String, + pub category: String, + pub python_module: String, +} + +/// Node input information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeInputInfo { + pub required: HashMap>, + #[serde(skip_serializing_if = "Option::is_none")] + pub optional: Option>>, +} + +/// Upload image response +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UploadImageResponse { + pub name: String, + pub subfolder: String, + #[serde(rename = "type")] + pub image_type: String, +} + +/// View metadata response +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ViewMetadata { + #[serde(flatten)] + pub metadata: HashMap, +} diff --git a/cargos/comfyui-sdk/types/events.rs b/cargos/comfyui-sdk/types/events.rs new file mode 100644 index 0000000..4165a47 --- /dev/null +++ b/cargos/comfyui-sdk/types/events.rs @@ -0,0 +1,160 @@ +//! Event types for ComfyUI SDK + +use serde::{Deserialize, Serialize}; +use chrono::{DateTime, Utc}; +use std::collections::HashMap; + +/// Execution progress information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecutionProgress { + pub node_id: String, + pub progress: u32, + pub max: u32, + pub timestamp: DateTime, +} + +/// Execution result information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecutionResult { + pub prompt_id: String, + pub outputs: HashMap, + pub execution_time: u64, + pub timestamp: DateTime, +} + +/// Execution error information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecutionError { + pub node_id: Option, + pub message: String, + pub details: Option, + pub timestamp: DateTime, +} + +/// Callback trait for execution events +pub trait ExecutionCallbacks: Send + Sync { + /// Called when progress is updated + fn on_progress(&self, _progress: ExecutionProgress) {} + + /// Called when a node starts executing + fn on_executing(&self, _node_id: String) {} + + /// Called when execution completes successfully + fn on_executed(&self, _result: ExecutionResult) {} + + /// Called when an error occurs + fn on_error(&self, _error: ExecutionError) {} +} + +/// WebSocket message types +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum WSMessage { + #[serde(rename = "progress")] + Progress { + data: WSProgressData, + }, + #[serde(rename = "executing")] + Executing { + data: WSExecutingData, + }, + #[serde(rename = "executed")] + Executed { + data: WSExecutedData, + }, + #[serde(rename = "execution_error")] + ExecutionError { + data: WSErrorData, + }, + #[serde(other)] + Unknown, +} + +/// WebSocket progress message data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WSProgressData { + pub value: u32, + pub max: u32, + pub node: Option, +} + +/// WebSocket executing message data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WSExecutingData { + pub node: Option, + pub prompt_id: String, +} + +/// WebSocket executed message data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WSExecutedData { + pub node: String, + pub output: HashMap, + pub prompt_id: String, +} + +/// WebSocket error message data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WSErrorData { + pub node_id: Option, + pub message: String, + pub details: Option, +} + +/// Execution options +#[derive(Debug, Clone)] +pub struct ExecutionOptions { + pub timeout: Option, + pub priority: Option, +} + +impl Default for ExecutionOptions { + fn default() -> Self { + Self { + timeout: Some(std::time::Duration::from_secs(300)), // 5 minutes default + priority: None, + } + } +} + +/// Template execution result +#[derive(Debug, Clone)] +pub struct TemplateExecutionResult { + pub prompt_id: String, + pub success: bool, + pub outputs: Option>, + pub error: Option, + pub execution_time: u64, +} + +impl TemplateExecutionResult { + /// Create a successful execution result + pub fn success( + prompt_id: String, + outputs: HashMap, + execution_time: u64, + ) -> Self { + Self { + prompt_id, + success: true, + outputs: Some(outputs), + error: None, + execution_time, + } + } + + /// Create a failed execution result + pub fn failure( + prompt_id: String, + error: ExecutionError, + execution_time: u64, + ) -> Self { + Self { + prompt_id, + success: false, + outputs: None, + error: Some(error), + execution_time, + } + } +} diff --git a/cargos/comfyui-sdk/types/mod.rs b/cargos/comfyui-sdk/types/mod.rs new file mode 100644 index 0000000..e3b9cd0 --- /dev/null +++ b/cargos/comfyui-sdk/types/mod.rs @@ -0,0 +1,192 @@ +//! Core types for ComfyUI SDK + +use std::collections::HashMap; +use serde::{Deserialize, Serialize}; +use chrono::{DateTime, Utc}; + +pub mod api; +pub mod events; + +// Re-export all types +pub use api::*; +pub use events::*; + +/// Base ComfyUI node definition +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ComfyUINode { + pub class_type: String, + pub inputs: HashMap, + #[serde(skip_serializing_if = "Option::is_none")] + pub _meta: Option, +} + +/// Node metadata +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeMeta { + #[serde(skip_serializing_if = "Option::is_none")] + pub title: Option, +} + +/// ComfyUI workflow definition +pub type ComfyUIWorkflow = HashMap; + +/// Parameter schema types +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum ParameterType { + String, + Number, + Boolean, + Array, + Object, +} + +/// Parameter schema definition +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ParameterSchema { + #[serde(rename = "type")] + pub param_type: ParameterType, + #[serde(skip_serializing_if = "Option::is_none")] + pub required: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub default: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub r#enum: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub min: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub pattern: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub items: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub properties: Option>, +} + +/// Template metadata +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TemplateMetadata { + pub id: String, + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub version: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub author: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tags: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub category: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub created_at: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub updated_at: Option>, +} + +/// Workflow template data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WorkflowTemplateData { + pub metadata: TemplateMetadata, + pub workflow: ComfyUIWorkflow, + pub parameters: HashMap, +} + +/// Client configuration +#[derive(Debug, Clone)] +pub struct ComfyUIClientConfig { + pub base_url: String, + pub timeout: Option, + pub retry_attempts: Option, + pub retry_delay: Option, + pub headers: Option>, +} + +impl Default for ComfyUIClientConfig { + fn default() -> Self { + Self { + base_url: "http://localhost:8188".to_string(), + timeout: Some(std::time::Duration::from_secs(30)), + retry_attempts: Some(3), + retry_delay: Some(std::time::Duration::from_millis(1000)), + headers: None, + } + } +} + +/// Parameter values type +pub type ParameterValues = HashMap; + +/// Variable syntax options +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum VariableSyntax { + DoubleBrace, // {{}} + DollarBrace, // ${} + AtBrace, // @{} +} + +/// Template validation result +#[derive(Debug, Clone)] +pub struct ValidationResult { + pub valid: bool, + pub errors: Vec, +} + +/// Validation error +#[derive(Debug, Clone)] +pub struct ValidationError { + pub path: String, + pub message: String, + pub value: Option, +} + +impl ValidationResult { + /// Create a successful validation result + pub fn success() -> Self { + Self { + valid: true, + errors: Vec::new(), + } + } + + /// Create a failed validation result with errors + pub fn failure(errors: Vec) -> Self { + Self { + valid: false, + errors, + } + } + + /// Add an error to the validation result + pub fn add_error(&mut self, error: ValidationError) { + self.valid = false; + self.errors.push(error); + } +} + +impl ValidationError { + /// Create a new validation error + pub fn new(path: impl Into, message: impl Into) -> Self { + Self { + path: path.into(), + message: message.into(), + value: None, + } + } + + /// Create a new validation error with value + pub fn with_value( + path: impl Into, + message: impl Into, + value: serde_json::Value, + ) -> Self { + Self { + path: path.into(), + message: message.into(), + value: Some(value), + } + } +} diff --git a/cargos/comfyui-sdk/utils/error_handling.rs b/cargos/comfyui-sdk/utils/error_handling.rs new file mode 100644 index 0000000..3100bda --- /dev/null +++ b/cargos/comfyui-sdk/utils/error_handling.rs @@ -0,0 +1,216 @@ +//! Error handling utilities + +use crate::error::{ComfyUIError, Result}; +use std::time::Duration; +use tokio::time::sleep; + +/// Retry configuration +#[derive(Debug, Clone)] +pub struct RetryConfig { + pub max_attempts: u32, + pub initial_delay: Duration, + pub max_delay: Duration, + pub backoff_multiplier: f64, +} + +impl Default for RetryConfig { + fn default() -> Self { + Self { + max_attempts: 3, + initial_delay: Duration::from_millis(1000), + max_delay: Duration::from_secs(30), + backoff_multiplier: 2.0, + } + } +} + +/// Retries an async operation with exponential backoff +pub async fn retry_with_backoff( + operation: F, + config: RetryConfig, +) -> Result +where + F: Fn() -> Fut, + Fut: std::future::Future>, +{ + let mut last_error = None; + let mut delay = config.initial_delay; + + for attempt in 1..=config.max_attempts { + match operation().await { + Ok(result) => return Ok(result), + Err(error) => { + last_error = Some(error); + + if attempt < config.max_attempts { + log::warn!("Attempt {} failed, retrying in {:?}", attempt, delay); + sleep(delay).await; + + // Exponential backoff + delay = std::cmp::min( + Duration::from_millis( + (delay.as_millis() as f64 * config.backoff_multiplier) as u64 + ), + config.max_delay, + ); + } + } + } + } + + Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error"))) +} + +/// Checks if an error is retryable +pub fn is_retryable_error(error: &ComfyUIError) -> bool { + match error { + ComfyUIError::Http(reqwest_error) => { + // Retry on network errors, timeouts, and 5xx status codes + reqwest_error.is_timeout() || + reqwest_error.is_connect() || + reqwest_error.status().map_or(false, |status| status.is_server_error()) + } + ComfyUIError::WebSocket(_) => true, // Most WebSocket errors are retryable + ComfyUIError::Connection(_) => true, + ComfyUIError::Timeout(_) => true, + ComfyUIError::Io(_) => true, + _ => false, // Don't retry validation errors, etc. + } +} + +/// Retries an operation only if the error is retryable +pub async fn retry_if_retryable( + operation: F, + config: RetryConfig, +) -> Result +where + F: Fn() -> Fut, + Fut: std::future::Future>, +{ + let mut last_error = None; + let mut delay = config.initial_delay; + + for attempt in 1..=config.max_attempts { + match operation().await { + Ok(result) => return Ok(result), + Err(error) => { + if !is_retryable_error(&error) { + return Err(error); + } + + last_error = Some(error); + + if attempt < config.max_attempts { + log::warn!("Attempt {} failed with retryable error, retrying in {:?}", attempt, delay); + sleep(delay).await; + + // Exponential backoff + delay = std::cmp::min( + Duration::from_millis( + (delay.as_millis() as f64 * config.backoff_multiplier) as u64 + ), + config.max_delay, + ); + } + } + } + } + + Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error"))) +} + +/// Wraps an operation with timeout +pub async fn with_timeout( + operation: F, + timeout: Duration, +) -> Result +where + F: std::future::Future>, +{ + match tokio::time::timeout(timeout, operation).await { + Ok(result) => result, + Err(_) => Err(ComfyUIError::timeout(format!( + "Operation timed out after {:?}", + timeout + ))), + } +} + +/// Error context helper +pub trait ErrorContext { + fn with_context(self, context: &str) -> Result; +} + +impl ErrorContext for Result { + fn with_context(self, context: &str) -> Result { + self.map_err(|error| { + ComfyUIError::new(format!("{}: {}", context, error)) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::Arc; + + #[tokio::test] + async fn test_retry_success_on_second_attempt() { + let counter = Arc::new(AtomicU32::new(0)); + let counter_clone = counter.clone(); + + let config = RetryConfig { + max_attempts: 3, + initial_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(100), + backoff_multiplier: 2.0, + }; + + let result = retry_with_backoff( + || { + let counter = counter_clone.clone(); + async move { + let count = counter.fetch_add(1, Ordering::SeqCst); + if count == 0 { + Err(ComfyUIError::connection("First attempt fails")) + } else { + Ok("Success") + } + } + }, + config, + ).await; + + assert!(result.is_ok()); + assert_eq!(result.unwrap(), "Success"); + assert_eq!(counter.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn test_retry_exhausted() { + let counter = Arc::new(AtomicU32::new(0)); + let counter_clone = counter.clone(); + + let config = RetryConfig { + max_attempts: 2, + initial_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(100), + backoff_multiplier: 2.0, + }; + + let result: Result<&str> = retry_with_backoff( + || { + let counter = counter_clone.clone(); + async move { + counter.fetch_add(1, Ordering::SeqCst); + Err(ComfyUIError::connection("Always fails")) + } + }, + config, + ).await; + + assert!(result.is_err()); + assert_eq!(counter.load(Ordering::SeqCst), 2); + } +} diff --git a/cargos/comfyui-sdk/utils/event_emitter.rs b/cargos/comfyui-sdk/utils/event_emitter.rs new file mode 100644 index 0000000..9b696af --- /dev/null +++ b/cargos/comfyui-sdk/utils/event_emitter.rs @@ -0,0 +1,171 @@ +//! Event emitter utilities + +use std::sync::Arc; +use tokio::sync::RwLock; +use std::collections::HashMap; +use crate::types::{ExecutionProgress, ExecutionResult, ExecutionError, ExecutionCallbacks}; + +/// Event emitter for execution events +#[derive(Clone)] +pub struct EventEmitter { + callbacks: Arc>>>, +} + +impl EventEmitter { + /// Creates a new event emitter + pub fn new() -> Self { + Self { + callbacks: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Registers callbacks with a unique ID + pub async fn register_callbacks(&self, id: String, callbacks: Arc) { + let mut callbacks_map = self.callbacks.write().await; + callbacks_map.insert(id, callbacks); + } + + /// Unregisters callbacks by ID + pub async fn unregister_callbacks(&self, id: &str) { + let mut callbacks_map = self.callbacks.write().await; + callbacks_map.remove(id); + } + + /// Emits a progress event + pub async fn emit_progress(&self, progress: ExecutionProgress) { + let callbacks_map = self.callbacks.read().await; + for callback in callbacks_map.values() { + callback.on_progress(progress.clone()); + } + } + + /// Emits an executing event + pub async fn emit_executing(&self, node_id: String) { + let callbacks_map = self.callbacks.read().await; + for callback in callbacks_map.values() { + callback.on_executing(node_id.clone()); + } + } + + /// Emits an executed event + pub async fn emit_executed(&self, result: ExecutionResult) { + let callbacks_map = self.callbacks.read().await; + for callback in callbacks_map.values() { + callback.on_executed(result.clone()); + } + } + + /// Emits an error event + pub async fn emit_error(&self, error: ExecutionError) { + let callbacks_map = self.callbacks.read().await; + for callback in callbacks_map.values() { + callback.on_error(error.clone()); + } + } + + /// Gets the number of registered callbacks + pub async fn callback_count(&self) -> usize { + let callbacks_map = self.callbacks.read().await; + callbacks_map.len() + } + + /// Clears all callbacks + pub async fn clear(&self) { + let mut callbacks_map = self.callbacks.write().await; + callbacks_map.clear(); + } +} + +impl Default for EventEmitter { + fn default() -> Self { + Self::new() + } +} + +/// Simple callback implementation for testing +pub struct SimpleCallbacks { + pub on_progress_fn: Option>, + pub on_executing_fn: Option>, + pub on_executed_fn: Option>, + pub on_error_fn: Option>, +} + +impl SimpleCallbacks { + /// Creates new simple callbacks + pub fn new() -> Self { + Self { + on_progress_fn: None, + on_executing_fn: None, + on_executed_fn: None, + on_error_fn: None, + } + } + + /// Sets the progress callback + pub fn with_progress(mut self, f: F) -> Self + where + F: Fn(ExecutionProgress) + Send + Sync + 'static, + { + self.on_progress_fn = Some(Box::new(f)); + self + } + + /// Sets the executing callback + pub fn with_executing(mut self, f: F) -> Self + where + F: Fn(String) + Send + Sync + 'static, + { + self.on_executing_fn = Some(Box::new(f)); + self + } + + /// Sets the executed callback + pub fn with_executed(mut self, f: F) -> Self + where + F: Fn(ExecutionResult) + Send + Sync + 'static, + { + self.on_executed_fn = Some(Box::new(f)); + self + } + + /// Sets the error callback + pub fn with_error(mut self, f: F) -> Self + where + F: Fn(ExecutionError) + Send + Sync + 'static, + { + self.on_error_fn = Some(Box::new(f)); + self + } +} + +impl ExecutionCallbacks for SimpleCallbacks { + fn on_progress(&self, progress: ExecutionProgress) { + if let Some(ref f) = self.on_progress_fn { + f(progress); + } + } + + fn on_executing(&self, node_id: String) { + if let Some(ref f) = self.on_executing_fn { + f(node_id); + } + } + + fn on_executed(&self, result: ExecutionResult) { + if let Some(ref f) = self.on_executed_fn { + f(result); + } + } + + fn on_error(&self, error: ExecutionError) { + if let Some(ref f) = self.on_error_fn { + f(error); + } + } +} + +impl Default for SimpleCallbacks { + fn default() -> Self { + Self::new() + } +} diff --git a/cargos/comfyui-sdk/utils/mod.rs b/cargos/comfyui-sdk/utils/mod.rs new file mode 100644 index 0000000..5795224 --- /dev/null +++ b/cargos/comfyui-sdk/utils/mod.rs @@ -0,0 +1,12 @@ +//! Utility functions for ComfyUI SDK + +pub mod validation; +pub mod template_parser; +pub mod event_emitter; +pub mod error_handling; + +// Re-export main utilities +pub use validation::*; +pub use template_parser::*; +pub use event_emitter::*; +pub use error_handling::*; diff --git a/cargos/comfyui-sdk/utils/template_parser.rs b/cargos/comfyui-sdk/utils/template_parser.rs new file mode 100644 index 0000000..74ae2d5 --- /dev/null +++ b/cargos/comfyui-sdk/utils/template_parser.rs @@ -0,0 +1,256 @@ +//! Template parsing utilities for parameter substitution + +use std::collections::HashMap; +use regex::Regex; +use crate::types::{ComfyUIWorkflow, ParameterValues, ComfyUINode, VariableSyntax}; +use crate::error::{ComfyUIError, Result}; + +/// Applies parameters to a workflow template +pub fn apply_parameters( + workflow: &ComfyUIWorkflow, + parameters: &ParameterValues, +) -> Result { + let mut resolved_workflow = HashMap::new(); + + for (node_id, node) in workflow { + let resolved_node = apply_parameters_to_node(node, parameters)?; + resolved_workflow.insert(node_id.clone(), resolved_node); + } + + Ok(resolved_workflow) +} + +/// Applies parameters to a single node +fn apply_parameters_to_node( + node: &ComfyUINode, + parameters: &ParameterValues, +) -> Result { + let mut resolved_inputs = HashMap::new(); + + for (input_name, input_value) in &node.inputs { + let resolved_value = apply_parameters_to_value(input_value, parameters)?; + resolved_inputs.insert(input_name.clone(), resolved_value); + } + + Ok(ComfyUINode { + class_type: node.class_type.clone(), + inputs: resolved_inputs, + _meta: node._meta.clone(), + }) +} + +/// Applies parameters to a JSON value +fn apply_parameters_to_value( + value: &serde_json::Value, + parameters: &ParameterValues, +) -> Result { + match value { + serde_json::Value::String(s) => { + let resolved_string = substitute_variables(s, parameters, VariableSyntax::DoubleBrace)?; + Ok(serde_json::Value::String(resolved_string)) + } + serde_json::Value::Array(arr) => { + let mut resolved_array = Vec::new(); + for item in arr { + resolved_array.push(apply_parameters_to_value(item, parameters)?); + } + Ok(serde_json::Value::Array(resolved_array)) + } + serde_json::Value::Object(obj) => { + let mut resolved_object = serde_json::Map::new(); + for (key, val) in obj { + let resolved_key = substitute_variables(key, parameters, VariableSyntax::DoubleBrace)?; + let resolved_val = apply_parameters_to_value(val, parameters)?; + resolved_object.insert(resolved_key, resolved_val); + } + Ok(serde_json::Value::Object(resolved_object)) + } + _ => Ok(value.clone()), + } +} + +/// Substitutes variables in a string using the specified syntax +pub fn substitute_variables( + template: &str, + parameters: &ParameterValues, + syntax: VariableSyntax, +) -> Result { + let pattern = match syntax { + VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}", + VariableSyntax::DollarBrace => r"\$\{([^}]+)\}", + VariableSyntax::AtBrace => r"@\{([^}]+)\}", + }; + + let regex = Regex::new(pattern) + .map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?; + + let mut result = template.to_string(); + let mut offset = 0i32; + + for captures in regex.captures_iter(template) { + let full_match = captures.get(0).unwrap(); + let var_name = captures.get(1).unwrap().as_str().trim(); + + // Get parameter value + let replacement = match parameters.get(var_name) { + Some(value) => value_to_string(value)?, + None => { + return Err(ComfyUIError::template_validation( + format!("Parameter '{}' not found", var_name) + )); + } + }; + + // Calculate positions with offset + let start = (full_match.start() as i32 + offset) as usize; + let end = (full_match.end() as i32 + offset) as usize; + + // Replace the variable + result.replace_range(start..end, &replacement); + + // Update offset + offset += replacement.len() as i32 - full_match.len() as i32; + } + + Ok(result) +} + +/// Converts a JSON value to string for substitution +fn value_to_string(value: &serde_json::Value) -> Result { + match value { + serde_json::Value::String(s) => Ok(s.clone()), + serde_json::Value::Number(n) => Ok(n.to_string()), + serde_json::Value::Bool(b) => Ok(b.to_string()), + serde_json::Value::Null => Ok("null".to_string()), + _ => { + // For complex types, serialize to JSON + serde_json::to_string(value) + .map_err(ComfyUIError::from) + } + } +} + +/// Extracts variable names from a template string +pub fn extract_variables(template: &str, syntax: VariableSyntax) -> Result> { + let pattern = match syntax { + VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}", + VariableSyntax::DollarBrace => r"\$\{([^}]+)\}", + VariableSyntax::AtBrace => r"@\{([^}]+)\}", + }; + + let regex = Regex::new(pattern) + .map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?; + + let mut variables = Vec::new(); + for captures in regex.captures_iter(template) { + if let Some(var_match) = captures.get(1) { + let var_name = var_match.as_str().trim().to_string(); + if !variables.contains(&var_name) { + variables.push(var_name); + } + } + } + + Ok(variables) +} + +/// Extracts all variables from a workflow +pub fn extract_workflow_variables(workflow: &ComfyUIWorkflow) -> Result> { + let mut all_variables = Vec::new(); + + for node in workflow.values() { + let node_variables = extract_node_variables(node)?; + for var in node_variables { + if !all_variables.contains(&var) { + all_variables.push(var); + } + } + } + + Ok(all_variables) +} + +/// Extracts variables from a single node +fn extract_node_variables(node: &ComfyUINode) -> Result> { + let mut variables = Vec::new(); + + for input_value in node.inputs.values() { + let value_variables = extract_value_variables(input_value)?; + for var in value_variables { + if !variables.contains(&var) { + variables.push(var); + } + } + } + + Ok(variables) +} + +/// Extracts variables from a JSON value +fn extract_value_variables(value: &serde_json::Value) -> Result> { + let mut variables = Vec::new(); + + match value { + serde_json::Value::String(s) => { + let string_vars = extract_variables(s, VariableSyntax::DoubleBrace)?; + variables.extend(string_vars); + } + serde_json::Value::Array(arr) => { + for item in arr { + let item_vars = extract_value_variables(item)?; + variables.extend(item_vars); + } + } + serde_json::Value::Object(obj) => { + for (key, val) in obj { + let key_vars = extract_variables(key, VariableSyntax::DoubleBrace)?; + variables.extend(key_vars); + + let val_vars = extract_value_variables(val)?; + variables.extend(val_vars); + } + } + _ => {} + } + + // Remove duplicates + variables.sort(); + variables.dedup(); + + Ok(variables) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_substitute_variables() { + let mut parameters = HashMap::new(); + parameters.insert("name".to_string(), json!("test")); + parameters.insert("value".to_string(), json!(42)); + + let template = "Hello {{name}}, value is {{value}}"; + let result = substitute_variables(template, ¶meters, VariableSyntax::DoubleBrace).unwrap(); + + assert_eq!(result, "Hello test, value is 42"); + } + + #[test] + fn test_extract_variables() { + let template = "{{var1}} and {{var2}} and {{var1}} again"; + let variables = extract_variables(template, VariableSyntax::DoubleBrace).unwrap(); + + assert_eq!(variables, vec!["var1", "var2"]); + } + + #[test] + fn test_missing_parameter() { + let parameters = HashMap::new(); + let template = "Hello {{missing}}"; + + let result = substitute_variables(template, ¶meters, VariableSyntax::DoubleBrace); + assert!(result.is_err()); + } +} diff --git a/cargos/comfyui-sdk/utils/validation.rs b/cargos/comfyui-sdk/utils/validation.rs new file mode 100644 index 0000000..81511eb --- /dev/null +++ b/cargos/comfyui-sdk/utils/validation.rs @@ -0,0 +1,276 @@ +//! Parameter validation utilities + +use std::collections::HashMap; +use crate::types::{ParameterSchema, ParameterValues, ValidationResult, ValidationError, ParameterType}; +use regex::Regex; + +/// Validates parameters against their schemas +pub fn validate_parameters( + parameters: &ParameterValues, + schemas: &HashMap, +) -> ValidationResult { + let mut errors = Vec::new(); + + // Check required parameters + for (name, schema) in schemas { + if schema.required.unwrap_or(false) { + if !parameters.contains_key(name) { + errors.push(ValidationError::new( + name, + "Required parameter is missing", + )); + continue; + } + } + + // Validate parameter if present + if let Some(value) = parameters.get(name) { + if let Err(error) = validate_parameter_value(name, value, schema) { + errors.push(error); + } + } + } + + // Check for unknown parameters + for name in parameters.keys() { + if !schemas.contains_key(name) { + errors.push(ValidationError::new( + name, + "Unknown parameter", + )); + } + } + + if errors.is_empty() { + ValidationResult::success() + } else { + ValidationResult::failure(errors) + } +} + +/// Validates a single parameter value against its schema +pub fn validate_parameter_value( + name: &str, + value: &serde_json::Value, + schema: &ParameterSchema, +) -> Result<(), ValidationError> { + // Check type + match (&schema.param_type, value) { + (ParameterType::String, serde_json::Value::String(s)) => { + validate_string_constraints(name, s, schema)?; + } + (ParameterType::Number, serde_json::Value::Number(n)) => { + validate_number_constraints(name, n, schema)?; + } + (ParameterType::Boolean, serde_json::Value::Bool(_)) => { + // Boolean values are always valid + } + (ParameterType::Array, serde_json::Value::Array(arr)) => { + validate_array_constraints(name, arr, schema)?; + } + (ParameterType::Object, serde_json::Value::Object(_)) => { + // Object validation could be more complex, but for now just check type + } + _ => { + return Err(ValidationError::with_value( + name, + format!("Expected type {:?}, got {:?}", schema.param_type, get_value_type(value)), + value.clone(), + )); + } + } + + // Check enum constraints + if let Some(enum_values) = &schema.r#enum { + if !enum_values.contains(value) { + return Err(ValidationError::with_value( + name, + format!("Value must be one of: {:?}", enum_values), + value.clone(), + )); + } + } + + Ok(()) +} + +/// Validates string constraints +fn validate_string_constraints( + name: &str, + value: &str, + schema: &ParameterSchema, +) -> Result<(), ValidationError> { + // Check pattern + if let Some(pattern) = &schema.pattern { + let regex = Regex::new(pattern).map_err(|_| { + ValidationError::new(name, "Invalid regex pattern in schema") + })?; + + if !regex.is_match(value) { + return Err(ValidationError::with_value( + name, + format!("String does not match pattern: {}", pattern), + serde_json::Value::String(value.to_string()), + )); + } + } + + // Check length constraints (using min/max as length bounds for strings) + if let Some(min) = schema.min { + if (value.len() as f64) < min { + return Err(ValidationError::with_value( + name, + format!("String length must be at least {}", min), + serde_json::Value::String(value.to_string()), + )); + } + } + + if let Some(max) = schema.max { + if (value.len() as f64) > max { + return Err(ValidationError::with_value( + name, + format!("String length must be at most {}", max), + serde_json::Value::String(value.to_string()), + )); + } + } + + Ok(()) +} + +/// Validates number constraints +fn validate_number_constraints( + name: &str, + value: &serde_json::Number, + schema: &ParameterSchema, +) -> Result<(), ValidationError> { + let num_value = value.as_f64().unwrap_or(0.0); + + if let Some(min) = schema.min { + if num_value < min { + return Err(ValidationError::with_value( + name, + format!("Number must be at least {}", min), + serde_json::Value::Number(value.clone()), + )); + } + } + + if let Some(max) = schema.max { + if num_value > max { + return Err(ValidationError::with_value( + name, + format!("Number must be at most {}", max), + serde_json::Value::Number(value.clone()), + )); + } + } + + Ok(()) +} + +/// Validates array constraints +fn validate_array_constraints( + name: &str, + value: &[serde_json::Value], + schema: &ParameterSchema, +) -> Result<(), ValidationError> { + // Check length constraints + if let Some(min) = schema.min { + if (value.len() as f64) < min { + return Err(ValidationError::with_value( + name, + format!("Array length must be at least {}", min), + serde_json::Value::Array(value.to_vec()), + )); + } + } + + if let Some(max) = schema.max { + if (value.len() as f64) > max { + return Err(ValidationError::with_value( + name, + format!("Array length must be at most {}", max), + serde_json::Value::Array(value.to_vec()), + )); + } + } + + // Validate items if schema is provided + if let Some(items_schema) = &schema.items { + for (index, item) in value.iter().enumerate() { + let item_name = format!("{}[{}]", name, index); + if let Err(error) = validate_parameter_value(&item_name, item, items_schema) { + return Err(error); + } + } + } + + Ok(()) +} + +/// Gets the type of a JSON value +fn get_value_type(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "boolean", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_validate_required_parameter() { + let mut schemas = HashMap::new(); + schemas.insert("test".to_string(), ParameterSchema { + param_type: ParameterType::String, + required: Some(true), + default: None, + description: None, + r#enum: None, + min: None, + max: None, + pattern: None, + items: None, + properties: None, + }); + + let parameters = HashMap::new(); + let result = validate_parameters(¶meters, &schemas); + + assert!(!result.valid); + assert_eq!(result.errors.len(), 1); + assert_eq!(result.errors[0].path, "test"); + } + + #[test] + fn test_validate_string_parameter() { + let mut schemas = HashMap::new(); + schemas.insert("test".to_string(), ParameterSchema { + param_type: ParameterType::String, + required: Some(true), + default: None, + description: None, + r#enum: None, + min: Some(3.0), + max: Some(10.0), + pattern: None, + items: None, + properties: None, + }); + + let mut parameters = HashMap::new(); + parameters.insert("test".to_string(), json!("hello")); + + let result = validate_parameters(¶meters, &schemas); + assert!(result.valid); + } +}