cmfy 0.4.1

A CLI companion app for Comfy UI
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use crate::{MarkAs, WithStatus};

use super::Prompt;
use chrono::{serde::ts_milliseconds, DateTime, Utc};
use itertools::Itertools;
use serde::{Deserialize, Serialize};
use std::collections::{btree_map::IntoValues, BTreeMap};

#[derive(Debug, Clone, Deserialize)]
#[serde(transparent)]
pub struct History(pub BTreeMap<String, HistoryLogEntry>);

#[derive(Debug, Clone, Deserialize)]
pub struct HistoryLogEntry {
    pub prompt: Prompt,
    pub outputs: Outputs,
    pub status: Status,
    pub meta: Meta,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(transparent)]
pub struct Outputs(pub BTreeMap<String, Output>);

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum Output {
    Images { images: Vec<Image> },
    Other(BTreeMap<String, Vec<serde_json::Value>>),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Image {
    pub filename: String,
    pub subfolder: String,
    #[serde(rename = "type")]
    pub _type: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Status {
    pub status_str: String,
    pub completed: bool,
    pub messages: Vec<Message>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Message {
    pub kind: MessageKind,
    pub data: MessageData,
}

#[derive(PartialEq, Eq, Debug, Clone, Deserialize)]
pub enum MessageKind {
    #[serde(rename = "execution_start")]
    Start,
    #[serde(rename = "execution_cached")]
    Cached,
    #[serde(rename = "execution_success")]
    Success,
    #[serde(rename = "execution_interrupted")]
    Interruped,
    #[serde(rename = "execution_error")]
    Error,
}

#[derive(Debug, Clone, Deserialize)]
pub struct MessageData {
    pub prompt_id: String,
    #[serde(with = "ts_milliseconds")]
    pub timestamp: DateTime<Utc>,
    #[serde(flatten)]
    pub other: BTreeMap<String, serde_json::Value>,
}

#[derive(Debug, Clone, Deserialize, Eq, PartialEq, PartialOrd, Ord)]
#[serde(untagged)]
pub enum StringOrNumber {
    String(String),
    Number(u32),
}

#[derive(Debug, Clone, Deserialize)]
#[serde(transparent)]
pub struct Meta(pub BTreeMap<StringOrNumber, Metadata>);

#[derive(Debug, Clone, Deserialize)]
pub struct Metadata {
    pub node_id: String,
    pub display_node: String,
    pub parent_node: Option<String>,
    pub real_node_id: String,
}

impl MarkAs for Prompt {}
pub type PromptBatchEntry = WithStatus<Prompt, Outputs>;
pub type PromptBatch = Vec<PromptBatchEntry>;

impl History {
    pub fn into_batch_entries(self) -> impl Iterator<Item = PromptBatchEntry> {
        self.into_iter().map(|entry| entry.into())
    }
}

impl IntoIterator for History {
    type Item = HistoryLogEntry;
    type IntoIter = IntoValues<String, HistoryLogEntry>;
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_values()
    }
}

impl From<HistoryLogEntry> for PromptBatchEntry {
    fn from(entry: HistoryLogEntry) -> Self {
        if entry.status.messages.iter().any(|msg| msg.kind == MessageKind::Interruped) {
            entry.prompt.mark_as(crate::Status::Cancelled)
        } else {
            entry.prompt.mark_as(crate::Status::Completed(entry.outputs))
        }
    }
}

impl From<History> for PromptBatch {
    fn from(history: History) -> Self {
        let mut result = history.into_batch_entries().collect_vec();
        result.sort_by(|l, r| l.inner.index.cmp(&r.inner.index));
        result
    }
}

impl Outputs {
    pub fn images(&self) -> impl Iterator<Item = &Image> {
        self.0
            .values()
            .filter_map(|output| {
                if let Output::Images { images } = output {
                    Some(images)
                } else {
                    None
                }
            })
            .flatten()
    }
}