421 lines
17 KiB
Python
421 lines
17 KiB
Python
from enum import Enum
|
|
import math
|
|
from typing import TypedDict
|
|
|
|
import nvvfx
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from typing_extensions import override
|
|
|
|
import comfy.model_management
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
|
|
MAX_BATCH_PIXELS = 1024 * 1024 * 16
|
|
VSR_QUALITY_OPTIONS = list(nvvfx.QualityLevel.__members__)
|
|
|
|
PQ_M1 = 2610 / 16384
|
|
PQ_M2 = 2523 / 4096 * 128
|
|
PQ_C1 = 3424 / 4096
|
|
PQ_C2 = 2413 / 4096 * 32
|
|
PQ_C3 = 2392 / 4096 * 32
|
|
|
|
HLG_A = 0.17883277
|
|
HLG_B = 1 - 4 * HLG_A
|
|
HLG_C = 0.5 - HLG_A * math.log(4 * HLG_A)
|
|
|
|
|
|
class UpscaleType(str, Enum):
|
|
SCALE_BY = "scale by multiplier"
|
|
TARGET_DIMENSIONS = "target dimensions"
|
|
|
|
|
|
class FrameGenerationType(str, Enum):
|
|
MULTIPLIER = "frame rate multiplier"
|
|
TIMESTEP = "specific timestep"
|
|
|
|
|
|
class ImageEncodingType(str, Enum):
|
|
RGB8 = "8-bit RGB"
|
|
RGB10A2 = "10-bit RGB (RGB10A2)"
|
|
|
|
|
|
class HDRColorSpace(str, Enum):
|
|
HLG = "HDR"
|
|
PQ = "HDR PQ"
|
|
|
|
|
|
def get_effect_device() -> tuple[torch.device, int]:
|
|
device = comfy.model_management.get_torch_device()
|
|
if device.type != "cuda":
|
|
raise RuntimeError("NVIDIA RTX Video Effects require a CUDA device.")
|
|
device_index = device.index if device.index is not None else torch.cuda.current_device()
|
|
return device, device_index
|
|
|
|
|
|
def pack_rgb10a2(image: torch.Tensor) -> torch.Tensor:
|
|
rgb = image.clamp(0.0, 1.0).mul(1023.0).round().to(torch.int32)
|
|
packed = rgb[..., 0] | (rgb[..., 1] << 10) | (rgb[..., 2] << 20) | -1073741824
|
|
return packed.view(torch.uint32).contiguous()
|
|
|
|
|
|
def unpack_rgb10a2(packed: torch.Tensor) -> torch.Tensor:
|
|
packed = packed.view(torch.int32)
|
|
red = packed.bitwise_and(0x3FF)
|
|
green = packed.bitwise_right_shift(10).bitwise_and(0x3FF)
|
|
blue = packed.bitwise_right_shift(20).bitwise_and(0x3FF)
|
|
return torch.stack((red, green, blue), dim=-1).to(torch.float32).div_(1023.0)
|
|
|
|
|
|
def pq_to_hlg(image: torch.Tensor, peak_luminance: float) -> torch.Tensor:
|
|
pq = image.clamp(0.0, 1.0).pow(1.0 / PQ_M2)
|
|
display_light = ((pq.sub(PQ_C1).clamp_min(0.0) / (PQ_C2 - PQ_C3 * pq)) ** (1.0 / PQ_M1)) * 10000.0
|
|
display_light.clamp_(max=peak_luminance)
|
|
|
|
display_luminance = 0.2627 * display_light[..., 0] + 0.6780 * display_light[..., 1] + 0.0593 * display_light[..., 2]
|
|
system_gamma = 1.2 + 0.42 * math.log10(peak_luminance / 1000.0)
|
|
luminance_ratio = (display_luminance / peak_luminance).clamp_min(torch.finfo(display_light.dtype).eps)
|
|
scene_light = (display_light / peak_luminance) * luminance_ratio.unsqueeze(-1).pow((1.0 - system_gamma) / system_gamma)
|
|
scene_light.clamp_(0.0, 1.0)
|
|
|
|
low = torch.sqrt(3.0 * scene_light)
|
|
high = HLG_A * torch.log((12.0 * scene_light - HLG_B).clamp_min(torch.finfo(scene_light.dtype).eps)) + HLG_C
|
|
return torch.where(scene_light <= 1.0 / 12.0, low, high)
|
|
|
|
|
|
def prepare_frame(image: torch.Tensor, device: torch.device, image_encoding: ImageEncodingType) -> torch.Tensor:
|
|
image = image.to(device=device, dtype=torch.float32)
|
|
if image_encoding == ImageEncodingType.RGB10A2:
|
|
return pack_rgb10a2(image)
|
|
return image.movedim(-1, 0).contiguous()
|
|
|
|
|
|
def copy_effect_output(dlpack_output, destination: torch.Tensor, image_encoding: ImageEncodingType) -> None:
|
|
output = torch.from_dlpack(dlpack_output).clone()
|
|
if image_encoding == ImageEncodingType.RGB10A2:
|
|
output = unpack_rgb10a2(output)
|
|
else:
|
|
output = output.movedim(0, -1)
|
|
destination.copy_(output)
|
|
|
|
|
|
class RTXVideoSuperResolution(io.ComfyNode):
|
|
class UpscaleTypedDict(TypedDict):
|
|
resize_type: UpscaleType
|
|
scale: float
|
|
width: int
|
|
height: int
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RTXVideoSuperResolution",
|
|
display_name="RTX Video Super Resolution",
|
|
category="image/upscaling",
|
|
description="Upscale, denoise, or deblur images with NVIDIA Video Effects.",
|
|
search_aliases=["rtx", "nvidia", "upscale", "super resolution", "vsr", "denoise", "deblur"],
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.DynamicCombo.Input(
|
|
"resize_type",
|
|
tooltip="Choose to scale by a multiplier or to exact target dimensions. Denoise and deblur modes keep the input dimensions.",
|
|
options=[
|
|
io.DynamicCombo.Option(UpscaleType.SCALE_BY, [
|
|
io.Float.Input("scale", default=2.0, min=1.0, max=4.0, step=0.01, tooltip="Scale factor (e.g., 2.0 doubles the size)."),
|
|
]),
|
|
io.DynamicCombo.Option(UpscaleType.TARGET_DIMENSIONS, [
|
|
io.Int.Input("width", default=1920, min=64, max=8192, step=8, tooltip="Target width in pixels."),
|
|
io.Int.Input("height", default=1080, min=64, max=8192, step=8, tooltip="Target height in pixels."),
|
|
]),
|
|
],
|
|
),
|
|
io.Combo.Input(
|
|
"quality",
|
|
options=VSR_QUALITY_OPTIONS,
|
|
default="ULTRA",
|
|
tooltip="Select the SDK model. DENOISE and DEBLUR process at the original resolution; HIGHBITRATE modes preserve clean sources.",
|
|
),
|
|
io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Effect strength."),
|
|
io.Combo.Input(
|
|
"image_encoding",
|
|
options=[encoding.value for encoding in ImageEncodingType],
|
|
default=ImageEncodingType.RGB8,
|
|
tooltip="SDK pixel encoding. 10-bit preserves more precision but does not by itself convert SDR content to HDR.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("upscaled_images"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
resize_type: UpscaleTypedDict,
|
|
quality: str,
|
|
strength: float = 1.0,
|
|
image_encoding: str = ImageEncodingType.RGB8,
|
|
) -> io.NodeOutput:
|
|
alpha = images[..., 3:4] if images.shape[-1] > 3 else None
|
|
images = images[..., :3]
|
|
b, h, w, c = images.shape
|
|
|
|
selected_type = resize_type["resize_type"]
|
|
if selected_type == UpscaleType.SCALE_BY:
|
|
scale = resize_type["scale"]
|
|
output_width = int(w * scale)
|
|
output_height = int(h * scale)
|
|
elif selected_type == UpscaleType.TARGET_DIMENSIONS:
|
|
output_width = resize_type["width"]
|
|
output_height = resize_type["height"]
|
|
else:
|
|
raise ValueError(f"Unsupported resize type: {selected_type}")
|
|
|
|
if quality.startswith(("DENOISE_", "DEBLUR_")):
|
|
output_width = w
|
|
output_height = h
|
|
else:
|
|
output_width = max(8, round(output_width / 8) * 8)
|
|
output_height = max(8, round(output_height / 8) * 8)
|
|
|
|
output = torch.empty((b, output_height, output_width, c), device=images.device, dtype=images.dtype)
|
|
device, device_index = get_effect_device()
|
|
batch_size = max(1, MAX_BATCH_PIXELS // (output_width * output_height))
|
|
selected_encoding = ImageEncodingType(image_encoding)
|
|
|
|
with torch.cuda.device(device), nvvfx.VideoSuperRes(
|
|
quality=nvvfx.QualityLevel[quality],
|
|
strength=strength,
|
|
device=device_index,
|
|
image_encoding=nvvfx.ImageEncoding.RGB10A2 if selected_encoding == ImageEncodingType.RGB10A2 else nvvfx.ImageEncoding.RGB8,
|
|
) as effect:
|
|
effect.output_width = output_width
|
|
effect.output_height = output_height
|
|
effect.load()
|
|
|
|
for i in range(0, images.shape[0], batch_size):
|
|
batch = images[i:i + batch_size]
|
|
batch_cuda = batch.to(device=device, dtype=torch.float32)
|
|
if selected_encoding == ImageEncodingType.RGB10A2:
|
|
batch_cuda = pack_rgb10a2(batch_cuda)
|
|
else:
|
|
batch_cuda = batch_cuda.movedim(-1, 1).contiguous()
|
|
for j, frame in enumerate(batch_cuda):
|
|
copy_effect_output(effect.run(frame).image, output[i + j], selected_encoding)
|
|
|
|
if alpha is not None:
|
|
upscaled_alpha = F.interpolate(
|
|
alpha.movedim(-1, 1),
|
|
size=(output_height, output_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
).movedim(1, -1)
|
|
output = torch.cat((output, upscaled_alpha), dim=-1)
|
|
|
|
return io.NodeOutput(output)
|
|
|
|
|
|
class RTXTrueHDR(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RTXTrueHDR",
|
|
display_name="RTX TrueHDR",
|
|
category="image/enhancement",
|
|
description="Convert SDR images to HLG or PQ HDR image data with NVIDIA RTX TrueHDR.",
|
|
search_aliases=["rtx", "nvidia", "true hdr", "hdr10", "hlg", "pq", "sdr to hdr"],
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.Int.Input("contrast", default=100, min=0, max=200),
|
|
io.Int.Input("saturation", default=100, min=0, max=200),
|
|
io.Int.Input("middle_gray", default=50, min=10, max=100, tooltip="Middle-grey reference level."),
|
|
io.Int.Input("luminance", default=650, min=400, max=2000, tooltip="Target HDR display peak luminance in nits."),
|
|
io.Boolean.Input("debanding", default=True, tooltip="Apply the SDK's DL-Debander pass."),
|
|
io.Combo.Input(
|
|
"output_colorspace",
|
|
options=[colorspace.value for colorspace in HDRColorSpace],
|
|
default=HDRColorSpace.HLG,
|
|
tooltip="HDR outputs BT.2020/HLG using the selected peak luminance. HDR PQ returns the SDK's BT.2020/PQ signal unchanged.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("hdr_images", tooltip="Normalized RGB values using the selected HDR colorspace."),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
contrast: int,
|
|
saturation: int,
|
|
middle_gray: int,
|
|
luminance: int,
|
|
debanding: bool,
|
|
output_colorspace: str = HDRColorSpace.HLG,
|
|
) -> io.NodeOutput:
|
|
alpha = images[..., 3:4] if images.shape[-1] > 3 else None
|
|
images = images[..., :3]
|
|
device, device_index = get_effect_device()
|
|
output = torch.empty(images.shape, device=images.device, dtype=torch.float32)
|
|
selected_colorspace = HDRColorSpace(output_colorspace)
|
|
|
|
with torch.cuda.device(device), nvvfx.TrueHDR(
|
|
contrast=contrast,
|
|
saturation=saturation,
|
|
middle_gray=middle_gray,
|
|
luminance=luminance,
|
|
debanding_off=int(not debanding),
|
|
device=device_index,
|
|
) as effect:
|
|
effect.load()
|
|
for i, image in enumerate(images):
|
|
result = effect.run(prepare_frame(image, device, ImageEncodingType.RGB8))
|
|
hdr_image = unpack_rgb10a2(torch.from_dlpack(result.image).clone())
|
|
if selected_colorspace == HDRColorSpace.HLG:
|
|
hdr_image = pq_to_hlg(hdr_image, luminance)
|
|
output[i].copy_(hdr_image)
|
|
|
|
if alpha is not None:
|
|
output = torch.cat((output, alpha.to(dtype=output.dtype)), dim=-1)
|
|
|
|
return io.NodeOutput(output)
|
|
|
|
|
|
class RTXVideoFrameGeneration(io.ComfyNode):
|
|
class GenerationTypedDict(TypedDict):
|
|
generation_type: FrameGenerationType
|
|
multiplier: int
|
|
timestep: float
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RTXVideoFrameGeneration",
|
|
display_name="RTX Video Frame Generation",
|
|
category="video",
|
|
description="Generate intermediate video frames with NVIDIA Video Effects.",
|
|
search_aliases=["rtx", "nvidia", "frame generation", "interpolation", "slow motion", "vfi"],
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.DynamicCombo.Input(
|
|
"generation_type",
|
|
options=[
|
|
io.DynamicCombo.Option(FrameGenerationType.MULTIPLIER, [
|
|
io.Int.Input("multiplier", default=2, min=2, max=16, tooltip="Output frame-rate multiplier."),
|
|
]),
|
|
io.DynamicCombo.Option(FrameGenerationType.TIMESTEP, [
|
|
io.Float.Input("timestep", default=0.5, min=0.01, max=0.99, step=0.01, tooltip="Position between each adjacent input pair."),
|
|
]),
|
|
],
|
|
),
|
|
io.Combo.Input("mode", options=["LOW", "MEDIUM", "HIGH"], default="MEDIUM", tooltip="Frame-generation quality mode."),
|
|
io.Boolean.Input("automatic_shot_change_detection", default=True),
|
|
io.Boolean.Input("shot_change", default=False, tooltip="Mark every submitted pair as a shot change. Useful when processing a single pair."),
|
|
io.Combo.Input(
|
|
"image_encoding",
|
|
options=[encoding.value for encoding in ImageEncodingType],
|
|
default=ImageEncodingType.RGB8,
|
|
tooltip="SDK pixel encoding. 10-bit preserves more precision but does not by itself convert SDR content to HDR.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("interpolated_images"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
generation_type: GenerationTypedDict,
|
|
mode: str,
|
|
automatic_shot_change_detection: bool,
|
|
shot_change: bool,
|
|
image_encoding: str = ImageEncodingType.RGB8,
|
|
) -> io.NodeOutput:
|
|
images = images[..., :3]
|
|
if images.shape[0] < 2:
|
|
return io.NodeOutput(images)
|
|
|
|
selected_type = generation_type["generation_type"]
|
|
if selected_type == FrameGenerationType.MULTIPLIER:
|
|
multiplier = generation_type["multiplier"]
|
|
generated_per_pair = multiplier - 1
|
|
elif selected_type == FrameGenerationType.TIMESTEP:
|
|
multiplier = 2
|
|
generated_per_pair = 1
|
|
else:
|
|
raise ValueError(f"Unsupported frame generation type: {selected_type}")
|
|
|
|
output_count = (images.shape[0] - 1) * (generated_per_pair + 1) + 1
|
|
output = torch.empty((output_count, *images.shape[1:]), device=images.device, dtype=images.dtype)
|
|
output[0].copy_(images[0])
|
|
|
|
_, height, width, _ = images.shape
|
|
device, device_index = get_effect_device()
|
|
selected_encoding = ImageEncodingType(image_encoding)
|
|
with torch.cuda.device(device), nvvfx.VideoFrameGeneration(
|
|
mode=nvvfx.VideoFrameGeneration.Mode[mode],
|
|
automatic_shot_change_detection_enabled=automatic_shot_change_detection,
|
|
device=device_index,
|
|
image_encoding=nvvfx.ImageEncoding.RGB10A2 if selected_encoding == ImageEncodingType.RGB10A2 else nvvfx.ImageEncoding.RGB8,
|
|
) as effect:
|
|
effect.input_width = width
|
|
effect.input_height = height
|
|
effect.frame_multiplier = multiplier
|
|
effect.load()
|
|
|
|
output_index = 1
|
|
previous_frame = prepare_frame(images[0], device, selected_encoding)
|
|
for current_image in images[1:]:
|
|
current_frame = prepare_frame(current_image, device, selected_encoding)
|
|
if selected_type == FrameGenerationType.MULTIPLIER:
|
|
for frame_index in range(1, multiplier):
|
|
result = effect.run(
|
|
previous_frame,
|
|
current_frame,
|
|
frame_index=frame_index,
|
|
shot_change=shot_change,
|
|
)
|
|
copy_effect_output(result.image, output[output_index], selected_encoding)
|
|
output_index += 1
|
|
else:
|
|
result = effect.run_at_timestep(
|
|
previous_frame,
|
|
current_frame,
|
|
generation_type["timestep"],
|
|
shot_change=shot_change,
|
|
)
|
|
copy_effect_output(result.image, output[output_index], selected_encoding)
|
|
output_index += 1
|
|
|
|
output[output_index].copy_(current_image)
|
|
output_index += 1
|
|
previous_frame = current_frame
|
|
|
|
return io.NodeOutput(output)
|
|
|
|
|
|
class NVVFXVideoExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
RTXVideoSuperResolution,
|
|
RTXTrueHDR,
|
|
RTXVideoFrameGeneration,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> NVVFXVideoExtension:
|
|
return NVVFXVideoExtension()
|
|
|
|
# Registry scanner compatibility.
|
|
if False:
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RTXVideoSuperResolution": RTXVideoSuperResolution,
|
|
"RTXTrueHDR": RTXTrueHDR,
|
|
"RTXVideoFrameGeneration": RTXVideoFrameGeneration,
|
|
}
|