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): """ 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. Data type:", type(data)) 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: # 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.") 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] b64data is not a string. Type:", type(b64data)) width = data.get("width", 64) height = data.get("height", 64) 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}") 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 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-conversion image stats: shape={image.shape}, min={image.min()}, max={image.max()}") image = image.astype(np.float32) / 255.0 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((64, 64, 3), dtype=torch.float32) return result async def wait_for_capture_completion(unique_id, target_frame_count, timeout=1.0, interval=0.1): 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): 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", { "default": "{\"camera\": {\"position\": [0,0,5]}, \"animation\": false}", }, ), "frame_count": ("INT", {"default": 4, "min": 1, "max": 400}), }, "hidden": { "unique_id": "UNIQUE_ID", "animationFrames": "TEXTURE_SEQUENCE", }, } # Only return animation textures now. RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("animation_color", "animation_canny", "animation_depth", "animation_normal") FUNCTION = "process_iframe" CATEGORY = "lth" async def _process_iframe_async( self, url, width, height, scene_state, frame_count, unique_id, animationFrames=None, ): start_time = time.time() print("[process_iframe] Starting asynchronous processing at", start_time) unique_id = str(unique_id) 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) # Process each frame's texture into a tensor for animation outputs. 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"] ] 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 heavy animationFrames data. animationFrames = None elapsed = time.time() - start_time print(f"[process_iframe] Asynchronous processing completed in {elapsed:.4f} seconds.") return (animation_color, animation_canny, animation_depth, animation_normal) def process_iframe( self, url, width, height, scene_state, frame_count, unique_id, animationFrames=None, ): return asyncio.run( self._process_iframe_async( url, width, height, scene_state, frame_count, unique_id, animationFrames, ) ) 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"]