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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
This commit is contained in:
Claude
2026-10-29 13:00:01 +00:00
parent 1b37f259ee
commit 97b85a3f09
25 changed files with 728 additions and 200 deletions
+1
View File
@@ -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.
+1
View File
@@ -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.
+104
View File
@@ -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(); });
},
};
+13
View File
@@ -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);
}
+33
View File
@@ -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);
});
+6 -4
View File
@@ -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);
}
+5 -24
View File
@@ -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);
}
+1
View File
@@ -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",
+5
View File
@@ -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 {
+36
View File
@@ -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)); };
+25 -2
View File
@@ -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
+28
View File
@@ -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).")
+41 -4
View File
@@ -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
+1
View File
@@ -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",
+17 -4
View File
@@ -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):
+154
View File
@@ -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
+3 -3
View File
@@ -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"]
+34
View File
@@ -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}
+76
View File
@@ -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
+55
View File
@@ -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"}]})
@@ -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
+4 -30
View File
@@ -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:
+5 -29
View File
@@ -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)]
@@ -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("{}"))
+2 -100
View File
@@ -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: