working with animation sequences

This commit is contained in:
nofunstudio
2025-02-02 17:37:00 -07:00
parent 3895b77048
commit c2876f5a6b
3 changed files with 364 additions and 167 deletions
Binary file not shown.
+163 -63
View File
@@ -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,
)
)
+201 -104
View File
@@ -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;
};