Spaces:
Running
Running
File size: 6,314 Bytes
613c9ab |
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 |
import server
import folder_paths
import os
import time
import subprocess
from .utils import is_url, get_sorted_dir_files_from_directory, ffmpeg_path, validate_sequence
from comfy.k_diffusion.utils import FolderOfImages
web = server.web
def is_safe(path):
if "VHS_STRICT_PATHS" not in os.environ:
return True
basedir = os.path.abspath('.')
try:
common_path = os.path.commonpath([basedir, path])
except:
#Different drive on windows
return False
return common_path == basedir
@server.PromptServer.instance.routes.get("/viewvideo")
async def view_video(request):
query = request.rel_url.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 = filename
else:
filename, output_dir = folder_paths.annotated_filepath(filename)
type = request.rel_url.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(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(output_dir):
return web.Response(status=403)
if "subfolder" in request.rel_url.query:
output_dir = os.path.join(output_dir, request.rel_url.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)
if query.get('format', 'video') == "folder":
#Check that folder contains some valid image file, get it's extension
#ffmpeg seems to not support list globs, so support for mixed extensions seems unfeasible
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))
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 = ["-an", "-i", file]
args = [ffmpeg_path, "-v", "error"] + in_args
vfilters = []
if int(query.get('force_rate',0)) != 0:
vfilters.append("fps=fps="+query['force_rate'] + ":round=up:start_time=0.001")
if int(query.get('skip_first_frames', 0)) > 0:
vfilters.append(f"select=gt(n\\,{int(query['skip_first_frames'])-1})")
if int(query.get('select_every_nth', 1)) > 1:
vfilters.append(f"select=not(mod(n\\,{query['select_every_nth']}))")
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}")
vfilters.append("setpts=PTS-STARTPTS")
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
args += ['-c:v', 'libvpx-vp9','-deadline', 'realtime', '-cpu-used', '8', '-f', 'webm', '-']
try:
with subprocess.Popen(args, stdout=subprocess.PIPE) as proc:
try:
resp = web.StreamResponse()
resp.content_type = 'video/webm'
resp.headers["Content-Disposition"] = f"filename=\"{filename}\""
await resp.prepare(request)
while True:
bytes_read = proc.stdout.read()
if bytes_read is None:
#TODO: check for timeout here
time.sleep(.1)
continue
if len(bytes_read) == 0:
break
await resp.write(bytes_read)
except ConnectionResetError as e:
#Kill ffmpeg before stdout closes
proc.kill()
except BrokenPipeError as e:
pass
return resp
@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)
path = os.path.abspath(query["path"])
if not os.path.exists(path) or not is_safe(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] in valid_extensions:
valid_items.append(item.name)
except OSError:
#Broken symlinks can throw a very unhelpful "Invalid argument"
pass
return web.json_response(valid_items)
|