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

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use super::{Image, Prompt};

/// Struct containing task results from the ComfyUI API `history` endpoint.
#[derive(Deserialize, Debug)]
#[serde(transparent)]
pub struct History {
    /// Completed tasks indexed by their uuid.
    pub tasks: HashMap<uuid::Uuid, Task>,
}

/// Struct representing a single task result from the ComfyUI API `history` endpoint.
#[derive(Serialize, Deserialize, Debug)]
pub struct Task {
    /// Outputs from the task.
    pub outputs: Outputs,
    /// Information about prompt execution.
    pub prompt: PromptResult,
}

/// Struct representing outputs from a task.
#[derive(Serialize, Deserialize, Debug)]
#[serde(transparent)]
pub struct Outputs {
    /// Outputs from the task indexed by node.
    pub nodes: HashMap<String, NodeOutputOrUnknown>,
}

/// Enumertion of all possible output types from a node.
#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum NodeOutputOrUnknown {
    /// Enum variant representing image outputs from a node.
    NodeOutput(NodeOutput),
    /// Struct capturing unknown outputs.
    Unknown(serde_json::Value),
}

/// Struct representing image outputs from a node.
#[derive(Serialize, Deserialize, Debug)]
pub struct NodeOutput {
    /// Images from the node.
    pub images: Vec<Image>,
}

/// Struct representing a prompt result.
#[derive(Serialize, Deserialize, Debug)]
#[serde(from = "(u64, uuid::Uuid, Prompt, ExtraData, OutputsToExecute)")]
pub struct PromptResult {
    /// The task number.
    pub num: u64,
    /// The task uuid.
    pub id: uuid::Uuid,
    /// The prompt that was executed.
    pub prompt: Prompt,
    /// Extra data about execution.
    pub extra_data: ExtraData,
    /// Outputs executed for this prompt.
    pub outputs_to_execute: OutputsToExecute,
}

impl From<(u64, uuid::Uuid, Prompt, ExtraData, OutputsToExecute)> for PromptResult {
    fn from(
        (num, id, prompt, extra_data, outputs_to_execute): (
            u64,
            uuid::Uuid,
            Prompt,
            ExtraData,
            OutputsToExecute,
        ),
    ) -> Self {
        Self {
            num,
            id,
            prompt,
            extra_data,
            outputs_to_execute,
        }
    }
}

impl From<PromptResult> for (u64, uuid::Uuid, Prompt, ExtraData, OutputsToExecute) {
    fn from(
        PromptResult {
            num,
            id,
            prompt,
            extra_data,
            outputs_to_execute,
        }: PromptResult,
    ) -> Self {
        (num, id, prompt, extra_data, outputs_to_execute)
    }
}

/// Struct representing extra data about prompt execution.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ExtraData {
    /// The client id that performed the request.
    pub client_id: uuid::Uuid,
}

/// Struct representing outputs to execute for a prompt.
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(transparent)]
pub struct OutputsToExecute {
    /// List of nodes which have outputs.
    pub nodes: Vec<String>,
}