//! 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( 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( 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( 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( 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(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(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 output in outputs.values() { 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!( "{base_url_str}/view?filename={filename}&subfolder={subfolder}&type={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(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(response.error_for_status().unwrap_err())); } Ok(()) } }