diff --git a/__init__.py b/__init__.py index d98c48c..89cc691 100644 --- a/__init__.py +++ b/__init__.py @@ -4,9 +4,11 @@ from .server import run_https_server import threading NODE_CLASS_MAPPINGS = { - "Environment Visualizer": EnvironmentVisualizer, - "Interpolate Edges": InterpolateEdges + "EnvironmentVisualizer": EnvironmentVisualizer, + "InterpolateEdges": InterpolateEdges } +WEB_DIRECTORY = "./web" + server_thread = threading.Thread(target=run_https_server) server_thread.start() diff --git a/node.py b/node.py index 3684ec5..9ffaf70 100644 --- a/node.py +++ b/node.py @@ -1,5 +1,7 @@ from .server import get_lan_ip from .server import SERVER_PORT +from server import PromptServer +from aiohttp import web from PIL import Image import os import numpy as np @@ -7,6 +9,12 @@ import webbrowser import time import re + +@PromptServer.instance.routes.post("/get_url") +async def get_url(_): + return web.json_response({"port": str(SERVER_PORT)}) + + class EnvironmentVisualizer: @classmethod @@ -15,7 +23,7 @@ class EnvironmentVisualizer: "required": { "texture": ("IMAGE", ), "name": ("STRING", ), - "open_visualizer": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "open_automatically": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), }, "optional": { "depth": ("IMAGE", ), @@ -49,7 +57,7 @@ class EnvironmentVisualizer: return new_name - def save_environment(self, texture, name, open_visualizer, depth=None): + def save_environment(self, texture, name, open_automatically, depth=None): if depth and texture.shape[0] != depth.shape[0]: raise Exception("Number of environment textures and depth maps must be equivalent.") @@ -69,7 +77,7 @@ class EnvironmentVisualizer: if depth: self.save_tensor_image(depth[batch_number], os.path.join(new_directory, 'depth.png')) - if open_visualizer: + if open_automatically: webbrowser.open(f"https://{get_lan_ip()}:{SERVER_PORT}/environments.html?env={new_name}") return {} diff --git a/web/environment_visualizer.js b/web/environment_visualizer.js new file mode 100644 index 0000000..88777f2 --- /dev/null +++ b/web/environment_visualizer.js @@ -0,0 +1,56 @@ +import { api } from "../../scripts/api.js"; +import { app } from "../../scripts/app.js"; + +let visualizerWindow = null; +let visualizerUrl = null; + +app.registerExtension({ + name: "ComfyUI.EnvironmentVisualizerExtension", + async setup() { + const menu = document.querySelector(".comfy-menu"); + const separator = document.createElement("hr"); + + separator.style.margin = "20px 0"; + separator.style.width = "100%"; + menu.append(separator); + + const managerButton = document.createElement("button"); + managerButton.textContent = "Environment Visualizer"; + managerButton.onclick = async function() { + if (visualizerUrl == null) { + try { + let responseData = await api.fetchApi("/get_url", { + method: "POST", + }); + + // Error handling + if (responseData.status != 200) { + console.log( + "Error [" + responseData.status + "] > " + responseData.statusText + ); + return; + } + responseData = await responseData?.json(); + if (!responseData || responseData == undefined || !responseData.hasOwnProperty("port")) { + console.log("Error: Could not get environment visualizer URL."); + return; + } + + const currentUrl = new URL(window.location.href); + visualizerUrl = `https://${currentUrl.hostname}:${responseData.port}`; + } catch (e) { + throw new Error(e); + } + } + + // Open or focus the visualizer window + if (!visualizerWindow || visualizerWindow.closed) { + visualizerWindow = window.open(visualizerUrl, '_blank'); + } else { + visualizerWindow.focus(); + } + + } + menu.append(managerButton); + }, +}); \ No newline at end of file