From 97b85a3f098d29fbe455aa9ae6f9956240b9b33b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 12:59:57 +0000 Subject: [PATCH] Interpolate: RTX Frame Generation as its own module and add-on, button and modes like Upscale Shared video stream and NVIDIA helpers move to core (video_pipe, nvidia); the add-on needs nvidia-vfx 0.2+, moves an older one up and never moves a newer one back; Upscale and Interpolate cannot feed each other in Always mode. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr --- CHANGELOG.md | 1 + README.md | 1 + app/ui/features/interpolate/index.js | 104 ++++++++++++ app/ui/features/interpolate/job.js | 13 ++ app/ui/features/interpolate/job.test.js | 33 ++++ app/ui/features/upscale/index.js | 10 +- app/ui/features/upscale/job.js | 29 +--- app/ui/modules.js | 1 + app/ui/shell/api.js | 5 + app/ui/shell/renders.js | 36 ++++ core/addons.py | 27 ++- core/nvidia.py | 28 ++++ core/tests/test_addons.py | 45 ++++- core/tests/test_shipped_nodes.py | 1 + core/update.py | 21 ++- core/video_pipe.py | 154 ++++++++++++++++++ modules/_core.py | 6 +- modules/output/interpolate/__init__.py | 34 ++++ modules/output/interpolate/framegen.py | 76 +++++++++ modules/output/interpolate/nodes.py | 55 +++++++ .../interpolate/tests/test_interpolate.py | 77 +++++++++ modules/output/render/files.py | 34 +--- modules/output/render/rtx.py | 34 +--- modules/output/render/tests/test_upscale.py | 1 + modules/output/render/upscale.py | 102 +----------- 25 files changed, 728 insertions(+), 200 deletions(-) create mode 100644 app/ui/features/interpolate/index.js create mode 100644 app/ui/features/interpolate/job.js create mode 100644 app/ui/features/interpolate/job.test.js create mode 100644 app/ui/shell/renders.js create mode 100644 core/nvidia.py create mode 100644 core/video_pipe.py create mode 100644 modules/output/interpolate/__init__.py create mode 100644 modules/output/interpolate/framegen.py create mode 100644 modules/output/interpolate/nodes.py create mode 100644 modules/output/interpolate/tests/test_interpolate.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 70333c62..2a330779 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- **Interpolate** (module of its own, add-on **NVIDIA RTX Frame Generation**): adds in-between frames to a finished render with NVIDIA's RTX frame generation (2 to 4 times the frames; the length and the sound stay, 24 fps becomes 48 fps). Settings > Editor > Interpolate: Never, a button beside the rating, or always when a render finishes. Quality Fast, Balanced or Best, and 10-bit processing. Works beside Upscale in either order and never runs twice on the same file. Needs an RTX card and nvidia-vfx 0.2 or newer: switching the add-on on moves an older nvidia-vfx up (nothing else is touched), and Super Resolution never moves a newer one back. ComfyUI's own RIFE/FILM interpolation nodes are a different thing and are not used. Untested on a GPU. - **Project settings > Resolution is the** (H3, LTX image to video): what Width and Height mean once the scene picture sets the shape. Short side (the smaller number is the clip's shorter side), Long side (the larger is its longer side) or Pixel count (the old rule: about Width x Height pixels, which turned 1024x1024 into 1312x768 for a wide picture). New projects start on Short side; projects saved earlier keep Pixel count, so nothing changes size by itself. The H3 preset no longer needs the KJNodes ImageTransformKJ node for this (saved H3 pipelines are switched over on load). - **VAE Loader: Decode-only VAE (Optional)**: a second VAE used only for the final picture decode while the first still encodes (references, start pictures). For decoders without an encoder: MiniMax H3's Sol-H3-VAE (`sol_h3_vae.safetensors`, whole clip at once, about 2.5x faster decode, needs a ComfyUI newer than 0.43.1) and the X2 Detail VAE. - **H3 loader: MXFP8** (Blackwell GPUs): stores the attention and MLP weights of the middle blocks as MXFP8 (first 2 and last 3 blocks stay bf16) for about twice the matmul rate. Needs torch 2.10+ with CUDA 13; skipped with a note elsewhere. Untested on a GPU; listen to the sound first. diff --git a/README.md b/README.md index a34e24e4..721c3955 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,7 @@ I express my deepest gratitude to: - [Inter](https://github.com/rsms/inter) and IBM Plex Mono — the fonts of the app (SIL Open Font License; copies in `app/assets/fonts`); - [musubi-tuner](https://github.com/kohya-ss/musubi-tuner) by kohya-ss (Apache-2.0) — the LoRA training add-on's loop (which layers get a LoRA, per-block checkpointing, noise sampling and loss weighting, optimizer and learning rate) and its block swap (a number of transformer blocks kept in system memory) follow its MiniMax H3 training recipe; - PTD few-step support for MiniMax H3 is ported from [ComfyUI-PTD-Compact](https://github.com/Nynxz/ComfyUI-PTD-Compact) by Nynxz (Apache-2.0), which builds on [ComfyUI-PTD](https://github.com/vita-epfl/PTD-inference) by EPFL VITA Lab (Apache-2.0; paper "Parametric Trajectory Distillation for Few-Step Video Generation", arXiv 2610.11498) and on ComfyUI-MiniMaxH3-PTD by taobao-content-ai (Apache-2.0); licence notice in `modules/sampling/ptd`; +- NVIDIA and the NVIDIA Video Effects SDK (`nvidia-vfx`, NVIDIA's own licence) — RTX Super Resolution in Upscale and RTX Frame Generation in Interpolate run on it; FunUI installs it only when you switch the add-on on; - NVIDIA and the authors of Sol-H3 (arXiv 2609.35110, NVlabs/Sana, Apache-2.0) — the MXFP8 loader option (which blocks stay in bf16) and the Two stages setting (the cross-resolution schedule) follow its recipe; Without all of you, this project would've been impossible. diff --git a/app/ui/features/interpolate/index.js b/app/ui/features/interpolate/index.js new file mode 100644 index 00000000..ad99482d --- /dev/null +++ b/app/ui/features/interpolate/index.js @@ -0,0 +1,104 @@ +// Interpolate a finished render with NVIDIA RTX Frame Generation: more frames, same length and sound; the result replaces the render. Never / button beside the rating / always. +import { composer as c } from "../../composer/composer.js"; +import { inPlace } from "../../shell/place.js"; +import { genUnitId, isGenerative, unitRoot } from "../../shell/scenes.js"; +import { derived, keyOf, swap, wentThrough } from "../../shell/renders.js"; +import { DEFAULTS, RTX, interpolate, modeLabel } from "./job.js"; + +export default { + id: "interpolate", + mount: "timeline.actions", + needs: ["project", "api", "selection"], + modules: ["interpolate"], + setup({ host, app }) { + const p = app.project, sel = app.selection; + const put = inPlace(host), busy = new Set(); + const render = (id) => ((p.project || {}).scene_renders || {})[id]; + const target = () => { // the picked clip's generation unit root, when it has a render + const open = p.project, sc = open && open.scenes.find((s) => s.id === sel.focus); + const root = sc && isGenerative(sc) ? unitRoot(open, genUnitId(sc)) : null; + const r = root && render(root.id); + return r && r.media ? r.media : null; + }; + const opts = () => ({ ...DEFAULTS, ...(p.pref("interpolateOpts", {}) || {}) }); + // Pressed with nothing to interpolate with: say what would fix it, and take the person there. + async function suggest() { + let all = []; try { all = await app.api.addons(); } catch { /* no add-on list */ } + const fg = all.find((a) => a.id === "nvidia_framegen"); + if (!fg) return void c.toast.warn({ text: `${RTX} needs an NVIDIA RTX card on Windows or Linux, so it is not available here.` }); + const pick = await c.modal.choice({ title: "Nothing to interpolate with yet", items: [fg.state === "pending" + ? { id: "modules", label: `${RTX} is switched on`, hint: "Restart ComfyUI to install it." } + : { id: "modules", label: `Install ${RTX}`, hint: "Switch it on in Settings ▸ Modules. It needs nvidia-vfx 0.2 or newer and moves an older one up." }] }).result; + if (pick) app.openSettings(pick); + } + + async function run(media, pressed = false) { + if (!media || wentThrough(media, "interpolate")) return; // never interpolate twice + const key = keyOf(media), pid = p.project.id; + if (busy.has(key)) return; + busy.add(key); draw(); + try { + // Only counts when this machine has it on right now: a project saved elsewhere may still say "always" + if (!(await app.api.frameGeneration().then((r) => Boolean(r.rtx)).catch(() => false))) + return void (pressed ? suggest() : c.toast.warn({ text: `${RTX} is not available on this machine (or is switched off in Settings ▸ Modules).` })); + c.toast.info({ text: `Interpolating with ${RTX}…` }); + const out = await interpolate(app.api, media, opts()); + derived(media, out, "interpolate"); + const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0); + if (!swapped) c.toast.warn({ text: "Interpolated, but that render was replaced meanwhile, so nothing was swapped." }); + else c.toast.good({ text: `Interpolated: the render now plays at ${opts().multiplier} times the frames.` }); + } catch (err) { c.toast.warn({ text: `Interpolate failed: ${err.message}` }); } + finally { busy.delete(key); draw(); } + } + + function draw() { + const media = p.pref("interpolateMode", "never") === "button" ? target() : null; + if (!media) return put(); + const done = wentThrough(media, "interpolate"), working = busy.has(keyOf(media)); + put(c.button.sm({ label: working ? "Interpolating…" : done ? "Interpolated" : "Interpolate", tone: "ghost", disabled: working || done, + title: "Add in-between frames to this video with NVIDIA RTX Frame Generation; it replaces the render", onClick: () => run(media, true) }).node); + } + + // "Always": a render that appears after the project was opened is interpolated. What was there on open is left alone. + let seen = null, obj = null; + const watch = () => { + const open = p.project; + if (!open) { seen = null; obj = null; return; } + const now = new Map(Object.entries(open.scene_renders || {}).filter(([, r]) => r && r.media).map(([id, r]) => [id, keyOf(r.media)])); + if (open !== obj || !seen) { obj = open; seen = now; return; } // another project, or undo/redo put an older copy back: not a new render + if (p.pref("interpolateMode", "never") === "always") { + const fresh = new Map([...now].filter(([id, k]) => seen.get(id) !== k)); + const pending = (m) => p.pref("upscaleMode", "never") === "always" && !wentThrough(m, "upscale"); // Upscale goes first; its result comes back as a new render + for (const media of new Map([...fresh.keys()].map((id) => [fresh.get(id), open.scene_renders[id].media])).values()) if (!pending(media)) run(media); + } + seen = now; + }; + + app.editorSettings.push(() => { + const page = c.region.stack({ gap: "md" }); + const set = (k, v) => { p.setPref("interpolateOpts", { ...opts(), [k]: v }); draw(); }; + const paint = (r) => { + const o = opts(); + page.set([ + c.label.section({ text: "Interpolate finished renders" }), + c.hint.default({ text: "Adds in-between frames to a finished video and replaces the render with the result. The length and the sound stay, the frame rate goes up." }), + c.settingsRow.default({ label: "Interpolate", hint: "Never, with a button beside the rating, or on every new render.", control: c.select.md({ label: "Interpolate", value: p.pref("interpolateMode", "never"), + options: [{ value: "never", label: "Never" }, { value: "button", label: "When I press Interpolate" }, { value: "always", label: "Always, when a render finishes" }], onChange: (v) => { p.setPref("interpolateMode", v); draw(); } }) }), + ...(r && r.rtx ? [ + c.settingsRow.default({ label: "Frames", hint: "2 times turns 24 fps into 48 fps.", control: c.select.md({ label: "Frames", value: String(o.multiplier), onChange: (v) => set("multiplier", Number(v)), + options: Array.from({ length: r.rtx.max_multiplier - 1 }, (_, i) => ({ value: String(i + 2), label: `${i + 2} times` })) }) }), + c.settingsRow.default({ label: "Quality", hint: "Fast is for a quick look, Best for the final video.", control: c.select.md({ label: "Quality", value: o.mode, onChange: (v) => set("mode", v), + options: r.rtx.modes.map((m) => ({ value: m, label: modeLabel(m) })) }) }), + c.toggle.default({ label: "10-bit processing", hint: "More precise colour inside the interpolation. The saved video stays 8-bit.", checked: o.tenbit, onChange: (v) => set("tenbit", v) }), + ] : r === null ? [] : [c.banner.warn({ text: !r.rtx_offered ? `${RTX} needs an NVIDIA RTX card on Windows or Linux, so it is not available here.` : `${RTX} is not installed or is switched off: switch it on in Settings ▸ Modules and restart ComfyUI.` })]), + ]); + }; + paint(null); + app.api.frameGeneration().then(paint).catch(() => page.set([c.banner.warn({ text: "Could not read the frame generation setup." })])); + return page; + }); + + draw(); watch(); + return app.on(() => { watch(); draw(); }); + }, +}; diff --git a/app/ui/features/interpolate/job.js b/app/ui/features/interpolate/job.js new file mode 100644 index 00000000..5934c586 --- /dev/null +++ b/app/ui/features/interpolate/job.js @@ -0,0 +1,13 @@ +// One interpolation job from queue to result. Pure of the page; the waiting and the swap are shared with Upscale (shell/renders.js). +import { poll } from "../../shell/renders.js"; + +export const RTX = "NVIDIA RTX Frame Generation"; // the same string as modules/output/interpolate/framegen.py NAME +export const DEFAULTS = { mode: "MEDIUM", multiplier: 2, tenbit: false }; +const MODE = { LOW: "Fast", MEDIUM: "Balanced", HIGH: "Best" }; +/** "MEDIUM" -> "Balanced": what the SDK calls a generation mode, said to a person. */ +export const modeLabel = (m) => MODE[m] || m; + +/** -> the interpolated media, or throws a sentence. `sleep` is injectable so a test need not wait. */ +export async function interpolate(api, media, opts, sleep) { + return poll(api, await api.queueInterpolate(media, opts), sleep); +} diff --git a/app/ui/features/interpolate/job.test.js b/app/ui/features/interpolate/job.test.js new file mode 100644 index 00000000..eb123992 --- /dev/null +++ b/app/ui/features/interpolate/job.test.js @@ -0,0 +1,33 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { interpolate, modeLabel } from "./job.js"; +import { derived, swap, wentThrough } from "../../shell/renders.js"; + +const m = (f, sub = "") => ({ filename: f, subfolder: sub, type: "output" }); + +test("the options reach the node, the job is waited for, and an error is thrown in words", async () => { + const seen = [], seq = [null, { videos: [m("fi.mp4", "funui_interpolated")] }]; + const api = { queueInterpolate: async (...a) => (seen.push(a), "id"), upscaleResult: async () => seq.shift() }; + assert.equal((await interpolate(api, m("x.mp4"), { multiplier: 2 }, async () => {})).filename, "fi.mp4"); + assert.deepEqual(seen[0][1], { multiplier: 2 }); + const bad = { queueInterpolate: async () => "id", upscaleResult: async () => ({ error: "no libnvidia-ngx" }) }; + await assert.rejects(interpolate(bad, m("x.mp4"), {}, async () => {}), /libnvidia-ngx/); +}); + +test("modes read as words", () => { + assert.deepEqual(["LOW", "MEDIUM", "HIGH"].map(modeLabel), ["Fast", "Balanced", "Best"]); +}); + +test("a file that went through a step is never offered it again, whichever order the two steps ran in", () => { + const raw = m("a.mp4"), up = m("a_up.mp4", "funui_upscaled"), fi = m("a_x2.mp4", "funui_interpolated"), both = m("a_up_x2.mp4", "funui_interpolated"); + assert.ok(!wentThrough(raw, "upscale") && !wentThrough(raw, "interpolate")); + assert.ok(wentThrough(up, "upscale") && !wentThrough(up, "interpolate")); + derived(up, both, "interpolate"); // upscaled first, then interpolated + assert.ok(wentThrough(both, "upscale") && wentThrough(both, "interpolate")); // so the upscale's "always" cannot take it again: no loop + assert.ok(wentThrough(fi, "interpolate") && !wentThrough(fi, "upscale")); // interpolated first: still upscalable once +}); + +test("the interpolated file replaces the render everywhere it played", () => { + const p = { scene_renders: { a: { media: m("x.mp4") } }, scene_ghosts: [] }; + assert.equal(swap(p, m("x.mp4"), m("fi.mp4", "funui_interpolated")), 1); +}); diff --git a/app/ui/features/upscale/index.js b/app/ui/features/upscale/index.js index 2258ef57..09baef73 100644 --- a/app/ui/features/upscale/index.js +++ b/app/ui/features/upscale/index.js @@ -2,7 +2,8 @@ import { composer as c } from "../../composer/composer.js"; import { inPlace } from "../../shell/place.js"; import { genUnitId, isGenerative, unitRoot } from "../../shell/scenes.js"; -import { RTX, RTX_DEFAULTS, UPSCALED, keepsSize, keyOf, qualityLabel, swap, upscale } from "./job.js"; +import { derived, wentThrough } from "../../shell/renders.js"; +import { RTX, RTX_DEFAULTS, keepsSize, keyOf, qualityLabel, swap, upscale } from "./job.js"; export default { id: "upscale", @@ -36,7 +37,7 @@ export default { } async function run(media, pressed = false) { - if (!media || media.subfolder === UPSCALED) return; // never upscale twice + if (!media || wentThrough(media, "upscale")) return; // never upscale twice, even after an interpolation in between const key = keyOf(media), pid = p.project.id; if (busy.has(key)) return; busy.add(key); draw(); @@ -48,6 +49,7 @@ export default { if (!model) return void (pressed ? suggest() : c.toast.warn({ text: "Upscale: pick a model or NVIDIA RTX Super Resolution in Settings ▸ Editor." })); c.toast.info({ text: `Upscaling with ${model}…` }); const out = await upscale(app.api, media, model, undefined, rtxOn() ? rtxOpts() : undefined); + derived(media, out, "upscale"); const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0); if (!swapped) c.toast.warn({ text: "Upscaled, but that render was replaced meanwhile, so nothing was swapped." }); else c.toast.good({ text: "Upscaled: the render now plays the larger video." }); @@ -58,8 +60,8 @@ export default { function draw() { const media = p.pref("upscaleMode", "never") === "button" ? target() : null; if (!media) return put(); - const working = busy.has(keyOf(media)) || media.subfolder === UPSCALED; - put(c.button.sm({ label: busy.has(keyOf(media)) ? "Upscaling…" : media.subfolder === UPSCALED ? "Upscaled" : "Upscale", tone: "ghost", disabled: working, + const working = busy.has(keyOf(media)) || wentThrough(media, "upscale"); + put(c.button.sm({ label: busy.has(keyOf(media)) ? "Upscaling…" : wentThrough(media, "upscale") ? "Upscaled" : "Upscale", tone: "ghost", disabled: working, title: "Upscale this video with the engine picked in Settings ▸ Editor; it replaces the render", onClick: () => run(media, true) }).node); } diff --git a/app/ui/features/upscale/job.js b/app/ui/features/upscale/job.js index 32c8ba0e..d7060ed6 100644 --- a/app/ui/features/upscale/job.js +++ b/app/ui/features/upscale/job.js @@ -1,5 +1,6 @@ -// One upscale job from queue to result, and the swap of the new file under every render that plays the old one. Pure of the page. -export const keyOf = (m) => `${m.type || "output"}/${m.subfolder || ""}/${m.filename}`; +// One upscale job from queue to result. Pure of the page. +import { poll } from "../../shell/renders.js"; +export { keyOf, swap } from "../../shell/renders.js"; export const UPSCALED = "funui_upscaled"; export const RTX = "NVIDIA RTX Super Resolution"; // sent as the model name for that engine; the same string as modules/output/render/rtx.py NAME export const RTX_DEFAULTS = { quality: "ULTRA", resize: "scale", scale: 2, width: 1920, height: 1080, strength: 1, tenbit: false }; @@ -16,26 +17,6 @@ export const qualityLabel = (q) => { export const keepsSize = (q) => q.startsWith("DENOISE_") || q.startsWith("DEBLUR_"); /** -> the upscaled media, or throws a sentence. `sleep` is injectable so a test need not wait. */ -export async function upscale(api, media, model, sleep = (ms) => new Promise((r) => setTimeout(r, ms)), rtx) { - const id = await api.queueUpscale(media, model, rtx); - let lost = 0; - for (;;) { - await sleep(2000); - const r = await api.upscaleResult(id); - if (!r) { lost = 0; continue; } - if (r.retry) { if (++lost < 30) continue; throw new Error("lost contact with ComfyUI"); } - if (r.gone) { if (++lost < 3) continue; throw new Error("the job left ComfyUI's queue without finishing (interrupted or ComfyUI restarted)"); } - if (r.error) throw new Error(r.error); - return r.videos[0]; - } -} - -/** One file backs a whole chain run, so every scene render and ghost that plays `from` gets `to`. -> how many were swapped. */ -export function swap(p, from, to) { - const same = (m) => !!m && keyOf(m) === keyOf(from); - let n = 0; - for (const [id, r] of Object.entries(p.scene_renders || {})) if (r && same(r.media)) { p.scene_renders[id] = { ...r, media: to }; n += 1; } - for (const list of Object.values(p.scene_variants || {})) list.forEach((t) => { if (same(t.media)) { t.media = to; n += 1; } }); // a take stepped back to must be the upscaled file too - p.scene_ghosts = (p.scene_ghosts || []).map((g) => (same(g.media) ? (n += 1, { ...g, media: to }) : g)); - return n; +export async function upscale(api, media, model, sleep, rtx) { + return poll(api, await api.queueUpscale(media, model, rtx), sleep); } diff --git a/app/ui/modules.js b/app/ui/modules.js index 57492d44..7f602aec 100644 --- a/app/ui/modules.js +++ b/app/ui/modules.js @@ -18,6 +18,7 @@ export default [ "./features/takes/index.js", "./features/subjects/index.js", "./features/upscale/index.js", + "./features/interpolate/index.js", "./features/revolver/index.js", "./features/autocomplete/index.js", "./features/autosave/index.js", diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index fbb6f015..4e651f37 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -143,6 +143,11 @@ export const api = { } return body.prompt_id; }, + frameGeneration: () => call("GET", "/api/m/interpolate/frame_generation"), + /** Interpolation is an ordinary ComfyUI job too; read it back with upscaleResult. -> the prompt id. */ + queueInterpolate(media, opts) { + return this.queueGraph({ 1: { class_type: "FunUIInterpolateVideo", inputs: { filename: media.filename, subfolder: media.subfolder || "", type: media.type === "temp" ? "temp" : "output", frame_options: JSON.stringify(opts || {}) } } }); + }, /** Whether ComfyUI's queue holds a run of this project (queued from any tab). Unreadable: false, the queue decides. */ async projectQueued(projectId) { try { diff --git a/app/ui/shell/renders.js b/app/ui/shell/renders.js new file mode 100644 index 00000000..9c7c5136 --- /dev/null +++ b/app/ui/shell/renders.js @@ -0,0 +1,36 @@ +// What the steps that rewrite a finished render (upscale, interpolate) share: naming a file, waiting for its job, swapping it under every +// scene that plays it, and remembering which steps a file has been through so two "always" modes can never feed each other forever. Pure of the page. +export const keyOf = (m) => `${m.type || "output"}/${m.subfolder || ""}/${m.filename}`; + +/** -> the finished job's first video, or throws a sentence. `sleep` is injectable so a test need not wait. */ +export async function poll(api, id, sleep = (ms) => new Promise((r) => setTimeout(r, ms))) { + let lost = 0; + for (;;) { + await sleep(2000); + const r = await api.upscaleResult(id); + if (!r) { lost = 0; continue; } + if (r.retry) { if (++lost < 30) continue; throw new Error("lost contact with ComfyUI"); } + if (r.gone) { if (++lost < 3) continue; throw new Error("the job left ComfyUI's queue without finishing (interrupted or ComfyUI restarted)"); } + if (r.error) throw new Error(r.error); + return r.videos[0]; + } +} + +/** One file backs a whole chain run, so every scene render and ghost that plays `from` gets `to`. -> how many were swapped. */ +export function swap(p, from, to) { + const same = (m) => !!m && keyOf(m) === keyOf(from); + let n = 0; + for (const [id, r] of Object.entries(p.scene_renders || {})) if (r && same(r.media)) { p.scene_renders[id] = { ...r, media: to }; n += 1; } + for (const list of Object.values(p.scene_variants || {})) list.forEach((t) => { if (same(t.media)) { t.media = to; n += 1; } }); // a take stepped back to must be the new file too + p.scene_ghosts = (p.scene_ghosts || []).map((g) => (same(g.media) ? (n += 1, { ...g, media: to }) : g)); + return n; +} + +/** Where each step writes its files (modules/output/render/upscale.py and modules/output/interpolate/nodes.py say the same). */ +export const FOLDERS = { upscale: "funui_upscaled", interpolate: "funui_interpolated" }; +const steps = new Map(); // file key -> the steps it has been through (kept for this page's life) +const doneOf = (m) => new Set([...(steps.get(keyOf(m)) || []), ...Object.keys(FOLDERS).filter((k) => FOLDERS[k] === m.subfolder)]); +/** Whether `media` already went through `step`. */ +export const wentThrough = (media, step) => doneOf(media).has(step); +/** `to` came out of `from` by `step`: it has been through all of `from`'s steps and this one. */ +export const derived = (from, to, step) => { steps.set(keyOf(to), doneOf(from).add(step)); }; diff --git a/core/addons.py b/core/addons.py index d169e7f7..177c255f 100644 --- a/core/addons.py +++ b/core/addons.py @@ -46,6 +46,9 @@ ADDONS = ( {"id": "nvidia_vsr", "title": "NVIDIA RTX Super Resolution", "hint": "Native upscaling for images and videos. Installs NVIDIA's own package, under its licence.", "requires": ("nvidia-vfx>=0.2",), "modules": ("nvvfx",), "offered": _nvidia_gpu}, + {"id": "nvidia_framegen", "title": "NVIDIA RTX Frame Generation", + "hint": "Adds in-between frames to finished videos (x2 and more). Installs NVIDIA's own package, under its licence; needs nvidia-vfx 0.2 or newer, and replaces an older one.", + "requires": ("nvidia-vfx>=0.2",), "modules": ("nvvfx",), "floors": {"nvidia-vfx": "0.2"}, "offered": _nvidia_gpu}, {"id": "camera", "title": "Automatic camera and focus", "hint": "Lets Shot camera aim camera moves, cuts, transitions and focus at the right things in your prompts, using spaCy.", "requires": ("spacy>=3.8,<3.9", PRESHED_PIN, _SPACY_MODEL), "modules": ("spacy", "en_core_web_sm"), "offered": lambda: True}, @@ -56,8 +59,28 @@ ADDONS = ( _BY_ID = {a["id"]: a for a in ADDONS} +def _version_tuple(v) -> tuple: + out = [] + for part in str(v).split("."): + digits = "".join(ch for ch in part if ch.isdigit()) + if not digits: + break + out.append(int(digits)) + return tuple(out) + + +def _old(addon) -> bool: + """True when a package the add-on needs is installed below the version it needs ("floors"): present, but not usable.""" + from . import update + for name, floor in (addon.get("floors") or {}).items(): + have = update._installed_version(name) + if have is not None and _version_tuple(have) < _version_tuple(floor): + return True + return False + + def _present(addon) -> bool: - return all(importlib.util.find_spec(m) is not None for m in addon["modules"]) + return all(importlib.util.find_spec(m) is not None for m in addon["modules"]) and not _old(addon) def _read() -> dict | None: @@ -167,7 +190,7 @@ def install_wanted() -> None: from . import update for a in todo: log.info("add-ons", f"installing {a['title']}…") - result = update.install_lines(a["requires"]) + result = update.install_lines(a["requires"], upgrade=tuple(a.get("floors") or ())) importlib.invalidate_caches() if _present(a): continue diff --git a/core/nvidia.py b/core/nvidia.py new file mode 100644 index 00000000..d5cfd6e2 --- /dev/null +++ b/core/nvidia.py @@ -0,0 +1,28 @@ +"""Helpers shared by the NVIDIA RTX video effects (nvidia-vfx): the packed 10-bit layout and the driver-library hint.""" + +def pack10(x): + """[3,H,W] float 0-1 -> (H,W) uint32 in the SDK's RGB10A2 layout: R bits 0-9, G 10-19, B 20-29, A 30-31.""" + import torch + v = (x.clamp(0, 1) * 1023).round().to(torch.int64) + return (v[0] | (v[1] << 10) | (v[2] << 20) | (3 << 30)).to(torch.uint32) + + +def unpack10(p): + """(H,W) uint32 -> [3,H,W] float 0-1.""" + import torch + v = p.to(torch.int64) + return torch.stack([v & 1023, (v >> 10) & 1023, (v >> 20) & 1023]).float() / 1023.0 + + +def ngx_hint(): + """Why load() can fail on a rental: the SDK needs the driver's own libnvidia-ngx.so.1, which a container host may not mount.""" + import ctypes + import sys + if sys.platform != "linux": + return "" + try: + ctypes.CDLL("libnvidia-ngx.so.1") + return " The driver's NGX library is present, so check the driver is 570.190+ / 580.82+ / 590.44+ and the card is Turing or newer." + except OSError: + return (" This machine has no libnvidia-ngx.so.1 (NVIDIA's driver library that the RTX video effects need; some rented GPUs do not mount it, " + "like libnvidia-encode). Use an upscale model instead, or rent a host that exposes the full NVIDIA driver (graphics/video capabilities).") diff --git a/core/tests/test_addons.py b/core/tests/test_addons.py index b615dbaf..04909df0 100644 --- a/core/tests/test_addons.py +++ b/core/tests/test_addons.py @@ -43,7 +43,8 @@ def test_switching_on_waits_for_the_restart_and_off_keeps_the_packages(here): def test_an_add_on_this_machine_cannot_run_is_not_offered(here, monkeypatch): - monkeypatch.setitem(addons._BY_ID["nvidia_vsr"], "offered", lambda: False) + for gone in ("nvidia_vsr", "nvidia_framegen"): + monkeypatch.setitem(addons._BY_ID[gone], "offered", lambda: False) assert [a["id"] for a in addons.listing()] == ["camera", "lora_training"] assert addons.set_wanted("nvidia_vsr", True) is False assert addons.set_wanted("no-such-addon", True) is False @@ -52,7 +53,7 @@ def test_an_add_on_this_machine_cannot_run_is_not_offered(here, monkeypatch): def test_the_start_installs_only_what_is_wanted_and_missing(here, monkeypatch): ran = [] - def fake(lines): + def fake(lines, upgrade=()): ran.append(list(lines)) here.add("nvidia_vsr") return {"ok": True, "detail": ""} @@ -68,7 +69,7 @@ def test_the_start_installs_only_what_is_wanted_and_missing(here, monkeypatch): def test_a_failed_install_is_said_and_not_retried_until_switched_off_and_on(here, monkeypatch): ran = [] - monkeypatch.setattr(update, "install_lines", lambda lines: ran.append(1) or {"ok": False, "detail": "no wheel for this Python"}) + monkeypatch.setattr(update, "install_lines", lambda lines, upgrade=(): ran.append(1) or {"ok": False, "detail": "no wheel for this Python"}) addons.set_wanted("nvidia_vsr", True) addons.install_wanted() row = next(a for a in addons.listing() if a["id"] == "nvidia_vsr") @@ -107,7 +108,7 @@ def test_the_routes_list_and_switch(comfyui, here): ok = await (await client.post("/fun/api/addons", json={"id": "camera", "enabled": True})).json() return got, bad.status, ok got, bad, ok = asyncio.run(go()) - assert {a["id"] for a in got["addons"]} == {"nvidia_vsr", "camera", "lora_training"} + assert {a["id"] for a in got["addons"]} == {"nvidia_vsr", "nvidia_framegen", "camera", "lora_training"} assert bad == 400 assert next(a["state"] for a in ok["addons"] if a["id"] == "camera") == "pending" assert json.loads(config.ADDONS_FILE.read_text())["wanted"] == ["camera"] @@ -139,3 +140,39 @@ def test_preshed_3_0_14_is_flagged_and_older_is_not(monkeypatch): for have, bad in (("3.0.14", True), ("3.0.13", False), ("3.1.0", True)): monkeypatch.setattr(metadata, "version", lambda _n, have=have: have) assert addons.preshed_broken() is bad + + +def test_frame_generation_is_a_separate_switch_on_the_same_package(): + vsr, fg = addons._BY_ID["nvidia_vsr"], addons._BY_ID["nvidia_framegen"] + assert vsr["requires"] == fg["requires"] == ("nvidia-vfx>=0.2",) # neither add-on can ask for the 0.1 that has no frame generation + assert fg["floors"] == {"nvidia-vfx": "0.2"} and "floors" not in vsr + + +def test_an_old_nvidia_vfx_is_present_for_super_resolution_but_not_for_frame_generation(monkeypatch): + monkeypatch.setattr(addons.importlib.util, "find_spec", lambda m: object()) + for have, ok in (("0.1.0.1", False), ("0.2.0.0", True), ("0.10.0", True)): + monkeypatch.setattr(update, "_installed_version", lambda n, have=have: have) + assert addons._present(addons._BY_ID["nvidia_framegen"]) is ok + assert addons._present(addons._BY_ID["nvidia_vsr"]) is True + + +def _pip(monkeypatch, installed, freeze): + seen = {} + monkeypatch.setattr(update, "_installed_version", lambda n: installed.get(n)) + monkeypatch.setattr(update, "_pip_freeze", lambda: freeze) + monkeypatch.setattr(update, "_install_each", lambda status, path, before, note, timeout, broken: + seen.update(lines=status["install"], pins=open(path).read()) or {"ok": True}) + return seen + + +def test_frame_generation_moves_an_old_nvidia_vfx_up_and_nothing_else(monkeypatch): + seen = _pip(monkeypatch, {"nvidia-vfx": "0.1.0.1"}, {"torch": "2.10.0", "nvidia-vfx": "0.1.0.1"}) + update.install_lines(["nvidia-vfx>=0.2"], upgrade=("nvidia-vfx",)) + assert seen["lines"] == ["nvidia-vfx>=0.2"] + assert "torch==2.10.0" in seen["pins"] and "nvidia-vfx" not in seen["pins"] + + +def test_super_resolution_never_moves_a_newer_nvidia_vfx_back(monkeypatch): + seen = _pip(monkeypatch, {"nvidia-vfx": "0.2.0.0"}, {"nvidia-vfx": "0.2.0.0"}) + assert update.install_lines(["nvidia-vfx>=0.2"])["ran"] is False and not seen # present: nothing to do, in either order + assert update.install_lines(["nvidia-vfx>=0.2"], upgrade=("nvidia-vfx",))["ran"] is False # already at the floor diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index b5830d7a..965f1be3 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -71,6 +71,7 @@ FROZEN = { "source_id", "phase_scale", "id_strength", "arcface_mode"], "FunUIDecode": ["samples", "vae", "model", "audio_vae", "tile_size"], "FunUIUpscaleVideo": ["filename", "subfolder", "type", "upscale_model", "rtx_options"], + "FunUIInterpolateVideo": ["filename", "subfolder", "type", "frame_options"], "FunUITrainLora": [ "model", "clip", "vae", "audio_vae", "set_name", "name", "use", "steps", "learning_rate", "rank", "algorithm", "optimizer", "loss", "batch_size", "grad_accumulation", "resolution", "aspect_buckets", "frames", "fps", "start", diff --git a/core/update.py b/core/update.py index a8cfe1d0..10681fec 100644 --- a/core/update.py +++ b/core/update.py @@ -449,17 +449,30 @@ def _constrained(before, run): os.unlink(constraints.name) -def install_lines(lines, timeout: int = 900) -> dict: +def _below_floor(name, spec) -> bool: + """True when `name` is installed but older than the ">=X" in `spec`.""" + have = _installed_version(name) + m = re.match(r"^>=\s*([0-9][0-9.]*)", spec or "") + if have is None or not m: + return False + key = lambda v: tuple(int(x) for x in re.findall(r"\d+", v)) + return key(have) < key(m.group(1)) + + +def install_lines(lines, timeout: int = 900, upgrade=()) -> dict: """Install the absent ones of these requirement lines, with the same guard as install_requirements: - nothing already present moves. Never raises.""" - todo = [l for l in lines if (p := parse_requirement(l)) and _installed_version(p[0]) is None] + nothing already present moves. Never raises. A name in `upgrade` that is installed below the line's + ">=" floor is moved up to it (and only that package: every other one stays pinned).""" + todo = [l for l in lines if (p := parse_requirement(l)) and (_installed_version(p[0]) is None or (p[0] in upgrade and _below_floor(*p)))] if not todo: return {"ran": False, "ok": True, "changed": [], "detail": "Nothing to install."} before = _pip_freeze() if not before: return {"ran": False, "ok": False, "detail": "Could not read the installed packages, so nothing was installed."} status = {"install": todo, "missing": [parse_requirement(l)[0] for l in todo]} - return _constrained(before, lambda path: _install_each(status, path, before, "", timeout, False)) + moving = {parse_requirement(l)[0].lower().replace("_", "-") for l in todo} + held = {n: v for n, v in before.items() if n.lower().replace("_", "-") not in moving} # the package being moved up is the one pin lifted + return _constrained(held, lambda path: _install_each(status, path, before, "", timeout, False)) def _install_each(status, constraints_path, before, floor_note, timeout, model_broken): diff --git a/core/video_pipe.py b/core/video_pipe.py new file mode 100644 index 00000000..6cf1db36 --- /dev/null +++ b/core/video_pipe.py @@ -0,0 +1,154 @@ +"""Stream a video file through a per-chunk function into a new file, keeping its sound. + +Frames go through ffmpeg in small chunks (never the whole clip in RAM). Shared by the post-render +steps that rewrite a finished video (upscale, interpolate), which cannot import each other. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess + +CHUNK = 8 # frames per call, at most +CHUNK_PIXELS = 2_000_000 # ...and at most this many input pixels per call: a big upscale of 1080p is gigabytes per frame + +FFMPEG_MISSING = "ffmpeg was not found on PATH: install it to preview, export or render clips." + + +class ClipError(Exception): + """A clip cannot be used, with the sentence the person should read.""" + + +def ffmpeg() -> str: + path = shutil.which("ffmpeg") + if not path: + raise ClipError(FFMPEG_MISSING) + return path + + +def comfy_dir(kind: str) -> str | None: + """ComfyUI's output or temp folder, or None when ComfyUI is not here to ask.""" + try: + import folder_paths + return folder_paths.get_output_directory() if kind == "output" else folder_paths.get_temp_directory() + except Exception: # noqa: BLE001 -- tests, headless + return None + + +def comfy_path(filename, subfolder="", kind="output", dir_of=None) -> str | None: + """The file ComfyUI calls (filename, subfolder, kind), or None if the name does not + resolve to a place under that folder.""" + if not isinstance(filename, str) or not filename: + return None + base = (dir_of or comfy_dir)("output" if kind != "temp" else "temp") + if not base: + return None + base = os.path.realpath(base) + full = os.path.realpath(os.path.join(base, subfolder if isinstance(subfolder, str) else "", + os.path.basename(filename))) + return full if os.path.commonpath((full, base)) == base else None + + +def probe(path): + """-> (width, height, fps string, has sound).""" + ffprobe = shutil.which("ffprobe") + if not ffprobe: + raise ClipError("ffprobe was not found on PATH: install ffmpeg to upscale videos.") + out = subprocess.run([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 ClipError(f"no video stream in {os.path.basename(path)}") + return (int(video["width"]), int(video["height"]), video.get("r_frame_rate") or "25", + any(s.get("codec_type") == "audio" for s in streams)) + + +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 scaled_fps(fps, k): + """"24/1" -> "48/1", "30000/1001" -> "60000/1001": the frame rate after every frame became k.""" + num, _, den = str(fps).partition("/") + return f"{int(num) * k}/{den or 1}" + + +def stream_file(src, dst, upscale, interrupted=lambda: None, fps_factor=1): + """Stream `src` through `upscale` ([B,H,W,3] float 0-1 -> same, larger, or more frames when `fps_factor` > 1) + into `dst`, played at `fps_factor` times the source's frame rate so the length and the sound stay in step. + Returns the frames written. `interrupted()` raises to stop between chunks.""" + import numpy as np + import torch + w, h, fps, has_sound = probe(src) + frame = w * h * 3 + chunk = max(1, min(CHUNK, CHUNK_PIXELS // max(1, w * h))) + exe = ffmpeg() + dec = subprocess.Popen([exe, "-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 = [exe, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", + "-s", f"{ow}x{oh}", "-r", scaled_fps(fps, fps_factor), "-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 += y.shape[0] + if enc is None: + raise ClipError(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 ClipError(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 p is not None: + p.wait() + for pipe in (p.stdin, p.stdout, p.stderr): + try: + if pipe: + pipe.close() + except OSError: # a killed encoder's unflushed pipe + pass + if os.path.exists(dst): + os.remove(dst) + raise + finally: + if dec.poll() is None: + dec.kill() + + +def free_name(out_dir, filename, model): + stem = os.path.splitext(os.path.basename(filename))[0] + tag = os.path.splitext(os.path.basename(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" + return name diff --git a/modules/_core.py b/modules/_core.py index 597fe5f1..286e2ec0 100644 --- a/modules/_core.py +++ b/modules/_core.py @@ -11,8 +11,8 @@ leading underscore keeps it out of the module scan. """ try: # root on sys.path: tests, dev server - from core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles + from core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia except ImportError: # inside ComfyUI: a subpackage of the pack - from ..core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles + from ..core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia -__all__ = ["addons", "chain", "config", "contract", "control", "dit_hooks", "graph", "input_steer", "log", "media", "patching", "projects", "rated_dial", "registry", "relations", "run", "schema", "shortcuts", "streams", "sysinfo", "temp_files", "traits", "vae_tiles"] +__all__ = ["addons", "chain", "config", "contract", "control", "dit_hooks", "graph", "input_steer", "log", "media", "patching", "projects", "rated_dial", "registry", "relations", "run", "schema", "shortcuts", "streams", "sysinfo", "temp_files", "traits", "vae_tiles", "video_pipe", "nvidia"] diff --git a/modules/output/interpolate/__init__.py b/modules/output/interpolate/__init__.py new file mode 100644 index 00000000..398b8bda --- /dev/null +++ b/modules/output/interpolate/__init__.py @@ -0,0 +1,34 @@ +"""Interpolate: more frames for a finished video, with NVIDIA RTX Frame Generation (the nvidia-vfx add-on). + +The length and the sound stay, the frame rate goes up (x2 turns 24 fps into 48 fps). It runs +beside ComfyUI on the file a render wrote, like Upscale, and is a module of its own. + +Not yet run on the real SDK (no RTX card in the test machine); the streaming around it is tested against real ffmpeg. +""" + +from . import framegen +from . import nodes + +try: + from aiohttp import web +except Exception: # noqa: BLE001 -- tests without aiohttp + web = None + +ID = "interpolate" +TITLE = "Interpolate" +STAGE = "post" +CATEGORY = "post" +STATUS = "experimental" +NODES = [nodes.FunUIInterpolateVideo] if nodes.io is not None else [] + + +def routes(table, base, web): + @table.get(base + "/frame_generation") + async def _frame_generation(_req): + from ..._core import addons + return web.json_response({ + "rtx": {"name": framegen.NAME, "modes": list(framegen.MODES), "max_multiplier": framegen.MAX_MULTIPLIER} if addons.enabled("nvidia_framegen") else None, + "rtx_offered": any(a["id"] == "nvidia_framegen" for a in addons.listing())}) # offered = this machine could switch it on in Settings ▸ Modules + + +PROVIDES = {"routes": routes} diff --git a/modules/output/interpolate/framegen.py b/modules/output/interpolate/framegen.py new file mode 100644 index 00000000..f1f56b7f --- /dev/null +++ b/modules/output/interpolate/framegen.py @@ -0,0 +1,76 @@ +"""NVIDIA RTX Video Frame Generation as an interpolation engine, through the nvidia-vfx add-on (the one RTX Super Resolution uses). + +`make(options)` returns a function [B,H,W,3] float 0-1 -> the same frames with the in-between ones added. +It remembers the last frame it saw, so chunks of one video are fed in order and the seams between chunks get their frames too. +""" + +import json + +NAME = "NVIDIA RTX Frame Generation" # shown as the engine's name (the UI holds the same string) +MODES = ("LOW", "MEDIUM", "HIGH") +DEFAULTS = {"mode": "MEDIUM", "multiplier": 2, "tenbit": False} +MAX_MULTIPLIER = 4 + + +def options(raw) -> dict: + """The user's choices, checked. A value out of range is refused by name, never quietly swapped.""" + try: + given = json.loads(raw) if isinstance(raw, str) and raw.strip() else (raw or {}) + except ValueError as exc: + raise ValueError(f"Frame generation options are not valid JSON: {exc}") from exc + if not isinstance(given, dict): + raise ValueError("Frame generation options must be an object.") + o = {**DEFAULTS, **given} + if o["mode"] not in MODES: + raise ValueError(f"Frame generation mode {o['mode']!r} is not one of: {', '.join(MODES)}") + if isinstance(o["multiplier"], bool) or not isinstance(o["multiplier"], int) or not 2 <= o["multiplier"] <= MAX_MULTIPLIER: + raise ValueError(f"Frame generation multiplier must be a whole number from 2 to {MAX_MULTIPLIER}.") + if not isinstance(o["tenbit"], bool): + raise ValueError("Frame generation tenbit must be true or false.") + return o + + +def make(o): + """-> the interpolation function for these (already checked) options. Raises a sentence when the add-on is missing.""" + from ..._core import addons + from ..._core import nvidia + ngx_hint, pack10, unpack10 = nvidia.ngx_hint, nvidia.pack10, nvidia.unpack10 + if not addons.enabled("nvidia_framegen"): # installed but switched off counts too: off means off + raise RuntimeError(f"{NAME} is not installed or is switched off: switch it on in Settings ▸ Modules and restart ComfyUI.") + try: + from nvvfx import VideoFrameGeneration + except ImportError as exc: + raise RuntimeError(f"{NAME} is not installed: switch it on in Settings ▸ Modules and restart ComfyUI.") from exc + if not hasattr(VideoFrameGeneration, "Mode"): + raise RuntimeError("This nvidia-vfx has no frame generation: update it (needs 0.2 or newer).") + import torch + mult = o["multiplier"] + effects, last = {}, [None] # one effect per frame size; the previous frame, on the GPU + + def effect(h, w): + if (h, w) not in effects: + kw = {"image_encoding": VideoFrameGeneration.ImageEncoding.RGB10A2} if o["tenbit"] else {} + e = VideoFrameGeneration(getattr(VideoFrameGeneration.Mode, o["mode"]), **kw) + e.input_width, e.input_height, e.frame_multiplier = w, h, mult + try: + e.load() + except Exception as exc: # noqa: BLE001 -- NvVFXError: say what to do, not just the SDK code + raise RuntimeError(f"{NAME} could not start ({exc}).{ngx_hint()}") from exc + effects[(h, w)] = e + return effects[(h, w)] + + def run(x): + out = [] + for frame in x: + h, w = frame.shape[0], frame.shape[1] + f = frame.to("cuda", torch.float32).permute(2, 0, 1).contiguous() + cur = pack10(f) if o["tenbit"] else f + if last[0] is not None: + for i in range(1, mult): + res = effect(h, w).run(last[0], cur, frame_index=i) + y = torch.from_dlpack(res.image).clone() # the effect reuses its buffer: copy before the next call + out.append((unpack10(y) if o["tenbit"] else y).permute(1, 2, 0)) + out.append(f.permute(1, 2, 0)) + last[0] = cur + return torch.stack(out) + return run diff --git a/modules/output/interpolate/nodes.py b/modules/output/interpolate/nodes.py new file mode 100644 index 00000000..28664352 --- /dev/null +++ b/modules/output/interpolate/nodes.py @@ -0,0 +1,55 @@ +"""Post-render interpolation: a finished video gets in-between frames (NVIDIA RTX Frame Generation), played at the multiplied frame rate. + +Same stream as the upscale (chunks through ffmpeg, the source's sound kept); only the engine and the frame rate differ. +""" + +import os + +from . import framegen +from ..._core import video_pipe + +SUBFOLDER = "funui_interpolated" + +try: + from comfy_api.latest import io +except Exception: # noqa: BLE001 -- tests without ComfyUI + io = None + +if io is not None: + class FunUIInterpolateVideo(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="FunUIInterpolateVideo", + display_name="FunUI Interpolate Video", + category="FunUI/Output", + description="Add in-between frames to a finished video file with NVIDIA RTX Frame Generation; the length and sound stay, the frame rate goes up.", + inputs=[ + io.String.Input("filename", default=""), + io.String.Input("subfolder", default=""), + io.Combo.Input("type", options=["output", "temp"], default="output"), + io.String.Input("frame_options", default="{}", optional=True, + tooltip="NVIDIA RTX Frame Generation options as JSON: mode (LOW, MEDIUM, HIGH), multiplier (2 to 4), tenbit."), + ], + outputs=[], + is_output_node=True, + ) + + @classmethod + def execute(cls, filename, subfolder, type, frame_options="{}") -> io.NodeOutput: + import comfy.model_management as mm + import folder_paths + + src = video_pipe.comfy_path(filename, subfolder, type) + if not src or not os.path.isfile(src): + raise RuntimeError(f"{filename} is not in the {type} folder (anymore).") + opts = framegen.options(frame_options) + fn = framegen.make(opts) + out_dir = os.path.join(folder_paths.get_output_directory(), SUBFOLDER) + os.makedirs(out_dir, exist_ok=True) + name = video_pipe.free_name(out_dir, filename, f"x{opts['multiplier']}") + frames = video_pipe.stream_file(src, os.path.join(out_dir, name), fn, mm.throw_exception_if_processing_interrupted, opts["multiplier"]) + from ..._core import log + log.info("FunUI Interpolate", f"{filename} x{opts['multiplier']} ({opts['mode']}): {frames} frames -> {SUBFOLDER}/{name}") + return io.NodeOutput(ui={"videos": [{"filename": name, "subfolder": SUBFOLDER, + "type": "output", "format": "video/h264-mp4"}]}) diff --git a/modules/output/interpolate/tests/test_interpolate.py b/modules/output/interpolate/tests/test_interpolate.py new file mode 100644 index 00000000..ab8b6202 --- /dev/null +++ b/modules/output/interpolate/tests/test_interpolate.py @@ -0,0 +1,77 @@ +import pytest + +from modules.output.interpolate import framegen +from modules.output.render import upscale +from modules.output.render.tests.clips import comfy, make_clip, needs_ffmpeg # noqa: F401 + +torch = pytest.importorskip("torch") + + +def test_scaled_fps_keeps_the_fraction(): + assert upscale.scaled_fps("24/1", 2) == "48/1" + assert upscale.scaled_fps("30000/1001", 2) == "60000/1001" + assert upscale.scaled_fps("25", 3) == "75/1" + + +@needs_ffmpeg +def test_interpolation_stream_adds_frames_and_keeps_length_and_sound(comfy): + src, dst = str(comfy.output / "a.mp4"), str(comfy.output / "fi.mp4") + make_clip(src, seconds=1.0, size="64x48", rate=10, audio=True) + last = [None] + + def between(x): # stands in for the effect: an average frame, carried across chunks + out = [] + for f in x: + if last[0] is not None: + out.append((last[0] + f) / 2) + out.append(f) + last[0] = f + return torch.stack(out) + + assert upscale.upscale_file(src, dst, between, fps_factor=2) == 19 # 10 frames -> 10 + 9 between, none lost at the chunk seam + w, h, fps, sound = upscale.probe(dst) + assert (w, h, fps, sound) == (64, 48, "20/1", True) + + +def test_frame_generation_options_are_checked_by_name(): + assert framegen.options("{}") == framegen.DEFAULTS + assert framegen.options({"multiplier": 4, "mode": "HIGH"})["multiplier"] == 4 + for bad in ({"mode": "ULTRA"}, {"multiplier": 1}, {"multiplier": 5}, {"multiplier": 2.5}, {"multiplier": True}, {"tenbit": 1}): + with pytest.raises(ValueError): + framegen.options(bad) + with pytest.raises(ValueError): + framegen.options("not json") + + +def test_frame_generation_without_the_add_on_says_how_to_get_it(monkeypatch): + monkeypatch.setitem(__import__("sys").modules, "nvvfx", None) + with pytest.raises(RuntimeError, match="Settings ▸ Modules"): + framegen.make(framegen.options("{}")) + + +def test_frame_generation_feeds_pairs_in_order_across_chunks(monkeypatch): + """A stand-in nvvfx (CPU) records the calls: every pair once, in order, indexes 1..N-1, the seam between chunks included.""" + import sys + import types + calls = [] + + class Effect: + class Mode: LOW = 0; MEDIUM = 1; HIGH = 2 + class ImageEncoding: RGB8 = 0; RGB10A2 = 1 + def __init__(self, mode, **kw): self.mode = mode + def load(self): pass + def run(self, a, b, frame_index=1): + calls.append((float(a.mean()), float(b.mean()), frame_index)) + return types.SimpleNamespace(image=(a + b) / 2) + + monkeypatch.setitem(sys.modules, "nvvfx", types.SimpleNamespace(VideoFrameGeneration=Effect)) + from core import addons + monkeypatch.setattr(addons, "enabled", lambda _id: True) + monkeypatch.setattr(torch.Tensor, "to", lambda self, *a, **k: self) # no GPU here: "cuda" is this device + monkeypatch.setattr(torch, "from_dlpack", lambda t: t) + fn = framegen.make(framegen.options({"multiplier": 3})) + frames = torch.stack([torch.full((2, 2, 3), v) for v in (0.0, 0.3, 0.6, 0.9)]) + a, b = fn(frames[:2]), fn(frames[2:]) + assert a.shape[0] == 4 and b.shape[0] == 6 # 1 + 3, then 3 + 3: the seam pair got its frames too + assert [round(c[1], 2) for c in calls[::2]] == [0.3, 0.6, 0.9] and [c[2] for c in calls] == [1, 2] * 3 + assert calls[2][0] == pytest.approx(0.3) and calls[2][1] == pytest.approx(0.6) # the seam pair was interpolated diff --git a/modules/output/render/files.py b/modules/output/render/files.py index ef79a69e..cef50744 100644 --- a/modules/output/render/files.py +++ b/modules/output/render/files.py @@ -20,47 +20,21 @@ import shutil import subprocess from pathlib import Path -from ..._core import media +from ..._core import media, video_pipe ISOBMFF = (".mp4", ".m4v", ".mov") # containers where moov/mdat order applies AUDIO_RATE = 48000 # every exported part is made at this rate, so parts can be joined -FFMPEG_MISSING = "ffmpeg was not found on PATH: install it to preview, export or render clips." - - -class ClipError(Exception): - """A clip cannot be used, with the sentence the person should read.""" - - -def ffmpeg() -> str: - path = shutil.which("ffmpeg") - if not path: - raise ClipError(FFMPEG_MISSING) - return path +ClipError, FFMPEG_MISSING, ffmpeg = video_pipe.ClipError, video_pipe.FFMPEG_MISSING, video_pipe.ffmpeg def comfy_dir(kind: str) -> str | None: - """ComfyUI's output or temp folder, or None when ComfyUI is not here to ask.""" - try: - import folder_paths - return folder_paths.get_output_directory() if kind == "output" else folder_paths.get_temp_directory() - except Exception: # noqa: BLE001 -- tests, headless - return None + return video_pipe.comfy_dir(kind) def comfy_path(filename, subfolder="", kind="output") -> str | None: - """The file ComfyUI calls (filename, subfolder, kind), or None if the name does not - resolve to a place under that folder.""" - if not isinstance(filename, str) or not filename: - return None - base = comfy_dir("output" if kind != "temp" else "temp") - if not base: - return None - base = os.path.realpath(base) - full = os.path.realpath(os.path.join(base, subfolder if isinstance(subfolder, str) else "", - os.path.basename(filename))) - return full if os.path.commonpath((full, base)) == base else None + return video_pipe.comfy_path(filename, subfolder, kind, comfy_dir) # through this module's comfy_dir, which tests replace def clip_path(clip: dict) -> str: diff --git a/modules/output/render/rtx.py b/modules/output/render/rtx.py index 63988544..9ecdd306 100644 --- a/modules/output/render/rtx.py +++ b/modules/output/render/rtx.py @@ -7,6 +7,10 @@ for ComfyUI (quality, scale or target size, strength, 8/10-bit), so a setup carr import json +from ..._core import nvidia + +ngx_hint, pack10, unpack10 = nvidia.ngx_hint, nvidia.pack10, nvidia.unpack10 # pack10/unpack10 stay importable from here + NAME = "NVIDIA RTX Super Resolution" # sent as the model name when this engine is picked (the UI holds the same string) _LEVELS = ("LOW", "MEDIUM", "HIGH", "ULTRA") QUALITIES = ("BICUBIC", *_LEVELS, @@ -46,34 +50,6 @@ def out_size(w, h, o): return snap(ow), snap(oh) -def pack10(x): - """[3,H,W] float 0-1 -> (H,W) uint32 in the SDK's RGB10A2 layout: R bits 0-9, G 10-19, B 20-29, A 30-31.""" - import torch - v = (x.clamp(0, 1) * 1023).round().to(torch.int64) - return (v[0] | (v[1] << 10) | (v[2] << 20) | (3 << 30)).to(torch.uint32) - - -def unpack10(p): - """(H,W) uint32 -> [3,H,W] float 0-1.""" - import torch - v = p.to(torch.int64) - return torch.stack([v & 1023, (v >> 10) & 1023, (v >> 20) & 1023]).float() / 1023.0 - - -def _ngx_hint(): - """Why load() can fail on a rental: the SDK needs the driver's own libnvidia-ngx.so.1, which a container host may not mount.""" - import ctypes - import sys - if sys.platform != "linux": - return "" - try: - ctypes.CDLL("libnvidia-ngx.so.1") - return " The driver's NGX library is present, so check the driver is 570.190+ / 580.82+ / 590.44+ and the card is Turing or newer." - except OSError: - return (" This machine has no libnvidia-ngx.so.1 (NVIDIA's driver library that RTX VSR needs; some rented GPUs do not mount it, " - "like libnvidia-encode). Use an upscale model instead, or rent a host that exposes the full NVIDIA driver (graphics/video capabilities).") - - def make(o): """-> the upscale function for these (already checked) options. Raises a sentence when the add-on is missing.""" from ..._core import addons @@ -101,7 +77,7 @@ def make(o): try: e.load() except Exception as exc: # noqa: BLE001 -- NvVFXError: say what to do, not just the SDK code - raise RuntimeError(f"{NAME} could not start ({exc}).{_ngx_hint()}") from exc + raise RuntimeError(f"{NAME} could not start ({exc}).{ngx_hint()}") from exc effects[(h, w)] = e return effects[(h, w)] diff --git a/modules/output/render/tests/test_upscale.py b/modules/output/render/tests/test_upscale.py index a00a1b7b..3e68c309 100644 --- a/modules/output/render/tests/test_upscale.py +++ b/modules/output/render/tests/test_upscale.py @@ -85,3 +85,4 @@ def test_rtx_refuses_when_switched_off_even_if_installed(monkeypatch): monkeypatch.setattr(addons, "enabled", lambda _id: False) with pytest.raises(RuntimeError, match="switched off"): rtx.make(rtx.options("{}")) + diff --git a/modules/output/render/upscale.py b/modules/output/render/upscale.py index 0e54ce48..0502ddc0 100644 --- a/modules/output/render/upscale.py +++ b/modules/output/render/upscale.py @@ -7,113 +7,15 @@ and the result is encoded with the source's frame rate and sound. from __future__ import annotations -import json import os -import shutil -import subprocess from . import files, rtx -from ..._core import addons +from ..._core import addons, video_pipe -CHUNK = 8 # frames per upscale call, at most -CHUNK_PIXELS = 2_000_000 # ...and at most this many input pixels per call: a big upscale of 1080p is gigabytes per frame SUBFOLDER = "funui_upscaled" -def probe(path): - """-> (width, height, fps string, has sound).""" - ffprobe = shutil.which("ffprobe") - if not ffprobe: - raise files.ClipError("ffprobe was not found on PATH: install ffmpeg to upscale videos.") - out = subprocess.run([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 files.ClipError(f"no video stream in {os.path.basename(path)}") - return (int(video["width"]), int(video["height"]), video.get("r_frame_rate") or "25", - any(s.get("codec_type") == "audio" for s in streams)) - - -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`. - Returns the frames written. `interrupted()` raises to stop between chunks.""" - import numpy as np - import torch - w, h, fps, has_sound = probe(src) - frame = w * h * 3 - chunk = max(1, min(CHUNK, CHUNK_PIXELS // max(1, w * h))) - ffmpeg = files.ffmpeg() - dec = subprocess.Popen([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, "-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 files.ClipError(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 files.ClipError(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 p is not None: - p.wait() - for pipe in (p.stdin, p.stdout, p.stderr): - try: - if pipe: - pipe.close() - except OSError: # a killed encoder's unflushed pipe - pass - if os.path.exists(dst): - os.remove(dst) - raise - finally: - if dec.poll() is None: - dec.kill() - - -def free_name(out_dir, filename, model): - stem = os.path.splitext(os.path.basename(filename))[0] - tag = os.path.splitext(os.path.basename(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" - return name +probe, upscale_file, scaled_fps, free_name = video_pipe.probe, video_pipe.stream_file, video_pipe.scaled_fps, video_pipe.free_name def models() -> list: