added iframe 3d viewer

__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"
 ]

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}<Loaded>{CLEAR}")
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

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 };
+