diff --git a/CHANGELOG.md b/CHANGELOG.md index d049bc94..bebbfc3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ ### Added +- **Upscale finished renders** (Settings ▸ Editor): run a `models/upscale_models` model over a + finished video, keeping its fps and sound; the result replaces the render. Never, always, or + on demand with the Upscale button beside the rating. - **Shot memory** (experimental, Chain Sampler ▸ Guidance): new gens can reuse the layout of a liked gen with fresh details. Ratings learn when to reuse, which shot and how strongly. No extra model call. diff --git a/__init__.py b/__init__.py index 3a818c02..7b079d90 100644 --- a/__init__.py +++ b/__init__.py @@ -48,6 +48,7 @@ if __package__: from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler from .templates import FunPackRefinementKeyLoader + from .video_upscale import FunPackUpscaleVideo try: from . import batch_training # noqa: F401 registers /funpack/batch/* routes except Exception as _e: @@ -104,6 +105,10 @@ else: from templates import FunPackRefinementKeyLoader except Exception: FunPackRefinementKeyLoader = None + try: + from video_upscale import FunPackUpscaleVideo + except Exception: + FunPackUpscaleVideo = None WEB_DIRECTORY = "./web" @@ -132,6 +137,7 @@ NODE_CLASS_MAPPINGS = { "FunPackCLIPLoader": FunPackCLIPLoader, "FunPackVAELoader": FunPackVAELoader, "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader, + "FunPackUpscaleVideo": FunPackUpscaleVideo, } NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None} @@ -160,6 +166,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FunPackCLIPLoader": "FunPack CLIP Loader", "FunPackVAELoader": "FunPack VAE Loader", "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader", + "FunPackUpscaleVideo": "FunPack Upscale Video", } __all__ = [ diff --git a/movie_editor/frontend/api.js b/movie_editor/frontend/api.js index 5bd9f183..79f7147a 100644 --- a/movie_editor/frontend/api.js +++ b/movie_editor/frontend/api.js @@ -185,6 +185,8 @@ generate: (id, onlyScene, sceneIds, resetSession, nodeOverrides, prevSceneMedia, project) => j("POST", API(`/projects/${id}/generate`), { only_scene: onlyScene || null, scene_ids: sceneIds || null, reset_session: !!resetSession, node_overrides: nodeOverrides || null, simple: !!window.FunPackMode?.isSimple(), prev_scene_media: prevSceneMedia || null, project: project || null }), status: (id, promptId) => j("GET", API(`/projects/${id}/status/${promptId}`)), + upscaleModels: () => j("GET", API("/upscale_models")), + upscale: (id, media, model) => j("POST", API(`/projects/${id}/upscale`), { media, model }), progress: () => j("GET", API("/progress")), // The editor's own in-flight generation, recovered from ComfyUI's queue (survives a UI reload). active: () => j("GET", API("/active")), diff --git a/movie_editor/frontend/editor_settings.js b/movie_editor/frontend/editor_settings.js index f33a0494..e4e3a35c 100644 --- a/movie_editor/frontend/editor_settings.js +++ b/movie_editor/frontend/editor_settings.js @@ -213,6 +213,33 @@ return row; } + // Post-render upscale: a models/upscale_models model run over a finished render, which it + // then replaces. Always = every render when it finishes; button = "Upscale" beside the rating. + function upscaleSection() { + const wrap = el("div", "es-section"); + wrap.append(el("div", "es-section-title", "Upscale finished renders")); + wrap.append(el("div", "es-hint", + "Runs an upscale model over a finished video and replaces the render with the result. " + + "Models come from ComfyUI/models/upscale_models.")); + const mode = el("select"); + [["never", "Never"], ["button", "When I press Upscale (beside the rating)"], ["always", "Always, when a render finishes"]] + .forEach(([v, t]) => mode.append(new Option(t, v, false, (S.getEditorSetting("upscaleMode") || "never") === v))); + mode.onchange = () => S.setEditorSetting("upscaleMode", mode.value); + wrap.append(labeled("Upscale", mode)); + const model = el("select"); + model.append(new Option("Loading…", "")); + wrap.append(labeled("Model", model)); + API.upscaleModels().then((r) => { + const list = (r && r.models) || []; + const cur = S.getEditorSetting("upscaleModel") || ""; + model.replaceChildren(new Option(list.length ? "— pick a model —" : "Nothing in models/upscale_models", "")); + list.forEach((m) => model.append(new Option(m, m, false, m === cur))); + if (cur && !list.includes(cur)) model.append(new Option(`${cur} (missing)`, cur, false, true)); + }).catch(() => model.replaceChildren(new Option("Could not list models", ""))); + model.onchange = () => S.setEditorSetting("upscaleModel", model.value); + return wrap; + } + function mount(body) { const content = el("div", "es-content"); content.append(toggleRow( @@ -229,13 +256,14 @@ "anchorEnabled")); content.append(revolverSection()); content.append(anchorGuideSection()); + content.append(upscaleSection()); body.append(content); } window.SettingsWindow.register({ id: "editor", group: "", title: "Editor", subtitle: "How the editor behaves. The open project remembers these, so they follow it to another machine.", - keywords: "autocomplete anchor prompt shortcuts ideas suggestions i2v bypass guide preferences revolver random replacements cycle", + keywords: "upscale upscaler esrgan resolution autocomplete anchor prompt shortcuts ideas suggestions i2v bypass guide preferences revolver random replacements cycle", iconBg: "linear-gradient(180deg,#6aa9ff,#3b6fd9)", icon: '', mount, diff --git a/movie_editor/frontend/store.js b/movie_editor/frontend/store.js index 82fd770b..7051368d 100644 --- a/movie_editor/frontend/store.js +++ b/movie_editor/frontend/store.js @@ -3464,6 +3464,67 @@ } } + // Post-render upscale (video_upscale.py): one ComfyUI job per file, then the upscaled file + // replaces the original under every scene and ghost that plays it (a chain run is one file). + const _upscaling = new Set(); + function _sameMedia(a, b) { + return !!(a && b && a.filename === b.filename && (a.subfolder || "") === (b.subfolder || "") + && (a.type || "output") === (b.type || "output")); + } + function _swapRenderMedia(from, to) { + let n = 0; + for (const id of Object.keys(state.sceneRenders || {})) { + const r = state.sceneRenders[id]; + if (r && _sameMedia(r.media, from)) { state.sceneRenders[id] = { ...r, media: to }; n++; } + } + state.sceneGhosts = (state.sceneGhosts || []).map((g) => (_sameMedia(g.media, from) ? { ...g, media: to } : g)); + if (n) { _validateRendersToken++; _syncEditorStateToProject(); scheduleSaveSilent(); } + return n; + } + async function upscaleMedia(media) { + const model = getEditorSetting("upscaleModel"); + if (!media || !media.filename) return; + if (!model) { + set({ notice: "Upscale: pick a model in Settings ▸ Editor ▸ Upscale finished renders." }); + return; + } + const key = `${media.type || "output"}/${media.subfolder || ""}/${media.filename}`; + const pid = state.project && state.project.id; + if (!pid || _upscaling.has(key)) return; + _upscaling.add(key); + set({ notice: `Upscaling with ${model}…` }); + let notice = ""; + try { + const { prompt_id: promptId } = await API.upscale(pid, media, model); + for (;;) { + await new Promise((res) => setTimeout(res, 2000)); + const s = await API.status(pid, promptId); + if (s.state === "error") throw new Error(s.error || "stopped or failed inside ComfyUI"); + if (s.state !== "completed") continue; + const out = (s.media || []).find((m) => m.kind === "videos"); + if (!out) throw new Error(s.error || "no video came out, check the ComfyUI terminal"); + if (state.project && state.project.id === pid && !_swapRenderMedia(media, out)) { + notice = "Upscaled, but that render was replaced meanwhile, so nothing was swapped."; + } + break; + } + } catch (e) { + notice = `Upscale failed: ${e && e.message ? e.message : e}`; + } finally { + _upscaling.delete(key); + set({ notice }); + } + } + function upscaleRender(sceneId) { + const r = (state.sceneRenders || {})[sceneId]; + return upscaleMedia(r && r.media); + } + function isUpscaling(sceneId) { + const r = (state.sceneRenders || {})[sceneId]; + const m = r && r.media; + return !!(m && _upscaling.has(`${m.type || "output"}/${m.subfolder || ""}/${m.filename}`)); + } + // Poll a single queued prompt to completion. Resolves true on success, false on error. function _pollPromise(promptId, targetSceneIds, prefix) { prefix = prefix || "Generating…"; @@ -3522,6 +3583,9 @@ _clearGenTimers(); _recordSegment(s.media, targetSceneIds, { sceneLayout: s.scene_layout }); set({ gen: { state: "done", promptId, media: s.media, msg: s.media.length ? "" : "Completed but no output media found — check ComfyUI terminal." } }); + if (getEditorSetting("upscaleMode") === "always") { + upscaleMedia((s.media || []).find((m) => m.kind === "videos" || m.kind === "gifs")); + } resolve(true); } else { // "pending" after "running" means the job left the queue without a history @@ -4739,6 +4803,8 @@ }); window.Store = { + upscaleRender, + isUpscaling, get, set, subscribe, notify, init, scheduleSaveFromHistory, notifyHistoryState, refreshProjectList, loadProject, newProject, deleteProject, downloadProject, importProject, diff --git a/movie_editor/frontend/timeline.js b/movie_editor/frontend/timeline.js index 77a9d120..10da977e 100644 --- a/movie_editor/frontend/timeline.js +++ b/movie_editor/frontend/timeline.js @@ -1791,9 +1791,18 @@ // it without making the user click it first. const scenes = p?.scenes || []; const sceneId = st.selectedSceneId || (scenes.length === 1 ? scenes[0].id : null); - if (!studioOn || !sceneId || !hasRender(st, sceneId) || !Picker) return wrap; + if (!sceneId || !hasRender(st, sceneId)) return wrap; const sc = S.scene(sceneId); if (!sc || !S.isGenerativeScene(sc)) return wrap; + if (S.getEditorSetting("upscaleMode") === "button") { + const busy = S.isUpscaling(sceneId); + const up = el("button", "btn ghost tiny tl-upscale-btn", busy ? "Upscaling…" : "Upscale"); + up.title = "Upscale this video with the model picked in Settings ▸ Editor; it replaces the render"; + up.disabled = busy; + up.onclick = (e) => { e.stopPropagation(); S.upscaleRender(sceneId); }; + wrap.append(up); + } + if (!studioOn || !Picker) return wrap; const sceneNo = p.scenes.indexOf(sc) + 1; const raw = sceneRatingRaw(st, sc); const rlabel = el("span", "tl-keys", `Scene ${sceneNo}`); @@ -1806,8 +1815,7 @@ S.setSceneRating(sc.id, val); }, { h3: !!window.PipelineCaps?.isH3(st) }); }; - wrap.append(rlabel); - wrap.append(btn); + wrap.prepend(rlabel, btn); return wrap; } diff --git a/movie_editor/server.py b/movie_editor/server.py index bbd60ee0..af2db475 100644 --- a/movie_editor/server.py +++ b/movie_editor/server.py @@ -2071,6 +2071,32 @@ if web is not None and PromptServer is not None: running = True return web.json_response({"state": "running" if running else "pending", "media": []}) + @routes.get(UI_PREFIX + "/api/upscale_models") + async def _upscale_models(_req): + try: + import folder_paths + return web.json_response({"models": folder_paths.get_filename_list("upscale_models")}) + except Exception as e: # noqa: BLE001 + return web.json_response({"models": [], "error": str(e)}) + + @routes.post(UI_PREFIX + "/api/projects/{pid}/upscale") + async def _upscale(req): + """Queue FunPackUpscaleVideo on one finished render; poll /status like a generation.""" + _project_or_404(req.match_info["pid"]) + body = await req.json() if req.can_read_body else {} + media, model = body.get("media") or {}, body.get("model") or "" + if not media.get("filename") or not model: + raise web.HTTPBadRequest(reason="pick an upscale model and a rendered video") + graph = {"1": {"class_type": "FunPackUpscaleVideo", "inputs": { + "filename": media["filename"], "subfolder": media.get("subfolder") or "", + "type": "temp" if media.get("type") == "temp" else "output", + "upscale_model": model}}} + try: + data = await bridge.queue_prompt(graph) + except Exception as e: # noqa: BLE001 + return web.json_response({"detail": bridge.format_funpack_error(e)}, status=502) + return web.json_response({"prompt_id": data.get("prompt_id")}) + @routes.get(UI_PREFIX + "/api/rating-labels") async def _rating_labels(_req): try: diff --git a/tests/test_video_upscale.py b/tests/test_video_upscale.py new file mode 100644 index 00000000..689fbf2b --- /dev/null +++ b/tests/test_video_upscale.py @@ -0,0 +1,66 @@ +"""video_upscale.upscale_file: streams a real clip through an upscaler in chunks and keeps +its frame count, frame rate and sound; a failed run leaves no half-written file.""" +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import video_upscale as vu # noqa: E402 + +pytestmark = pytest.mark.skipif(not shutil.which("ffmpeg"), reason="needs ffmpeg") + + +def _clip(path, frames=21, fps=24, sound=True): + cmd = ["ffmpeg", "-y", "-v", "error", "-f", "lavfi", "-i", + f"testsrc=size=64x48:rate={fps}:duration={frames / fps}"] + if sound: + cmd += ["-f", "lavfi", "-i", f"sine=duration={frames / fps}", "-c:a", "aac"] + subprocess.run(cmd + ["-c:v", "libx264", "-pix_fmt", "yuv420p", str(path)], check=True) + + +def _count(path): + out = subprocess.run(["ffprobe", "-v", "error", "-count_frames", "-select_streams", "v:0", + "-show_entries", "stream=nb_read_frames", "-of", "csv=p=0", str(path)], + capture_output=True, text=True, check=True) + return int(out.stdout.strip()) + + +def _twice(x): + return x.repeat_interleave(2, dim=1).repeat_interleave(2, dim=2) + + +def test_upscales_every_frame_and_keeps_fps_and_sound(tmp_path): + src, dst = tmp_path / "in.mp4", tmp_path / "out.mp4" + _clip(src) + calls = [] + frames = vu.upscale_file(str(src), str(dst), lambda x: calls.append(x.shape[0]) or _twice(x)) + assert frames == 21 == _count(dst) + assert max(calls) <= vu.CHUNK # streamed, not the whole clip + w, h, fps, sound = vu.probe(str(dst)) + assert (w, h, fps, sound) == (128, 96, "24/1", True) + + +def test_a_silent_clip_stays_silent(tmp_path): + src, dst = tmp_path / "in.mp4", tmp_path / "out.mp4" + _clip(src, sound=False) + vu.upscale_file(str(src), str(dst), _twice) + assert vu.probe(str(dst))[3] is False + + +def test_a_stopped_run_leaves_no_file(tmp_path): + src, dst = tmp_path / "in.mp4", tmp_path / "out.mp4" + _clip(src) + seen = [] + + def stop(): + seen.append(1) + if len(seen) > 1: + raise RuntimeError("interrupted") + with pytest.raises(RuntimeError, match="interrupted"): + vu.upscale_file(str(src), str(dst), _twice, stop) + assert not dst.exists() diff --git a/video_upscale.py b/video_upscale.py new file mode 100644 index 00000000..8b75c994 --- /dev/null +++ b/video_upscale.py @@ -0,0 +1,153 @@ +"""Upscale a finished video with a models/upscale_models model (ESRGAN-style, via spandrel). + +The Editor queues this on a render once it is ready: frames stream through ffmpeg in small +chunks (never the whole clip in RAM), each chunk goes through ComfyUI's own +ImageUpscaleWithModel (tiling, OOM back-off, device handling), and the result is encoded +with the source's frame rate and sound. +""" + +import json +import os +import shutil +import subprocess + +import numpy as np +import torch + +CHUNK = 8 # frames per upscale call +SUBFOLDER = "funpack_upscaled" + + +def _ffmpeg(name): + path = shutil.which(name) + if not path: + raise RuntimeError(f"{name} not found on PATH; install ffmpeg to upscale videos.") + return path + + +def probe(path): + """-> (width, height, fps string, has sound).""" + out = subprocess.run( + [_ffmpeg("ffprobe"), "-v", "error", "-show_entries", + "stream=codec_type,width,height,r_frame_rate", "-of", "json", path], + capture_output=True, text=True, check=True) + streams = json.loads(out.stdout).get("streams") or [] + video = next((s for s in streams if s.get("codec_type") == "video"), None) + if not video: + raise RuntimeError(f"no video stream in {os.path.basename(path)}") + has_sound = any(s.get("codec_type") == "audio" for s in streams) + return int(video["width"]), int(video["height"]), video.get("r_frame_rate") or "25", has_sound + + +def resolve(filename, subfolder="", type_="output"): + """Absolute path of a ComfyUI output/temp file; refuses anything outside that folder.""" + import folder_paths + base = folder_paths.get_output_directory() if type_ == "output" else folder_paths.get_temp_directory() + base = os.path.realpath(base) + path = os.path.realpath(os.path.join(base, subfolder or "", filename)) + if os.path.commonpath([base, path]) != base: + raise ValueError("video path escapes the output folder") + if not os.path.isfile(path): + raise FileNotFoundError(f"{filename} is not in the {type_} folder (anymore)") + return path + + +def _read(stream, n): + buf = bytearray() + while len(buf) < n: + part = stream.read(n - len(buf)) + if not part: + break + buf += part + return bytes(buf) + + +def upscale_file(src, dst, upscale, interrupted=lambda: None): + """Stream `src` through `upscale` ([B,H,W,3] float 0-1 -> same, larger) into `dst`. + -> frames written. `interrupted()` raises to stop between chunks.""" + w, h, fps, has_sound = probe(src) + frame = w * h * 3 + dec = subprocess.Popen([_ffmpeg("ffmpeg"), "-v", "error", "-i", src, "-f", "rawvideo", + "-pix_fmt", "rgb24", "-"], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) + enc, frames = None, 0 + try: + while True: + interrupted() + raw = _read(dec.stdout, frame * CHUNK) + n = len(raw) // frame + if n == 0: + break + x = torch.from_numpy(np.frombuffer(raw[:n * frame], dtype=np.uint8) + .reshape(n, h, w, 3).copy()).float() / 255.0 + y = (upscale(x).clamp(0, 1) * 255.0).round().to(torch.uint8).cpu().numpy() + if enc is None: + oh, ow = y.shape[1], y.shape[2] + cmd = [_ffmpeg("ffmpeg"), "-y", "-v", "error", "-f", "rawvideo", + "-pix_fmt", "rgb24", "-s", f"{ow}x{oh}", "-r", fps, "-i", "-"] + if has_sound: + cmd += ["-i", src, "-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac", "-b:a", "192k"] + cmd += ["-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-c:v", "libx264", + "-crf", "16", "-pix_fmt", "yuv420p", "-movflags", "+faststart", dst] + enc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE) + enc.stdin.write(y.tobytes()) + frames += n + if enc is None: + raise RuntimeError(f"no frames could be read from {os.path.basename(src)}") + enc.stdin.close() + err = enc.stderr.read().decode(errors="replace") + if enc.wait() != 0: + raise RuntimeError(f"ffmpeg could not encode the upscaled video: {err[-500:]}") + return frames + except BaseException: + for p in (dec, enc): + if p is not None and p.poll() is None: + p.kill() + if os.path.exists(dst): + os.remove(dst) + raise + finally: + if dec.poll() is None: + dec.kill() + + +class FunPackUpscaleVideo: + """Upscale a finished video file with an upscale model; keeps its fps and sound.""" + + @classmethod + def INPUT_TYPES(cls): + import folder_paths + return {"required": { + "filename": ("STRING", {"default": ""}), + "subfolder": ("STRING", {"default": ""}), + "type": (["output", "temp"], {"default": "output"}), + "upscale_model": (folder_paths.get_filename_list("upscale_models"),), + }} + + RETURN_TYPES = () + FUNCTION = "run" + OUTPUT_NODE = True + CATEGORY = "FunPack/video" + + def run(self, filename, subfolder, type, upscale_model): + import comfy.model_management as mm + import folder_paths + from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel, UpscaleModelLoader + + src = resolve(filename, subfolder, type) + model = UpscaleModelLoader.execute(upscale_model).result[0] + out_dir = os.path.join(folder_paths.get_output_directory(), SUBFOLDER) + os.makedirs(out_dir, exist_ok=True) + stem = os.path.splitext(os.path.basename(filename))[0] + tag = os.path.splitext(os.path.basename(upscale_model))[0] + n, name = 0, f"{stem}_{tag}.mp4" + while os.path.exists(os.path.join(out_dir, name)): + n += 1 + name = f"{stem}_{tag}_{n}.mp4" + frames = upscale_file( + src, os.path.join(out_dir, name), + lambda x: ImageUpscaleWithModel.execute(model, x).result[0], + mm.throw_exception_if_processing_interrupted) + print(f"[FunPack] upscaled {filename} with {upscale_model}: {frames} frames -> {SUBFOLDER}/{name}") + return {"ui": {"videos": [{"filename": name, "subfolder": SUBFOLDER, "type": "output", + "format": "video/h264-mp4"}]}}