ShotSplitter: add Download from URL button + yt-dlp route

ComfyUI-ShotSplitter/js/shot_splitter_ui.js

@@ -1,6 +1,7 @@
 // Oz_ShotSplitter UI:
 // 1. Upload Video button (uses ComfyUI /upload/image endpoint, type=input)
-// 2. Preview stack of <video> tags from ui.videos payload
+// 2. Download from URL button (POST /oz_shotsplitter/download_url, uses yt-dlp)
+// 3. Preview stack of <video> tags from ui.videos payload
 // Adapted from ComfyUI-VideoHelperSuite VHS.core.js (addVideoPreview pattern).
 
 import { app } from "../../scripts/app.js";
@@ -38,8 +39,24 @@ function clearContainer(container) {
     container.innerHTML = "";
 }
 
-// ---------- FEATURE 1: Upload Video Button ----------
-function addUploadButton(nodeType) {
+function setVideoWidget(node, filename) {
+    const videoWidget = node.widgets && node.widgets.find((w) => w.name === "video");
+    if (!videoWidget) return false;
+    if (!videoWidget.options) videoWidget.options = {};
+    if (!videoWidget.options.values) videoWidget.options.values = [];
+    if (!videoWidget.options.values.includes(filename)) {
+        videoWidget.options.values = videoWidget.options.values.filter(
+            (v) => v !== "<no videos in input/>"
+        );
+        videoWidget.options.values.push(filename);
+        videoWidget.options.values.sort();
+    }
+    videoWidget.value = filename;
+    return true;
+}
+
+// ---------- FEATURE 1+2: Upload Video / Download from URL ----------
+function addInputButtons(nodeType) {
     chainCallback(nodeType.prototype, "onNodeCreated", function () {
         if (this._ozUploadAdded) return;
         this._ozUploadAdded = true;
@@ -48,14 +65,25 @@ function addUploadButton(nodeType) {
         wrapper.style.cssText =
             "display:flex;flex-direction:column;gap:4px;padding:4px;width:100%;box-sizing:border-box;";
 
-        const btn = document.createElement("button");
-        btn.textContent = "📤 Upload Video";
-        btn.type = "button";
-        btn.style.cssText =
-            "width:100%;padding:6px 8px;cursor:pointer;background:#3a3a3a;color:#fff;" +
+        const btnStyle =
+            "width:100%;padding:6px 8px;cursor:pointer;color:#fff;" +
             "border:1px solid #555;border-radius:4px;font-size:12px;";
-        btn.addEventListener("mouseenter", () => (btn.style.background = "#4a4a4a"));
-        btn.addEventListener("mouseleave", () => (btn.style.background = "#3a3a3a"));
+
+        // -- Upload button --
+        const uploadBtn = document.createElement("button");
+        uploadBtn.textContent = "📤 Upload Video";
+        uploadBtn.type = "button";
+        uploadBtn.style.cssText = btnStyle + "background:#3a3a3a;";
+        uploadBtn.addEventListener("mouseenter", () => (uploadBtn.style.background = "#4a4a4a"));
+        uploadBtn.addEventListener("mouseleave", () => (uploadBtn.style.background = "#3a3a3a"));
+
+        // -- URL download button --
+        const urlBtn = document.createElement("button");
+        urlBtn.textContent = "🔗 Download from URL";
+        urlBtn.type = "button";
+        urlBtn.style.cssText = btnStyle + "background:#2c4a6e;";
+        urlBtn.addEventListener("mouseenter", () => (urlBtn.style.background = "#345a86"));
+        urlBtn.addEventListener("mouseleave", () => (urlBtn.style.background = "#2c4a6e"));
 
         const status = document.createElement("div");
         status.style.cssText = "font-size:10px;color:#888;text-align:center;min-height:12px;";
@@ -65,53 +93,36 @@ function addUploadButton(nodeType) {
         fileInput.accept = "video/*";
         fileInput.style.display = "none";
 
-        btn.addEventListener("click", () => fileInput.click());
+        uploadBtn.addEventListener("click", () => fileInput.click());
 
         const node = this;
+
         fileInput.addEventListener("change", async (e) => {
             const file = e.target.files && e.target.files[0];
             if (!file) return;
 
-            btn.disabled = true;
-            btn.textContent = `⏳ Uploading...`;
+            uploadBtn.disabled = true;
+            urlBtn.disabled = true;
+            uploadBtn.textContent = `⏳ Uploading...`;
             status.textContent = `${file.name} (${(file.size / 1024 / 1024).toFixed(1)} MB)`;
 
             try {
                 const fd = new FormData();
-                // ComfyUI's /upload/image accepts arbitrary file types
                 fd.append("image", file, file.name);
                 fd.append("type", "input");
                 fd.append("subfolder", "");
                 fd.append("overwrite", "true");
 
-                const res = await fetch("/upload/image", {
-                    method: "POST",
-                    body: fd,
-                });
-                if (!res.ok) {
-                    throw new Error(`HTTP ${res.status} ${res.statusText}`);
-                }
+                const res = await fetch("/upload/image", { method: "POST", body: fd });
+                if (!res.ok) throw new Error(`HTTP ${res.status} ${res.statusText}`);
                 const data = await res.json();
                 const uploadedName = data.name || file.name;
 
-                // Update the video widget options + value
-                const videoWidget = node.widgets && node.widgets.find((w) => w.name === "video");
-                if (videoWidget) {
-                    if (!videoWidget.options) videoWidget.options = {};
-                    if (!videoWidget.options.values) videoWidget.options.values = [];
-                    if (!videoWidget.options.values.includes(uploadedName)) {
-                        // Drop the placeholder if present
-                        videoWidget.options.values = videoWidget.options.values.filter(
-                            (v) => v !== "<no videos in input/>"
-                        );
-                        videoWidget.options.values.push(uploadedName);
-                        videoWidget.options.values.sort();
-                    }
-                    videoWidget.value = uploadedName;
-                    btn.textContent = `✅ ${uploadedName}`;
+                if (setVideoWidget(node, uploadedName)) {
+                    uploadBtn.textContent = `✅ ${uploadedName}`;
                     status.textContent = "Video selected in dropdown.";
                 } else {
-                    btn.textContent = `⚠️ video widget not found`;
+                    uploadBtn.textContent = `⚠️ video widget not found`;
                 }
 
                 if (node.graph) {
@@ -120,18 +131,68 @@ function addUploadButton(nodeType) {
                 }
             } catch (err) {
                 console.error("[Oz ShotSplitter] Upload failed:", err);
-                btn.textContent = `❌ Upload failed`;
+                uploadBtn.textContent = `❌ Upload failed`;
                 status.textContent = String(err.message || err);
             } finally {
-                btn.disabled = false;
+                uploadBtn.disabled = false;
+                urlBtn.disabled = false;
                 fileInput.value = "";
                 setTimeout(() => {
-                    btn.textContent = "📤 Upload Video";
+                    uploadBtn.textContent = "📤 Upload Video";
                 }, 6000);
             }
         });
 
-        wrapper.appendChild(btn);
+        urlBtn.addEventListener("click", async () => {
+            const url = (prompt("Paste video URL (Instagram / TikTok / YouTube / …):") || "").trim();
+            if (!url) return;
+            if (!/^https?:\/\//i.test(url)) {
+                status.textContent = "Must start with http(s)://";
+                return;
+            }
+
+            uploadBtn.disabled = true;
+            urlBtn.disabled = true;
+            urlBtn.textContent = "⏳ Downloading...";
+            status.textContent = "yt-dlp running (may take a minute)...";
+
+            try {
+                const res = await fetch("/oz_shotsplitter/download_url", {
+                    method: "POST",
+                    headers: { "Content-Type": "application/json" },
+                    body: JSON.stringify({ url }),
+                });
+                const data = await res.json().catch(() => ({ ok: false, error: `HTTP ${res.status}` }));
+                if (!res.ok || !data.ok) {
+                    throw new Error(data.error || `HTTP ${res.status}`);
+                }
+                const filename = data.filename;
+                if (setVideoWidget(node, filename)) {
+                    urlBtn.textContent = `✅ ${filename}${data.cached ? " (cached)" : ""}`;
+                    status.textContent = "Video selected in dropdown.";
+                } else {
+                    urlBtn.textContent = `⚠️ video widget not found`;
+                }
+
+                if (node.graph) {
+                    node.setSize(node.computeSize());
+                    node.graph.setDirtyCanvas(true, true);
+                }
+            } catch (err) {
+                console.error("[Oz ShotSplitter] URL download failed:", err);
+                urlBtn.textContent = "❌ Download failed";
+                status.textContent = String(err.message || err);
+            } finally {
+                uploadBtn.disabled = false;
+                urlBtn.disabled = false;
+                setTimeout(() => {
+                    urlBtn.textContent = "🔗 Download from URL";
+                }, 8000);
+            }
+        });
+
+        wrapper.appendChild(uploadBtn);
+        wrapper.appendChild(urlBtn);
         wrapper.appendChild(status);
         wrapper.appendChild(fileInput);
 
@@ -144,7 +205,7 @@ function addUploadButton(nodeType) {
             setValue() {},
         });
         widget.computeSize = function (width) {
-            return [width, 56];
+            return [width, 90];
         };
     });
 }
@@ -233,7 +294,7 @@ app.registerExtension({
     name: "oz.shot_splitter",
     async beforeRegisterNodeDef(nodeType, nodeData) {
         if (nodeData && nodeData.name === "Oz_ShotSplitter") {
-            addUploadButton(nodeType);
+            addInputButtons(nodeType);
             addShotSplitterPreview(nodeType);
         }
     },

ComfyUI-ShotSplitter/server_hooks.py

@@ -1,14 +1,17 @@
 """ComfyUI server hooks for Oz ShotSplitter.
 
-Exposes a lightweight debug endpoint that lists active WebSocket client IDs
-so a driver script can route prompt submissions to the user's browser session
-(required because ComfyUI gates `executed` events on a non-None client_id in
-execution.py — events submitted without client_id are silently dropped).
+Routes:
+  GET  /oz_shotsplitter/active_clients  — debug list of active websocket clients
+  POST /oz_shotsplitter/download_url    — yt-dlp a public URL into ComfyUI input/
 """
 
 from __future__ import annotations
 
+import asyncio
+import hashlib
 import logging
+import os
+import re
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +21,33 @@ except Exception as e:  # pragma: no cover
     PromptServer = None  # type: ignore
     logger.warning("Oz ShotSplitter: PromptServer unavailable (%s)", e)
 
+DOWNLOAD_TIMEOUT_SEC = 600
+DOWNLOAD_MAX_FILESIZE = "500M"
+VIDEO_EXTS = ("mp4", "mkv", "webm", "mov", "m4v")
+_INPUT_DIR_CACHE: str | None = None
+
+
+def _resolve_input_dir() -> str:
+    global _INPUT_DIR_CACHE
+    if _INPUT_DIR_CACHE:
+        return _INPUT_DIR_CACHE
+    try:
+        import folder_paths  # type: ignore
+        d = folder_paths.get_input_directory()
+    except Exception:
+        d = "/workspace/ComfyUI/input"
+    os.makedirs(d, exist_ok=True)
+    _INPUT_DIR_CACHE = d
+    return d
+
+
+def _find_cached(input_dir: str, stem: str) -> str | None:
+    for ext in VIDEO_EXTS:
+        p = os.path.join(input_dir, f"{stem}.{ext}")
+        if os.path.exists(p) and os.path.getsize(p) > 0:
+            return p
+    return None
+
 
 def register_routes() -> None:
     if PromptServer is None:
@@ -35,7 +65,98 @@ def register_routes() -> None:
         sids = list(getattr(server, "sockets", {}).keys())
         return web.json_response({"client_ids": sids, "count": len(sids)})
 
-    logger.info("Oz ShotSplitter: registered /oz_shotsplitter/active_clients")
+    @server.routes.post("/oz_shotsplitter/download_url")
+    async def _download_url(request):
+        try:
+            data = await request.json()
+        except Exception:
+            return web.json_response({"ok": False, "error": "invalid JSON body"}, status=400)
+
+        url = (data.get("url") or "").strip()
+        if not url or not re.match(r"^https?://", url):
+            return web.json_response({"ok": False, "error": "invalid URL"}, status=400)
+
+        input_dir = _resolve_input_dir()
+        url_hash = hashlib.md5(url.encode("utf-8")).hexdigest()[:12]
+        stem = f"yt_{url_hash}"
+
+        cached = _find_cached(input_dir, stem)
+        if cached:
+            return web.json_response({
+                "ok": True,
+                "filename": os.path.basename(cached),
+                "cached": True,
+            })
+
+        out_template = os.path.join(input_dir, f"{stem}.%(ext)s")
+        cmd = [
+            "yt-dlp",
+            url,
+            "-o", out_template,
+            "--no-playlist",
+            "--max-filesize", DOWNLOAD_MAX_FILESIZE,
+            "-f", "mp4/best[ext=mp4]/bestvideo*+bestaudio/best",
+            "--merge-output-format", "mp4",
+            "--no-warnings",
+            "--no-progress",
+            "--quiet",
+        ]
+        logger.info("Oz ShotSplitter: yt-dlp %s -> %s", url, stem)
+
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+        except FileNotFoundError:
+            return web.json_response({
+                "ok": False,
+                "error": "yt-dlp not installed. Run: /venv/main/bin/pip install yt-dlp",
+            }, status=500)
+        except Exception as e:
+            return web.json_response({"ok": False, "error": f"spawn failed: {e}"}, status=500)
+
+        try:
+            stdout, stderr = await asyncio.wait_for(
+                proc.communicate(), timeout=DOWNLOAD_TIMEOUT_SEC
+            )
+        except asyncio.TimeoutError:
+            try:
+                proc.kill()
+                await proc.wait()
+            except Exception:
+                pass
+            return web.json_response({
+                "ok": False,
+                "error": f"yt-dlp timed out after {DOWNLOAD_TIMEOUT_SEC}s",
+            }, status=504)
+
+        if proc.returncode != 0:
+            err_tail = (stderr or b"").decode("utf-8", errors="replace").strip()
+            if len(err_tail) > 600:
+                err_tail = err_tail[-600:]
+            return web.json_response({
+                "ok": False,
+                "error": f"yt-dlp exit {proc.returncode}: {err_tail or 'no stderr'}",
+            }, status=500)
+
+        saved = _find_cached(input_dir, stem)
+        if not saved:
+            return web.json_response({
+                "ok": False,
+                "error": "yt-dlp succeeded but output file missing",
+            }, status=500)
+
+        return web.json_response({
+            "ok": True,
+            "filename": os.path.basename(saved),
+            "cached": False,
+        })
+
+    logger.info(
+        "Oz ShotSplitter: registered /oz_shotsplitter/{active_clients,download_url}"
+    )
 
 
 register_routes()