diff --git a/__pycache__/iframe_view.cpython-312.pyc b/__pycache__/iframe_view.cpython-312.pyc index b1c2773..0b6c23e 100644 Binary files a/__pycache__/iframe_view.cpython-312.pyc and b/__pycache__/iframe_view.cpython-312.pyc differ diff --git a/iframe_view.py b/iframe_view.py index 0d0a857..e85085a 100644 --- a/iframe_view.py +++ b/iframe_view.py @@ -28,33 +28,56 @@ def get_global_texture_value(unique_id, key): def process_texture_data(data): + """ + Process incoming texture data from the iframe. + We expect data to be a dict containing a base64-encoded image string (under the "base64" key) + along with width, height, and other metadata. + """ start_time = time.time() - print("[process_texture_data] Started processing texture data.") + print("[process_texture_data] Started processing texture data. Data type:", type(data)) - if not data: - print("[process_texture_data] No texture data received.") - result = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + if data is None: + print("[process_texture_data] No texture data received (data is None).") + result = torch.zeros((64, 64, 3), dtype=torch.float32) print(f"[process_texture_data] Returning default tensor with shape: {result.shape}") return result + # Log the structure of the data if it's a dict. + if isinstance(data, dict): + print("[process_texture_data] Texture data keys:", list(data.keys())) + if "base64" in data: + base64_string = data.get("base64") + print("[process_texture_data] Found 'base64' key with string length:", len(base64_string)) + elif "data" in data: + base64_string = data.get("data") + print("[process_texture_data] Found 'data' key with string length:", len(base64_string) if isinstance(base64_string, str) else "non-string") + else: + print("[process_texture_data] Unexpected data format. Data:", data) + 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"] + # Process if the data dict has either "base64" or "data" + if isinstance(data, dict) and ("data" in data or "base64" in data): + b64data = data.get("base64", data.get("data")) + if isinstance(b64data, str): + print("[process_texture_data] b64data is string. Checking for header...") if b64data.startswith("data:"): print("[process_texture_data] Stripping data URL header from base64 string.") - b64data = b64data.split(",")[-1] + parts = b64data.split(",") + print("[process_texture_data] Header part:", parts[0]) + b64data = parts[-1] + print("[process_texture_data] Final base64 string length:", len(b64data)) image_bytes = base64.b64decode(b64data) + print("[process_texture_data] Decoded base64, received", len(image_bytes), "bytes.") pil_image = Image.open(BytesIO(image_bytes)).convert("RGB") + print("[process_texture_data] PIL image mode:", pil_image.mode) 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.") + print("[process_texture_data] b64data is not a string. Type:", type(b64data)) width = data.get("width", 64) height = data.get("height", 64) - pixel_array = np.array(data["data"], dtype=np.uint8) + pixel_array = np.array(data.get("data"), dtype=np.uint8) + print("[process_texture_data] Raw pixel_array shape from 'data':", pixel_array.shape) try: image = pixel_array.reshape(height, width, 3) print(f"[process_texture_data] Reshaped image to: {image.shape}") @@ -62,24 +85,39 @@ def process_texture_data(data): print(f"[process_texture_data] Error reshaping array: {reshape_err}") raise else: - print("[process_texture_data] Data format not recognized, returning default.") - result = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + print("[process_texture_data] Data format not recognized, returning default tensor.") + result = torch.zeros((64, 64, 3), dtype=torch.float32) print(f"[process_texture_data] Returning default tensor with shape: {result.shape}") return result - print(f"[process_texture_data] Post reshape image stats: shape={image.shape}, min={image.min()}, max={image.max()}") + print(f"[process_texture_data] Post-conversion 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}]") + + tensor = torch.from_numpy(image) + print(f"[process_texture_data] Final tensor shape: {tensor.shape}, range=[{tensor.min().item():.3f}, {tensor.max().item():.3f}]") elapsed = time.time() - start_time print(f"[process_texture_data] Completed processing in {elapsed:.4f} seconds.") return tensor except Exception as e: print(f"[process_texture_data] Error processing texture data: {e}") traceback.print_exc() - result = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + result = torch.zeros((64, 64, 3), dtype=torch.float32) return result + +async def wait_for_capture_completion(unique_id, target_frame_count, timeout=20.0, interval=0.5): + waited = 0.0 + while waited < timeout: + frames = get_global_texture_value(unique_id, "animationFrames") + frame_count = len(frames) if frames is not None else 0 + print(f"[wait_for_capture_completion] {waited:.2f}s elapsed: Found {frame_count} frames for unique_id '{unique_id}'") + if frames and frame_count >= target_frame_count: + return frames + await asyncio.sleep(interval) + waited += interval + print("[wait_for_capture_completion] Timeout reached, returning empty list.") + return [] + class IframeView: @classmethod def INPUT_TYPES(cls): @@ -90,71 +128,133 @@ class IframeView: "height": ("INT", {"default": 512, "min": 64, "max": 2048}), "scene_state": ( "STRING", - {"multiline": True, "default": "{\"camera\": {\"position\": [0,0,5]}, \"animation\": false}"} + { + "multiline": True, + "default": "{\"camera\": {\"position\": [0,0,5]}, \"animation\": false}", + }, ), + "frame_count": ("INT", {"default": 4, "min": 1, "max": 60}), }, "hidden": { "unique_id": "UNIQUE_ID", "color": "TEXTURE", "canny": "TEXTURE", "depth": "TEXTURE", - "normal": "TEXTURE" - } + "normal": "TEXTURE", + "animationFrames": "TEXTURE_SEQUENCE", + }, } - RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") - RETURN_NAMES = ("color", "canny", "depth", "normal") + RETURN_TYPES = ( + "IMAGE", "IMAGE", "IMAGE", "IMAGE", # base outputs + "IMAGE", "IMAGE", "IMAGE", "IMAGE", # animation outputs (one per channel) + ) + RETURN_NAMES = ( + "color", "canny", "depth", "normal", + "animation_color", "animation_canny", "animation_depth", "animation_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): + async def _process_iframe_async( + self, + url, + width, + height, + scene_state, + frame_count, + unique_id, + color=None, + canny=None, + depth=None, + normal=None, + animationFrames=None, + ): start_time = time.time() print("[process_iframe] Starting asynchronous processing at", start_time) - 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) + print("frame_count:", frame_count) + + unique_id = str(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") + # Process base textures first. + color_tensor = process_texture_data(color) + canny_tensor = process_texture_data(canny) + depth_tensor = process_texture_data(depth) + normal_tensor = process_texture_data(normal) - try: - state_data = json.loads(scene_state) - except json.JSONDecodeError as e: - print(f"[process_iframe] Invalid scene state JSON: {e}") - state_data = {} + # Use provided hidden animation frames if available; otherwise, use fallback. + if animationFrames is not None and len(animationFrames) >= frame_count: + print("[process_iframe] Using provided animationFrames from hidden data.") + else: + print("[process_iframe] Animation frames not provided or incomplete; using fallback.") + animationFrames = await wait_for_capture_completion(unique_id, frame_count) - color_tensor = process_texture_data(color) - canny_tensor = process_texture_data(canny) - depth_tensor = process_texture_data(depth) - normal_tensor = process_texture_data(normal) - - elapsed = time.time() - start_time - print(f"[process_iframe] Asynchronous processing completed in {elapsed:.4f} seconds.") - 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) + # Process each frame's texture into a tensor. + color_frames = [ + process_texture_data(frame["textures"]["color"]) + for frame in animationFrames + if "textures" in frame and "color" in frame["textures"] + ] + canny_frames = [ + process_texture_data(frame["textures"]["canny"]) + for frame in animationFrames + if "textures" in frame and "canny" in frame["textures"] + ] + depth_frames = [ + process_texture_data(frame["textures"]["depth"]) + for frame in animationFrames + if "textures" in frame and "depth" in frame["textures"] + ] + normal_frames = [ + process_texture_data(frame["textures"]["normal"]) + for frame in animationFrames + if "textures" in frame and "normal" in frame["textures"] + ] - 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. - """ + # Convert the processed frames into batched tensors. + animation_color = torch.stack(color_frames, dim=0) if color_frames else None + animation_canny = torch.stack(canny_frames, dim=0) if canny_frames else None + animation_depth = torch.stack(depth_frames, dim=0) if depth_frames else None + animation_normal = torch.stack(normal_frames, dim=0) if normal_frames else None + + # Clear the local reference to the heavy animationFrames data. + # This lets garbage-collection release the base64 strings from memory. + animationFrames = None + + elapsed = time.time() - start_time + print(f"[process_iframe] Asynchronous processing completed in {elapsed:.4f} seconds.") + return ( + color_tensor, canny_tensor, depth_tensor, normal_tensor, + animation_color, animation_canny, animation_depth, animation_normal, + ) + + def process_iframe( + self, + url, + width, + height, + scene_state, + frame_count, + unique_id, + color=None, + canny=None, + depth=None, + normal=None, + animationFrames=None, + ): return asyncio.run( self._process_iframe_async( - url, width, height, scene_state, unique_id, color, canny, depth, normal + url, + width, + height, + scene_state, + frame_count, + unique_id, + color, + canny, + depth, + normal, + animationFrames, ) ) diff --git a/js/iframe_view.js b/js/iframe_view.js index 4df4f28..7f6023e 100644 --- a/js/iframe_view.js +++ b/js/iframe_view.js @@ -8,7 +8,7 @@ window.TEXTURE_STORE = window.TEXTURE_STORE || {}; // Function to create the iframe widget async function widgetIframe(node, nodeData, inputData, app) { console.log("[widgetIframe] Initializing iframe widget for node:", node); - + node.isCaptureComplete = false; // Create iframe wrapper with some styling const iframeWrapper = $el("div.iframeWrapper", { style: { @@ -44,6 +44,22 @@ async function widgetIframe(node, nodeData, inputData, app) { iframe.src = urlWidget.value; } + // Ensure we have a "frame_count" widget to specify animation frame count + let frameCountWidget = node.widgets.find((w) => w.name === "frame_count"); + if (!frameCountWidget) { + frameCountWidget = { + name: "frame_count", + value: "4", // default number of frames as a string + isHidden: false, // visible so user can change it + callback: function () { + // Reset frame buffer when the frame_count changes. + node.frameBuffer = []; + node.isCaptureComplete = false; + }, + }; + node.widgets.push(frameCountWidget); + } + // Message handler for communication from the iframe function messageHandler(event) { // Ensure the message comes from our iframe @@ -55,7 +71,9 @@ async function widgetIframe(node, nodeData, inputData, app) { case "IFRAME_READY": console.log("[messageHandler] IFRAME_READY received"); sendParamsToIframe(); - break; + if (iframe.contentWindow) + // iframe.contentWindow.postMessage({ type: "CAPTURE_TEXTURES" }, "*"); + break; case "TEXTURE_OUTPUT": console.log( "[messageHandler] TEXTURE_OUTPUT received with data:", @@ -96,6 +114,7 @@ async function widgetIframe(node, nodeData, inputData, app) { width: parseInt(widthWidget.value, 10) || 512, height: parseInt(heightWidget.value, 10) || 512, sceneState, + unique_id: node.unique_id, }, }; @@ -111,66 +130,73 @@ async function widgetIframe(node, nodeData, inputData, app) { } } - // Function to handle texture output coming from the iframe - async function handleTextureOutput(data) { - const startTime = performance.now(); - console.log( - "[handleTextureOutput] Processing texture output for node:", - node.unique_id, - "at", - startTime - ); + window.addEventListener("message", (event) => { + if (event.data && event.data.type === "TEXTURE_OUTPUT") { + console.log( + "[iframe_view] Received TEXTURE_OUTPUT message:", + event.data.data + ); + handleTextureOutput(event.data.data); + } + }); - // Make sure the node has a valid unique_id - if (!node.unique_id) { + // Updated handleTextureOutput: Capture frames until target is reached, then aggregate them. + function handleTextureOutput(data) { + if (node.isCaptureComplete) return; + // Validate that we received all expected textures. + if (!data || !data.color || !data.depth || !data.normal || !data.canny) { console.warn( - "[handleTextureOutput] Node unique_id is not set, cannot update texture store." + "[handleTextureOutput] Incomplete texture data received; skipping frame." ); 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, - }); - } + + // Ensure the frame buffer exists. + if (!node.frameBuffer) node.frameBuffer = []; + + // Append the complete frame data (with base64 textures) into the buffer. + node.frameBuffer.push({ + timestamp: Date.now(), + textures: { + color: data.color, + depth: data.depth, + normal: data.normal, + canny: data.canny, + }, }); console.log( - "[handleTextureOutput] Global texture store for unique_id", - node.unique_id, - window.TEXTURE_STORE[node.unique_id] + "[handleTextureOutput] Captured frame", + node.frameBuffer.length ); - // Optionally clear heavy data from the global store if it is no longer needed - // Delay the prompt so that the global store has updated - setTimeout(() => { - 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]; - } + // Use the cached target frame count value. + let targetFrameCount = node.targetFrameCount || 4; + console.log("[handleTextureOutput] Target frame count:", targetFrameCount); - if (app) { - app.queuePrompt(); + // If we haven't reached the target, request the next capture. + if (node.frameBuffer.length < targetFrameCount) { + if (iframe && iframe.contentWindow) { + console.log("[handleTextureOutput] Requesting next frame capture..."); + if (node.frameBuffer.length === 0) { + // iframe.contentWindow.postMessage({ type: "CAPTURE_TEXTURES" }, "*"); + // console.log("[handleTextureOutput] starting capture..."); + } } else { - console.error( - "queuePrompt is not available in the current app context." + console.warn( + "[handleTextureOutput] No iframe contentWindow available." ); } - }, 50); - - const elapsedTime = performance.now() - startTime; - console.log( - "[handleTextureOutput] Processing completed in", - elapsedTime.toFixed(2), - "ms" - ); + } else { + // When all frames are captured, update the hidden widget and mark the capture complete. + node.setHiddenWidgetValue("animationFrames", node.frameBuffer); + iframe.contentWindow.postMessage({ type: "STOP_CAPTURE" }, "*"); + node.isCaptureComplete = true; + console.log( + "[handleTextureOutput] Completed capturing", + targetFrameCount, + "frames. Capture complete flag set." + ); + } } // Create the DOM widget for this node and register cleanup @@ -270,20 +296,43 @@ async function widgetIframe(node, nodeData, inputData, app) { // Updated onSerialize: Strip out the heavy "data" field to avoid localStorage quota problems. node.onSerialize = (nodeData) => { - const textureKeys = ["color", "canny", "depth", "normal"]; + // Always ensure a hidden data object exists nodeData.hidden = nodeData.hidden || {}; + + // Serialize the basic texture keys (keep any properties like base64 data) + const textureKeys = ["color", "canny", "depth", "normal"]; 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; - } + nodeData.hidden[key] = widgetEntry.value; } }); + + // Serialize the animation frames in full. + const animWidget = node.widgets.find((w) => w.name === "animationFrames"); + nodeData.hidden["animationFrames"] = + animWidget && animWidget.value ? animWidget.value : []; + console.log("[onSerialize] Serialized animation frames:", nodeData.hidden); + }; + + // Update triggerCapture to reset the frame buffer state and cache the dynamic frame count. + node.triggerCapture = function () { + // Reset previous capture state. + node.frameBuffer = []; + node.isCaptureComplete = false; + // Clear out previous animation frames from the hidden widget/global store. + node.setHiddenWidgetValue("animationFrames", []); + // Cache the target frame count from the widget (defaulting to 4 if not provided). + node.targetFrameCount = + parseInt(node.widgets.find((w) => w.name === "frame_count")?.value, 10) || + 4; + console.log( + "[triggerCapture] Reset capture state. Target frame count:", + node.targetFrameCount + ); + if (iframe.contentWindow) + iframe.contentWindow.postMessage({ type: "CAPTURE_TEXTURES" }, "*"); + console.log(" starting capture..."); }; return widget; @@ -326,60 +375,108 @@ app.registerExtension({ 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; + // Create and initialize the Iframe widget. 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, - }); + // Use node.id as the unique id to ensure consistency between JS and Python. + this.unique_id = this.id; + console.log( + "[widgetIframe] Using node id as unique id:", + this.unique_id + ); + + // Update or create the hidden unique id widget to be consistent in the global texture store. + 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; + + return originalOnNodeCreated + ? originalOnNodeCreated.apply(this, arguments) + : undefined; + }; + + // Here is the key: hook into the node processing + nodeType.prototype.onProcess = function () { + console.log("[onProcess] Node is now processing. Triggering capture."); + // Call the capture method only when the node is actually queued. + if (this.triggerCapture) this.triggerCapture(); }; } }, }); +function generateFallbackUniqueId() { + // Only used if no unique id is provided + return Math.random().toString(36).substr(2, 9); +} + +export function onNodeCreated(node) { + // Try to read the unique id from node.settings or node.hiddenInputs + let providedUniqueId = undefined; + if (node.settings && node.settings.unique_id) { + providedUniqueId = node.settings.unique_id; + } else if (node.hiddenInputs && node.hiddenInputs.unique_id) { + providedUniqueId = node.hiddenInputs.unique_id; + } + + // Force it to a string so that Python (which does str(unique_id)) can match correctly. + if (providedUniqueId) { + node.unique_id = String(providedUniqueId); + } else { + node.unique_id = generateFallbackUniqueId(); + } + + console.log("[widgetIframe] onNodeCreated: Using unique id:", node.unique_id); + + // Update or create a hidden widget for unique id to ensure consistency. + let uniqueIdWidget = node.widgets.find((w) => w.name === "unique_id"); + if (uniqueIdWidget) { + uniqueIdWidget.value = node.unique_id; + console.log( + "[widgetIframe] Updated hidden widget 'unique_id' to:", + node.unique_id + ); + } else { + node.widgets.push({ + name: "unique_id", + value: node.unique_id, + isHidden: true, + }); + console.log( + "[widgetIframe] Created hidden widget 'unique_id' with:", + node.unique_id + ); + } +} + export { widgetIframe }; + +// At the end of your file, or in a separate initialization block, +// override the global queuePrompt to trigger capture on IframeView nodes. + +const originalQueuePrompt = app.queuePrompt; +app.queuePrompt = function (...args) { + console.log("[QueuePrompt] Triggering capture for IframeView nodes."); + // Call the original queuePrompt functionality. + const result = originalQueuePrompt.apply(this, args); + // Find all nodes in the graph and trigger capture on those with comfyClass "IframeView". + if (app.graph && app.graph._nodes_by_id) { + Object.values(app.graph._nodes_by_id).forEach((node) => { + if ( + node.comfyClass === "IframeView" && + typeof node.triggerCapture === "function" + ) { + console.log(`[QueuePrompt] Triggering capture on node ${node.id}`); + node.triggerCapture(); + } + }); + } + return result; +};