comfyui_rs/
types.rs

1use serde::{Deserialize, Serialize};
2use std::time::Duration;
3
4/// Real-time progress update from ComfyUI's WebSocket.
5#[derive(Debug, Clone)]
6pub struct ProgressUpdate {
7    pub current_step: u32,
8    pub total_steps: u32,
9}
10
11/// Richer progress type with node and prompt context.
12#[derive(Debug, Clone)]
13pub struct ComfyProgress {
14    /// Current sampling step.
15    pub current_step: u32,
16    /// Total sampling steps.
17    pub total_steps: u32,
18    /// The node currently executing, if known.
19    pub node_id: Option<String>,
20    /// The prompt this progress belongs to.
21    pub prompt_id: String,
22}
23
24/// High-level status of a ComfyUI generation.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum ComfyStatus {
27    /// Prompt is queued but not yet executing.
28    Queued,
29    /// A specific node is executing.
30    Running { node_id: String },
31    /// Sampling progress update.
32    Progress { step: u32, total: u32 },
33    /// Generation completed successfully.
34    Completed,
35    /// Generation failed with an error.
36    Failed { error: String },
37}
38
39/// WebSocket connection configuration for recovery and limits.
40#[derive(Debug, Clone)]
41pub struct WsConfig {
42    /// Number of reconnection attempts before falling back to polling (default: 3).
43    pub reconnect_attempts: u32,
44    /// Delay between reconnection attempts (default: 1s).
45    pub reconnect_delay: Duration,
46    /// Timeout for receiving a WebSocket message (default: 30s).
47    pub message_timeout: Duration,
48    /// Maximum messages for the tracked prompt before fallback (default: 10,000).
49    pub max_messages_per_prompt: usize,
50    /// Maximum total messages before fallback (default: 50,000).
51    pub max_total_messages: usize,
52}
53
54impl Default for WsConfig {
55    fn default() -> Self {
56        Self {
57            reconnect_attempts: 3,
58            reconnect_delay: Duration::from_secs(1),
59            message_timeout: Duration::from_secs(30),
60            max_messages_per_prompt: 10_000,
61            max_total_messages: 50_000,
62        }
63    }
64}
65
66/// Limits for image download operations.
67#[derive(Debug, Clone)]
68pub struct DownloadLimits {
69    /// Maximum image size in bytes (default: 100 MB).
70    pub max_image_bytes: usize,
71    /// Per-image download timeout (default: 60s).
72    pub download_timeout: Duration,
73}
74
75impl Default for DownloadLimits {
76    fn default() -> Self {
77        Self {
78            max_image_bytes: 100 * 1024 * 1024,
79            download_timeout: Duration::from_secs(60),
80        }
81    }
82}
83
84/// Reference to an image stored in ComfyUI's output directory.
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct ImageRef {
87    pub filename: String,
88    pub subfolder: String,
89    pub img_type: String,
90}
91
92/// Parsed history entry for a completed prompt.
93#[derive(Debug, Clone)]
94pub struct PromptHistory {
95    pub status: String,
96    pub completed: bool,
97    pub images: Vec<ImageRef>,
98}
99
100/// Snapshot of ComfyUI's queue state.
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct QueueStatus {
103    pub running: u32,
104    pub pending: u32,
105}
106
107/// Outcome of waiting for a generation to finish.
108#[derive(Debug, Clone)]
109pub enum GenerationOutcome {
110    /// Generation completed successfully with output images.
111    Completed { images: Vec<ImageRef> },
112    /// ComfyUI reported an execution-level failure.
113    Failed { error: String },
114    /// Timed out before completion.
115    TimedOut,
116}
comfyui_rs/types.rs

comfyui_rs/
types.rs
