File size: 5,080 Bytes
16b2f87 a4f9f3b 16b2f87 a4f9f3b 16b2f87 a4f9f3b 16b2f87 a4f9f3b 16b2f87 a4f9f3b 16b2f87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """ComfyUI server hooks for Oz ShotSplitter.
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__)
try:
from server import PromptServer # type: ignore
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:
return
try:
server = PromptServer.instance
except Exception as e:
logger.warning("Oz ShotSplitter: no PromptServer.instance (%s)", e)
return
from aiohttp import web
@server.routes.get("/oz_shotsplitter/active_clients")
async def _list_active_clients(_request):
sids = list(getattr(server, "sockets", {}).keys())
return web.json_response({"client_ids": sids, "count": len(sids)})
@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()
|