feat: upscale finished renders with an upscale model (never / always / Upscale button)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
c072a088ec
commit
81d8a9b094
@@ -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.
|
||||
|
||||
@@ -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__ = [
|
||||
|
||||
@@ -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")),
|
||||
|
||||
@@ -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: '<svg viewBox="0 0 16 16" width="13" height="13" fill="none" stroke="#fff" stroke-width="1.5" stroke-linecap="round"><path d="M2 4.5h12M2 8h12M2 11.5h12"/><circle cx="6" cy="4.5" r="1.7" fill="#fff" stroke="none"/><circle cx="11" cy="8" r="1.7" fill="#fff" stroke="none"/><circle cx="5" cy="11.5" r="1.7" fill="#fff" stroke="none"/></svg>',
|
||||
mount,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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"}]}}
|
||||
Reference in New Issue
Block a user