diff --git a/__init__.py b/__init__.py index dabfaab..c3cb656 100644 --- a/__init__.py +++ b/__init__.py @@ -2,19 +2,25 @@ from .layered_infinite_zoom import LayeredInfiniteZoom from .replicate_Flux_Pro_Ultra import ReplicateAPI_flux_1_1_pro_ultra from .replicate_Flux_Fill import ReplicateAPI_flux_fill_pro +from .iframe_view import IframeView + +WEB_DIRECTORY = "./js" NODE_CLASS_MAPPINGS = { "LayeredInfiniteZoom": LayeredInfiniteZoom, "ReplicateAPI_flux_1_1_pro_ultra": ReplicateAPI_flux_1_1_pro_ultra, - "ReplicateAPI_flux_fill_pro": ReplicateAPI_flux_fill_pro + "ReplicateAPI_flux_fill_pro": ReplicateAPI_flux_fill_pro, + "IframeView": IframeView } NODE_DISPLAY_NAME_MAPPINGS = { "LayeredInfiniteZoom": "Layered Infinite Zoom", "ReplicateAPI_flux_1_1_pro_ultra": "Replicate flux 1.1 pro ultra", - "ReplicateAPI_flux_fill_pro": "Replicate flux fill pro" + "ReplicateAPI_flux_fill_pro": "Replicate flux fill pro", + "IframeView": "Iframe View" } __all__ = [ "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY" ] diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..03eebb9 Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/__pycache__/iframe_view.cpython-312.pyc b/__pycache__/iframe_view.cpython-312.pyc new file mode 100644 index 0000000..8642895 Binary files /dev/null and b/__pycache__/iframe_view.cpython-312.pyc differ diff --git a/__pycache__/layered_infinite_zoom.cpython-312.pyc b/__pycache__/layered_infinite_zoom.cpython-312.pyc new file mode 100644 index 0000000..d02303a Binary files /dev/null and b/__pycache__/layered_infinite_zoom.cpython-312.pyc differ diff --git a/__pycache__/replicate_Flux_Fill.cpython-312.pyc b/__pycache__/replicate_Flux_Fill.cpython-312.pyc new file mode 100644 index 0000000..5ee81f8 Binary files /dev/null and b/__pycache__/replicate_Flux_Fill.cpython-312.pyc differ diff --git a/__pycache__/replicate_Flux_Pro_Ultra.cpython-312.pyc b/__pycache__/replicate_Flux_Pro_Ultra.cpython-312.pyc new file mode 100644 index 0000000..eb0427a Binary files /dev/null and b/__pycache__/replicate_Flux_Pro_Ultra.cpython-312.pyc differ diff --git a/iframe_view.py b/iframe_view.py new file mode 100644 index 0000000..8c2fe7c --- /dev/null +++ b/iframe_view.py @@ -0,0 +1,165 @@ +import os +import PIL +from PIL import Image, ImageOps +import hashlib +import base64 +from io import BytesIO +import torch +import numpy as np +import time +import json +import traceback +from aiohttp import web +import folder_paths +import asyncio + + +DEBUG = True +THREEVIEW_DICT = {} # Global dict (can be retained for compatibility if needed) + + +# Global dictionary to store texture data (populated from the JS side via postMessage) +TEXTURE_STORE = {} + +def get_global_texture_value(unique_id, key): + if unique_id in TEXTURE_STORE: + return TEXTURE_STORE[unique_id].get(key, None) + return None + + +def process_texture_data(data): + if not data: + print("[process_texture_data] No texture data received.") + return torch.zeros((1, 64, 64, 3), dtype=torch.float32) + try: + print("Raw data received:") + if isinstance(data, dict) and "data" in data: + if isinstance(data["data"], str): + print("[process_texture_data] Received base64 encoded texture data.") + b64data = data["data"] + if b64data.startswith("data:"): + print("[process_texture_data] Stripping data URL header from base64 string.") + b64data = b64data.split(",")[-1] + image_bytes = base64.b64decode(b64data) + pil_image = Image.open(BytesIO(image_bytes)).convert("RGB") + image = np.array(pil_image) + print(f"[process_texture_data] PIL image shape: {image.shape}") + else: + print("[process_texture_data] Received raw numeric texture data.") + width = data.get("width", 64) + height = data.get("height", 64) + pixel_array = np.array(data["data"], dtype=np.uint8) + try: + image = pixel_array.reshape(height, width, 3) + print(f"[process_texture_data] Reshaped image to: {image.shape}") + except Exception as reshape_err: + print(f"[process_texture_data] Error reshaping array: {reshape_err}") + raise + else: + print("[process_texture_data] Data format not recognized, returning default.") + return torch.zeros((1, 64, 64, 3), dtype=torch.float32) + + print(f"[process_texture_data] Post reshape image stats: shape={image.shape}, min={image.min()}, max={image.max()}") + image = image.astype(np.float32) / 255.0 + tensor = torch.from_numpy(image)[None,] + print(f"[process_texture_data] Tensor shape: {tensor.shape}, range=[{tensor.min().item():.3f}, {tensor.max().item():.3f}]") + return tensor + except Exception as e: + print(f"[process_texture_data] Error processing texture data: {e}") + traceback.print_exc() + return torch.zeros((1, 64, 64, 3), dtype=torch.float32) + +class IframeView: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "https://example.com"}), + "width": ("INT", {"default": 512, "min": 64, "max": 2048}), + "height": ("INT", {"default": 512, "min": 64, "max": 2048}), + "scene_state": ( + "STRING", + {"multiline": True, "default": "{\"camera\": {\"position\": [0,0,5]}, \"animation\": false}"} + ), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "color": "TEXTURE", + "canny": "TEXTURE", + "depth": "TEXTURE", + "normal": "TEXTURE" + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("color", "canny", "depth", "normal") + FUNCTION = "process_iframe" + CATEGORY = "lth" + + async def _process_iframe_async(self, url, width, height, scene_state, unique_id, color=None, canny=None, depth=None, normal=None): + try: + print("\n[process_iframe] Called with inputs:") + print(" url:", url) + print(" width:", width, "height:", height) + print(" scene_state:", scene_state) + print(" unique_id:", unique_id) + + # Since the iframe is always sending fully processed texture data, + # there is no need to trigger an intermediate save check. + # Directly retrieve the texture data. + if color is None: + color = get_global_texture_value(unique_id, "color") + if canny is None: + canny = get_global_texture_value(unique_id, "canny") + if depth is None: + depth = get_global_texture_value(unique_id, "depth") + if normal is None: + normal = get_global_texture_value(unique_id, "normal") + + try: + state_data = json.loads(scene_state) + except json.JSONDecodeError as e: + print(f"[process_iframe] Invalid scene state JSON: {e}") + state_data = {} + + color_tensor = process_texture_data(color) + canny_tensor = process_texture_data(canny) + depth_tensor = process_texture_data(depth) + normal_tensor = process_texture_data(normal) + + return (color_tensor, canny_tensor, depth_tensor, normal_tensor) + except Exception as e: + print(f"[process_iframe] Error: {e}") + empty = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty, empty, empty, empty) + + def process_iframe(self, url, width, height, scene_state, unique_id, color=None, canny=None, depth=None, normal=None): + """ + Synchronous wrapper for _process_iframe_async. + """ + return asyncio.run( + self._process_iframe_async( + url, width, height, scene_state, unique_id, color, canny, depth, normal + ) + ) + +NODE_CLASS_MAPPINGS = { + "IframeView": IframeView +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IframeView": "Iframe View" +} + +WEB_DIRECTORY = "./js" + +GREEN = "\033[92m" +LIGHT_YELLOW = "\033[93m" +MAGENTA = "\033[95m" +BLUE = "\033[94m" +CLEAR = "\033[0m" + +nodesNames = ", ".join(NODE_DISPLAY_NAME_MAPPINGS.values()) +print(f"\n{MAGENTA}* {GREEN}lo-th -> {LIGHT_YELLOW}{nodesNames} {BLUE}{CLEAR}") + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/js/iframe_view.js b/js/iframe_view.js new file mode 100644 index 0000000..ef6ad2f --- /dev/null +++ b/js/iframe_view.js @@ -0,0 +1,296 @@ +// iframe_widget.js +import { app } from "../../scripts/app.js"; +import { $el } from "../../scripts/ui.js"; + +// Global store for texture data, keyed by unique_id +window.TEXTURE_STORE = window.TEXTURE_STORE || {}; + +// Function to create the iframe widget +async function widgetIframe(node, nodeData, inputData, app) { + // Create iframe wrapper with some styling + const iframeWrapper = $el("div.iframeWrapper", { + style: { + width: "100%", + height: "100%", + overflow: "hidden", + border: "1px solid var(--border-color)", + borderRadius: "4px" + } + }); + + // Create the iframe element + const iframe = $el("iframe", { + style: { + width: "100%", + height: "100%", + border: "none" + }, + sandbox: "allow-scripts allow-same-origin allow-forms", + loading: "lazy" + }); + + iframeWrapper.appendChild(iframe); + + // Locate our parameter widgets + const urlWidget = node.widgets.find((w) => w.name === "url"); + const widthWidget = node.widgets.find((w) => w.name === "width"); + const heightWidget = node.widgets.find((w) => w.name === "height"); + const stateWidget = node.widgets.find((w) => w.name === "scene_state"); + + // Set initial URL if provided + if (urlWidget.value) { + iframe.src = urlWidget.value; + } + + // Message handler for communication from the iframe + function messageHandler(event) { + // Ensure the message comes from our iframe + if (event.source !== iframe.contentWindow) return; + + console.log("[messageHandler] Received message:", event.data); + + const { type, data } = event.data; + switch (type) { + case 'IFRAME_READY': + sendParamsToIframe(); + break; + case 'TEXTURE_OUTPUT': + console.log("[messageHandler] TEXTURE_OUTPUT received with data:", data); + handleTextureOutput(data); + break; + case 'SCENE_STATE_UPDATED': + if (data.state) { + stateWidget.value = JSON.stringify(data.state, null, 2); + } + break; + default: + console.warn("[messageHandler] Unhandled message type:", type); + } + } + + // Attach the message handler + window.addEventListener("message", messageHandler, false); + + // Function to send parameters (width, height, scene state) to the iframe + function sendParamsToIframe() { + let sceneState; + try { + sceneState = JSON.parse(stateWidget.value || "{}"); + } catch (e) { + console.error("Error parsing scene state, using default:", e); + sceneState = {}; + } + const params = { + type: 'PARAMS_UPDATE', + data: { + width: parseInt(widthWidget.value, 10) || 512, + height: parseInt(heightWidget.value, 10) || 512, + sceneState, + } + }; + + if (iframe.contentWindow) { + // If you know the target origin, replace '*' with that string + iframe.contentWindow.postMessage(params, '*'); + } + } + + // Function to handle texture output coming from the iframe + async function handleTextureOutput(data) { + // Make sure the node has a valid unique_id + if (!node.unique_id) { + console.warn("[handleTextureOutput] Node unique_id is not set, cannot update texture store."); + return; + } + // For each expected texture key, update the hidden widget and the global texture store. + const textureKeys = ["color", "canny", "depth", "normal"]; + textureKeys.forEach((key) => { + if (data[key]) { + node.setHiddenWidgetValue(key, { + data: data[key].data, + width: data[key].width, + height: data[key].height, + format: data[key].format, + type: data[key].type, + shape: data[key].shape + }); + } + }); + + // Optionally clear heavy data from the global store if it is no longer needed + if (node.unique_id && window.TEXTURE_STORE[node.unique_id]) { + // You can conditionally delete or nullify large arrays after processing. + delete window.TEXTURE_STORE[node.unique_id]; + } + // Delay the prompt so that the global store has updated + setTimeout(() => { + console.log("[handleTextureOutput] Final global texture store for unique_id", node.unique_id, window.TEXTURE_STORE[node.unique_id]); + app.queuePrompt(); + }, 50); + } + + // Create the DOM widget for this node and register cleanup + const widget = node.addDOMWidget("iframe", "custom_widget", iframeWrapper, { + getValue() { + return { + url: urlWidget.value, + width: widthWidget.value, + height: heightWidget.value, + scene_state: stateWidget.value + }; + }, + setValue(v) { + if (v.url) { + iframe.src = v.url; + urlWidget.value = v.url; + } + } + }); + + // Attach a reference to the iframe on the widget so we can access it later + widget.iframe = iframe; + + // Handle resizing: update both node size and iframe dimensions, then re-send params + node.onResize = function () { + const width = parseInt(widthWidget.value, 10) || 512; + const height = parseInt(heightWidget.value, 10) || 512; + node.size = [width, Math.max(height + 40, 140)]; + iframe.style.width = width + "px"; + iframe.style.height = (height - 40) + "px"; + sendParamsToIframe(); + }; + + // When the URL widget changes, update the iframe's src + urlWidget.callback = function () { + iframe.src = this.value; + }; + + // When any of the parameter widgets change, update the iframe with the new parameters + widthWidget.callback = heightWidget.callback = stateWidget.callback = function () { + sendParamsToIframe(); + }; + + // Override the default setHiddenWidgetValue to also update our global TEXTURE_STORE. + node.setHiddenWidgetValue = function (name, value) { + let widgetEntry = this.widgets.find(w => w.name === name); + if (widgetEntry) { + widgetEntry.value = value; + console.log(`[setHiddenWidgetValue] Updated widget '${name}' to:`, value); + } else { + // Create the widget if it does not exist. + widgetEntry = { name, value, isHidden: true }; + this.widgets.push(widgetEntry); + console.log(`[setHiddenWidgetValue] Created and set widget '${name}' to:`, value); + } + if (this.unique_id) { + window.TEXTURE_STORE = window.TEXTURE_STORE || {}; + window.TEXTURE_STORE[this.unique_id] = window.TEXTURE_STORE[this.unique_id] || {}; + window.TEXTURE_STORE[this.unique_id][name] = value; + console.log(`[setHiddenWidgetValue] Global texture store for unique_id '${this.unique_id}':`, window.TEXTURE_STORE[this.unique_id]); + } + }; + + // When creating the widget +node.cleanup = () => { + window.removeEventListener("message", messageHandler, false); + if (node.unique_id && window.TEXTURE_STORE[node.unique_id]) { + delete window.TEXTURE_STORE[node.unique_id]; + } + console.log("[cleanup] Cleaned up resources for node:", node.unique_id); + }; + + // Updated onSerialize: Strip out the heavy "data" field to avoid localStorage quota problems. + node.onSerialize = (nodeData) => { + const textureKeys = ["color", "canny", "depth", "normal"]; + nodeData.hidden = nodeData.hidden || {}; + textureKeys.forEach((key) => { + const widgetEntry = node.widgets.find(w => w.name === key); + if (widgetEntry && widgetEntry.value) { + // If the widget's value includes a "data" field, remove it from the serialization. + if ("data" in widgetEntry.value) { + const { data, ...metadata } = widgetEntry.value; + nodeData.hidden[key] = metadata; + } else { + nodeData.hidden[key] = widgetEntry.value; + } + } + }); + }; + + return widget; +} + +// Register the extension and add extra menu options for the IframeView node +app.registerExtension({ + name: "Iframe View", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeType.comfyClass === "IframeView") { + // Add extra menu options to refresh the iframe and trigger texture capture + nodeType.prototype.getExtraMenuOptions = function (_, options) { + options.unshift( + { + content: "🔄 Refresh Iframe", + callback: () => { + const domWidget = this.widgets.find(w => w.name === "iframe"); + if (domWidget && domWidget.iframe) { + const url = this.widgets.find(w => w.name === "url").value; + domWidget.iframe.src = url; + } + } + }, + { + content: "📸 Capture Textures", + callback: () => { + const domWidget = this.widgets.find(w => w.name === "iframe"); + if (domWidget?.iframe?.contentWindow) { + domWidget.iframe.contentWindow.postMessage( + { type: 'CAPTURE_TEXTURES' }, + '*' + ); + } + } + } + ); + }; + + const originalOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = async function () { + // If nodeData.hidden definitions exist, add any missing widgets. + if (nodeData && nodeData.hidden) { + Object.keys(nodeData.hidden).forEach((key) => { + if (!this.widgets.find(widget => widget.name === key)) { + this.widgets.push({ name: key, value: nodeData.hidden[key] || null, isHidden: true }); + console.log(`[onNodeCreated] Hidden widget '${key}' added.`); + } + }); + } + const r = originalOnNodeCreated ? originalOnNodeCreated.apply(this, arguments) : undefined; + const widget = await widgetIframe(this, nodeData, nodeData, app); + this.onResize(); + + // Synchronize the node's unique_id with the one passed from Python. + if (nodeData && nodeData.hidden && nodeData.hidden.unique_id) { + this.unique_id = nodeData.hidden.unique_id; + console.log(`[onNodeCreated] Using provided unique_id from nodeData: ${this.unique_id}`); + let uniqueWidget = this.widgets.find(w => w.name === "unique_id"); + if (uniqueWidget) { + uniqueWidget.value = this.unique_id; + } else { + this.widgets.push({ name: "unique_id", value: this.unique_id, isHidden: true }); + } + } else { + // Fallback: generate a unique id. + this.unique_id = Math.random().toString(36).substr(2, 9); + console.log(`[onNodeCreated] Generated fallback unique_id: ${this.unique_id}`); + let uniqueWidget = this.widgets.find(w => w.name === "unique_id"); + if (uniqueWidget) uniqueWidget.value = this.unique_id; + else this.widgets.push({ name: "unique_id", value: this.unique_id, isHidden: true }); + } + return r; + }; + } + } +}); + +export { widgetIframe }; +