digital-garbageandClaude Sonnet 5.5
aef6e08e55
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
101 lines
4.7 KiB
Python
101 lines
4.7 KiB
Python
"""NVIDIA RTX Video Super Resolution as an upscale engine, through the nvidia-vfx add-on.
|
|
|
|
Same contract as a model upscale: `make(options)` returns a function [B,H,W,3] float 0-1 -> [B,h,w,3],
|
|
which upscale_file streams a video through. The option names and ranges follow NVIDIA's RTX node
|
|
for ComfyUI (quality, scale or target size, strength, 8/10-bit), so a setup carries over.
|
|
"""
|
|
|
|
import json
|
|
|
|
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,
|
|
*(f"DENOISE_{l}" for l in _LEVELS), *(f"DEBLUR_{l}" for l in _LEVELS),
|
|
*(f"HIGHBITRATE_{l}" for l in _LEVELS), "STREAMING_MEDIUM", "STREAMING_ULTRA")
|
|
SAME_SIZE = ("DENOISE_", "DEBLUR_") # these clean a video at its own size
|
|
DEFAULTS = {"quality": "ULTRA", "resize": "scale", "scale": 2.0, "width": 1920, "height": 1080, "strength": 1.0, "tenbit": False}
|
|
|
|
|
|
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"RTX options are not valid JSON: {exc}") from exc
|
|
if not isinstance(given, dict):
|
|
raise ValueError("RTX options must be an object.")
|
|
o = {**DEFAULTS, **given}
|
|
if o["quality"] not in QUALITIES:
|
|
raise ValueError(f"RTX quality {o['quality']!r} is not one of: {', '.join(QUALITIES)}")
|
|
if o["resize"] not in ("scale", "size"):
|
|
raise ValueError("RTX resize must be 'scale' or 'size'.")
|
|
for key, lo, hi in (("scale", 1.0, 4.0), ("strength", 0.0, 1.0), ("width", 64, 8192), ("height", 64, 8192)):
|
|
if not isinstance(o[key], (int, float)) or isinstance(o[key], bool) or not lo <= o[key] <= hi:
|
|
raise ValueError(f"RTX {key} must be a number from {lo} to {hi}.")
|
|
if not isinstance(o["tenbit"], bool):
|
|
raise ValueError("RTX tenbit must be true or false.")
|
|
return o
|
|
|
|
|
|
def out_size(w, h, o):
|
|
"""Output size: the input's own for denoise/deblur, else the scale or the target, each snapped to a multiple of 8."""
|
|
if o["quality"].startswith(SAME_SIZE):
|
|
return w, h
|
|
ow, oh = (int(w * o["scale"]), int(h * o["scale"])) if o["resize"] == "scale" else (int(o["width"]), int(o["height"]))
|
|
snap = lambda v: max(8, round(v / 8) * 8)
|
|
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 make(o):
|
|
"""-> the upscale function for these (already checked) options. Raises a sentence when the add-on is missing."""
|
|
from ..._core import addons
|
|
if not addons.enabled("nvidia_vsr"): # 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 VideoSuperRes
|
|
except ImportError as exc:
|
|
raise RuntimeError(f"{NAME} is not installed: switch it on in Settings ▸ Modules and restart ComfyUI.") from exc
|
|
import torch
|
|
quality = getattr(VideoSuperRes.QualityLevel, o["quality"], None)
|
|
if quality is None:
|
|
raise RuntimeError(f"This nvidia-vfx has no quality level {o['quality']}.")
|
|
effects = {}
|
|
|
|
def effect(h, w):
|
|
if (h, w) not in effects:
|
|
kw = {"quality": quality}
|
|
if o["strength"] < 1.0:
|
|
kw["strength"] = o["strength"]
|
|
if o["tenbit"]:
|
|
kw["image_encoding"] = VideoSuperRes.ImageEncoding.RGB10A2
|
|
e = VideoSuperRes(**kw)
|
|
e.output_width, e.output_height = out_size(w, h, o)
|
|
e.load()
|
|
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()
|
|
res = effect(h, w).run(pack10(f) if o["tenbit"] else f)
|
|
y = torch.from_dlpack(res.image).clone() # the effect reuses its buffer: copy before the next frame
|
|
out.append((unpack10(y) if o["tenbit"] else y).permute(1, 2, 0))
|
|
return torch.stack(out)
|
|
return run
|