From 91590adf04e644304267a7a296710b4f7ae80bb2 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Fri, 17 May 2024 18:16:08 +0100
Subject: [PATCH] Add webcam node (#3497)

* Add webcam node

* unused import
---
 comfy_extras/nodes_webcam.py         |  33 ++++++++
 nodes.py                             |   1 +
 web/extensions/core/webcamCapture.js | 120 +++++++++++++++++++++++++++
 3 files changed, 154 insertions(+)
 create mode 100644 comfy_extras/nodes_webcam.py
 create mode 100644 web/extensions/core/webcamCapture.js

diff --git a/comfy_extras/nodes_webcam.py b/comfy_extras/nodes_webcam.py
new file mode 100644
index 00000000..32a0ba2f
--- /dev/null
+++ b/comfy_extras/nodes_webcam.py
@@ -0,0 +1,33 @@
+import nodes
+import folder_paths
+
+MAX_RESOLUTION = nodes.MAX_RESOLUTION
+
+
+class WebcamCapture(nodes.LoadImage):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("WEBCAM", {}),
+                "width": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
+                "height": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
+                "capture_on_queue": ("BOOLEAN", {"default": True}),
+            }
+        }
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "load_capture"
+
+    CATEGORY = "image"
+
+    def load_capture(s, image, **kwargs):
+        return super().load_image(folder_paths.get_annotated_filepath(image))
+
+
+NODE_CLASS_MAPPINGS = {
+    "WebcamCapture": WebcamCapture,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WebcamCapture": "Webcam Capture",
+}
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
index 37e7d734..34821ca3 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1961,6 +1961,7 @@ def init_custom_nodes():
         "nodes_align_your_steps.py",
         "nodes_attention_multiply.py",
         "nodes_advanced_samplers.py",
+        "nodes_webcam.py",
     ]
 
     import_failed = []
diff --git a/web/extensions/core/webcamCapture.js b/web/extensions/core/webcamCapture.js
new file mode 100644
index 00000000..ea1f597e
--- /dev/null
+++ b/web/extensions/core/webcamCapture.js
@@ -0,0 +1,120 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+const WEBCAM_READY = Symbol();
+
+app.registerExtension({
+	name: "Comfy.WebcamCapture",
+	getCustomWidgets(app) {
+		return {
+			WEBCAM(node, inputName) {
+				let res;
+				node[WEBCAM_READY] = new Promise((resolve) => (res = resolve));
+
+				const container = document.createElement("div");
+				container.style.background = "rgba(0,0,0,0.25)";
+				container.style.textAlign = "center";
+
+				const video = document.createElement("video");
+				video.style.height = video.style.width = "100%";
+
+				const loadVideo = async () => {
+					try {
+						const stream = await navigator.mediaDevices.getUserMedia({ video: true, audio: false });
+						container.replaceChildren(video);
+
+						setTimeout(() => res(video), 500); // Fallback as loadedmetadata doesnt fire sometimes?
+						video.addEventListener("loadedmetadata", () => res(video), false);
+						video.srcObject = stream;
+						video.play();
+					} catch (error) {
+						const label = document.createElement("div");
+						label.style.color = "red";
+						label.style.overflow = "auto";
+						label.style.maxHeight = "100%";
+						label.style.whiteSpace = "pre-wrap";
+						label.textContent = "Unable to load webcam, please ensure access is granted:\n" + error.message;
+						container.replaceChildren(label);
+					}
+				};
+
+				loadVideo();
+
+				return { widget: node.addDOMWidget(inputName, "WEBCAM", container) };
+			},
+		};
+	},
+	nodeCreated(node) {
+		if ((node.type, node.constructor.comfyClass !== "WebcamCapture")) return;
+
+		let video;
+		const camera = node.widgets.find((w) => w.name === "image");
+		const w = node.widgets.find((w) => w.name === "width");
+		const h = node.widgets.find((w) => w.name === "height");
+		const captureOnQueue = node.widgets.find((w) => w.name === "capture_on_queue");
+
+		const canvas = document.createElement("canvas");
+
+		const capture = () => {
+			canvas.width = w.value;
+			canvas.height = h.value;
+			const ctx = canvas.getContext("2d");
+			ctx.drawImage(video, 0, 0, w.value, h.value);
+			const data = canvas.toDataURL("image/png");
+
+			const img = new Image();
+			img.onload = () => {
+				node.imgs = [img];
+				app.graph.setDirtyCanvas(true);
+				requestAnimationFrame(() => {
+					node.setSizeForImage?.();
+				});
+			};
+			img.src = data;
+		};
+
+		const btn = node.addWidget("button", "waiting for camera...", "capture", capture);
+		btn.disabled = true;
+		btn.serializeValue = () => undefined;
+
+		camera.serializeValue = async () => {
+			if (captureOnQueue.value) {
+				capture();
+			} else if (!node.imgs?.length) {
+				const err = `No webcam image captured`;
+				alert(err);
+				throw new Error(err);
+			}
+
+			// Upload image to temp storage
+			const blob = await new Promise((r) => canvas.toBlob(r));
+			const name = `${+new Date()}.png`;
+			const file = new File([blob], name);
+			const body = new FormData();
+			body.append("image", file);
+			body.append("subfolder", "webcam");
+			body.append("type", "temp");
+			const resp = await api.fetchApi("/upload/image", {
+				method: "POST",
+				body,
+			});
+			if (resp.status !== 200) {
+				const err = `Error uploading camera image: ${resp.status} - ${resp.statusText}`;
+				alert(err);
+				throw new Error(err);
+			}
+			return `webcam/${name} [temp]`;
+		};
+
+		node[WEBCAM_READY].then((v) => {
+			video = v;
+			// If width isnt specified then use video output resolution
+			if (!w.value) {
+				w.value = video.videoWidth || 640;
+				h.value = video.videoHeight || 480; 
+			}
+			btn.disabled = false;
+			btn.label = "capture";
+		});
+	},
+});