import server import folder_paths import os import subprocess import re import asyncio from .utils import is_url, get_sorted_dir_files_from_directory, ffmpeg_path, \ validate_sequence, is_safe_path, strip_path, try_download_video, ENCODE_ARGS from comfy.k_diffusion.utils import FolderOfImages web = server.web @server.PromptServer.instance.routes.get("/vhs/viewvideo") @server.PromptServer.instance.routes.get("/viewvideo") async def view_video(request): query = request.rel_url.query path_res = await resolve_path(query) if isinstance(path_res, web.Response): return path_res file, filename, output_dir = path_res if ffmpeg_path is None: #Don't just return file, that provides arbitrary read access to any file if is_safe_path(output_dir, strict=True): return web.FileResponse(path=file) frame_rate = query.get('frame_rate', 8) if query.get('format', 'video') == "folder": os.makedirs(folder_paths.get_temp_directory(), exist_ok=True) concat_file = os.path.join(folder_paths.get_temp_directory(), "image_sequence_preview.txt") skip_first_images = int(query.get('skip_first_images', 0)) select_every_nth = int(query.get('select_every_nth', 1)) or 1 valid_images = get_sorted_dir_files_from_directory(file, skip_first_images, select_every_nth, FolderOfImages.IMG_EXTENSIONS) if len(valid_images) == 0: return web.Response(status=400) with open(concat_file, "w") as f: f.write("ffconcat version 1.0\n") for path in valid_images: f.write("file '" + os.path.abspath(path) + "'\n") f.write("duration 0.125\n") in_args = ["-safe", "0", "-i", concat_file] else: in_args = ["-i", file] if '%' in file: in_args = ['-framerate', str(frame_rate)] + in_args #Do prepass to pull info #breaks skip_first frames if this default is ever actually needed base_fps = 30 try: proc = await asyncio.create_subprocess_exec(ffmpeg_path, *in_args, '-t', '0','-f', 'null','-', stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.DEVNULL) _, res_stderr = await proc.communicate() match = re.search(': Video: (\\w+) .+, (\\d+) fps,', res_stderr.decode(*ENCODE_ARGS)) if match: base_fps = float(match.group(2)) if match.group(1) == 'vp9': #force libvpx for transparency in_args = ['-c:v', 'libvpx-vp9'] + in_args except subprocess.CalledProcessError as e: print("An error occurred in the ffmpeg prepass:\n" \ + e.stderr.decode(*ENCODE_ARGS)) return web.Response(status=500) vfilters = [] target_rate = float(query.get('force_rate', 0)) or base_fps modified_rate = target_rate / float(query.get('select_every_nth',1) or 1) start_time = 0 if 'start_time' in query: start_time = float(query['start_time']) elif int(query.get('skip_first_frames', 0)) > 0: start_time = float(query.get('skip_first_frames'))/target_rate if start_time > 1/modified_rate: start_time += 1/modified_rate if start_time > 0: if start_time > 4: post_seek = ['-ss', '4'] pre_seek = ['-ss', str(start_time - 4)] else: post_seek = ['-ss', str(start_time)] pre_seek = [] else: pre_seek = [] post_seek = [] args = [ffmpeg_path, "-v", "error"] + pre_seek + in_args + post_seek if int(query.get('force_rate',0)) != 0: args += ['-r', str(modified_rate)] if query.get('force_size','Disabled') != "Disabled": size = query['force_size'].split('x') if size[0] == '?' or size[1] == '?': size[0] = "-2" if size[0] == '?' else f"'min({size[0]},iw)'" size[1] = "-2" if size[1] == '?' else f"'min({size[1]},ih)'" else: #Aspect ratio is likely changed. A more complex command is required #to crop the output to the new aspect ratio ar = float(size[0])/float(size[1]) vfilters.append(f"crop=if(gt({ar}\\,a)\\,iw\\,ih*{ar}):if(gt({ar}\\,a)\\,iw/{ar}\\,ih)") size = ':'.join(size) vfilters.append(f"scale={size}") if len(vfilters) > 0: args += ["-vf", ",".join(vfilters)] if int(query.get('frame_load_cap', 0)) > 0: args += ["-frames:v", query['frame_load_cap']] #TODO:reconsider adding high frame cap/setting default frame cap on node if query.get('deadline', 'realtime') == 'good': deadline = 'good' else: deadline = 'realtime' args += ['-c:v', 'libvpx-vp9','-deadline', deadline, '-cpu-used', '8', '-f', 'webm', '-'] try: proc = await asyncio.create_subprocess_exec(*args, stdout=subprocess.PIPE, stdin=subprocess.DEVNULL) try: resp = web.StreamResponse() resp.content_type = 'video/webm' resp.headers["Content-Disposition"] = f"filename=\"{filename}\"" await resp.prepare(request) while len(bytes_read := await proc.stdout.read(2**20)) != 0: await resp.write(bytes_read) #Of dubious value given frequency of kill calls, but more correct await proc.wait() except (ConnectionResetError, ConnectionError) as e: proc.kill() except BrokenPipeError as e: pass return resp query_cache = {} @server.PromptServer.instance.routes.get("/vhs/queryvideo") async def query_video(request): query = request.rel_url.query filepath = await resolve_path(query) #TODO: cache lookup if isinstance(filepath, web.Response): return filepath filepath = filepath[0] if filepath.endswith(".webp"): # ffmpeg doesn't support decoding animated WebP https://trac.ffmpeg.org/ticket/4907 return web.json_response({}) if filepath in query_cache and query_cache[filepath][0] == os.stat(filepath).st_mtime: source = query_cache[filepath][1] else: args_dummy = [ffmpeg_path, "-i", filepath, '-c', 'copy', '-frames:v', '1', "-f", "null", "-"] try: dummy_res = subprocess.run(args_dummy, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, check=True) except subprocess.CalledProcessError as e: raise Exception("An error occurred in the ffmpeg subprocess:\n" \ + e.stderr.decode(*ENCODE_ARGS)) lines = dummy_res.stderr.decode(*ENCODE_ARGS) source = {} for line in lines.split('\n'): match = re.search("^ *Stream .* Video.*, ([1-9]|\\d{2,})x(\\d+)", line) if match is not None: source['size'] = [int(match.group(1)), int(match.group(2))] fps_match = re.search(", ([\\d\\.]+) fps", line) if not fps_match: return web.json_response({}) source['fps'] = float(fps_match.group(1)) if re.search("(yuva|rgba)", line): source['alpha'] = True break else: raise Exception("Failed to parse video/image information. FFMPEG output:\n" + lines) durs_match = re.search("Duration: (\\d+:\\d+:\\d+\\.\\d+),", lines) if not (durs_match and 'fps' in source): return web.json_response({}) durs = durs_match.group(1).split(':') duration = int(durs[0])*360 + int(durs[1])*60 + float(durs[2]) source['duration'] = duration source['frames'] = int(duration*source['fps']) query_cache[filepath] = (os.stat(filepath).st_mtime, source) loaded = {} if 'duration' not in source: return web.json_response({}) loaded['duration'] = source['duration'] loaded['duration'] -= float(query.get('start_time',0)) loaded['fps'] = float(query.get('force_rate', 0)) or source['fps'] loaded['duration'] -= int(query.get('skip_first_frames', 0)) / loaded['fps'] loaded['fps'] /= int(query.get('select_every_nth', 1)) or 1 loaded['frames'] = loaded['duration'] * loaded['fps'] return web.json_response({'source': source, 'loaded': loaded}) async def resolve_path(query): if "filename" not in query: return web.Response(status=404) filename = query["filename"] #Path code misformats urls on windows and must be skipped if is_url(filename): file = await asyncio.to_thread(try_download_video, filename) or file filname, output_dir = os.path.split(file) return file, filename, output_dir else: filename, output_dir = folder_paths.annotated_filepath(filename) type = query.get("type", "output") if type == "path": #special case for path_based nodes #NOTE: output_dir may be empty, but non-None output_dir, filename = os.path.split(strip_path(filename)) if output_dir is None: output_dir = folder_paths.get_directory_by_type(type) if output_dir is None: return web.Response(status=400) if not is_safe_path(output_dir): return web.Response(status=403) if "subfolder" in query: output_dir = os.path.join(output_dir, query["subfolder"]) filename = os.path.basename(filename) file = os.path.join(output_dir, filename) if query.get('format', 'video') == 'folder': if not os.path.isdir(file): return web.Response(status=404) else: if not os.path.isfile(file) and not validate_sequence(file): return web.Response(status=404) return file, filename, output_dir @server.PromptServer.instance.routes.get("/vhs/getpath") @server.PromptServer.instance.routes.get("/getpath") async def get_path(request): query = request.rel_url.query if "path" not in query: return web.Response(status=404) #NOTE: path always ends in `/`, so this is functionally an lstrip path = os.path.abspath(strip_path(query["path"])) if not os.path.exists(path) or not is_safe_path(path): return web.json_response([]) #Use get so None is default instead of keyerror valid_extensions = query.get("extensions") valid_items = [] for item in os.scandir(path): try: if item.is_dir(): valid_items.append(item.name + "/") continue if valid_extensions is None or item.name.split(".")[-1].lower() in valid_extensions: valid_items.append(item.name) except OSError: #Broken symlinks can throw a very unhelpful "Invalid argument" pass valid_items.sort(key=lambda f: os.stat(os.path.join(path,f)).st_mtime) return web.json_response(valid_items)