comfystream/nodes/web/js/comfystream_ui_preview_node.js at d4384412d6a434d63c3b86e735001518f28a87c6 · livepeer/comfystream · GitHub

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
// ComfyStream Node - A custom JavaScript node for ComfyUI
// This node displays the ComfyStream UI in an iframe

const app = window.comfyAPI?.app?.app;

// Register our extension
app.registerExtension({
    name: "ComfyStream.Node",

    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "ComfyStreamUIPreview") {
            // Set default size for the node type
            nodeType.size = [700, 800];

            // Make node resizable
            nodeType.resizable = true;

            // Save the original onNodeCreated method
            const onNodeCreated = nodeType.prototype.onNodeCreated;

            // Override the onNodeCreated method
            nodeType.prototype.onNodeCreated = function() {
                // Set node properties before calling original method
                this.title = "ComfyStream UI";
                this.color = "#4B9CD3"; // Blue color for the node

                // Set initial size
                this.size = [700, 800];

                // Make the node resizable
                this.resizable = true;
                this.flags.resizable = true;

                // Call the original onNodeCreated method if it exists
                const result = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;

                // Create iframe element
                this.iframe = document.createElement("iframe");
                this.iframe.style.width = "100%";
                this.iframe.style.height = "100%";
                this.iframe.style.border = "none";
                this.iframe.style.borderRadius = "8px";

                // Function to load or refresh the iframe
                this.loadIframe = () => {
                    fetch('/comfystream/extension_info')
                        .then(response => response.json())
                        .then(data => {
                            if (data.success) {
                                // Use the current origin with the static route from extension_info
                                this.iframe.src = `${window.location.origin}${data.static_route}/index.html`;
                            } else {
                                console.error("[ComfyStream] Error getting extension info:", data.error);
                                // Fallback to hardcoded path
                                const extensionName = "comfystream";
                                this.iframe.src = `${window.location.origin}/extensions/${extensionName}/static/index.html`;
                            }
                        })
                        .catch(error => {
                            console.error("[ComfyStream] Error fetching extension info:", error);
                            // Fallback to hardcoded path
                            const extensionName = "comfystream";
                            this.iframe.src = `${window.location.origin}/extensions/${extensionName}/static/index.html`;
                        });
                };

                // Initial load of the iframe
                this.loadIframe();

                // Add the iframe as a DOM widget
                this.iframeWidget = this.addDOMWidget("iframe", "UI", this.iframe, {
                    serialize: false,
                    width: this.size[0],
                    height: this.size[1] - 40
                });

                // Add a button to refresh the iframe
                this.addWidget("button", "Refresh UI", null, () => {
                    console.log("[ComfyStream] Refreshing UI...");
                    if (this.iframe) {
                        // If iframe already has a src, we can just reload it
                        if (this.iframe.src) {
                            this.iframe.src = this.iframe.src;
                        } else {
                            // Otherwise load it using our function
                            this.loadIframe();
                        }
                    }
                });

                // Add a button to launch the UI in a new tab
                this.addWidget("button", "Open in New Tab", null, () => {
                    // Get extension info which contains the correct UI URL
                    fetch('/comfystream/extension_info', {
                        method: 'GET',
                        headers: {
                            'Accept': 'application/json'
                        }
                    })
                    .then(response => response.json())
                    .then(data => {
                        if (data.success) {
                            // Use the current origin with the static route
                            const uiUrl = `${window.location.origin}${data.static_route}/index.html`;
                            // Open the URL in a new tab
                            window.open(uiUrl, '_blank');
                        }
                    })
                    .catch(error => {
                        console.error("[ComfyStream] Error launching UI:", error);
                    });
                });


                return result;
            };

            // Add a helper method to update iframe size
            nodeType.prototype.updateIframeSize = function() {
                if (this.iframeWidget) {
                    this.iframeWidget.width = this.size[0];
                    this.iframeWidget.height = this.size[1] - 40;

                    // Also update the iframe element directly
                    if (this.iframe) {
                        this.iframe.style.width = this.size[0] + "px";
                        this.iframe.style.height = (this.size[1] - 40) + "px";
                    }

                    // Force a canvas update
                    this.setDirtyCanvas(true, true);
                }
            };

            // Override the onExecute method
            nodeType.prototype.onExecute = function() {
                // Trigger the output
                this.triggerSlot(0);
            };
        }
    }
});