109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
import torch
|
|
import nvvfx
|
|
from typing_extensions import override
|
|
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
|
|
MAX_PIXELS = 1024 * 1024 * 16
|
|
|
|
QUALITY_MAPPING = {
|
|
"LOW": nvvfx.effects.QualityLevel.LOW,
|
|
"MEDIUM": nvvfx.effects.QualityLevel.MEDIUM,
|
|
"HIGH": nvvfx.effects.QualityLevel.HIGH,
|
|
"ULTRA": nvvfx.effects.QualityLevel.ULTRA,
|
|
}
|
|
|
|
DEBLUR_MAPPING = {
|
|
"LOW": nvvfx.effects.QualityLevel.DEBLUR_LOW,
|
|
"MEDIUM": nvvfx.effects.QualityLevel.DEBLUR_MEDIUM,
|
|
"HIGH": nvvfx.effects.QualityLevel.DEBLUR_HIGH,
|
|
"ULTRA": nvvfx.effects.QualityLevel.DEBLUR_ULTRA,
|
|
}
|
|
|
|
|
|
def _run_effect(images: torch.Tensor, quality, output_width: int, output_height: int) -> torch.Tensor:
|
|
out_pixels = output_width * output_height
|
|
batch_size = max(1, MAX_PIXELS // out_pixels)
|
|
|
|
out_batches = []
|
|
|
|
with nvvfx.VideoSuperRes(quality) as sr:
|
|
sr.output_width = output_width
|
|
sr.output_height = output_height
|
|
sr.load()
|
|
|
|
for i in range(0, images.shape[0], batch_size):
|
|
batch = images[i:i + batch_size]
|
|
|
|
batch_cuda = batch.cuda().permute(0, 3, 1, 2).contiguous()
|
|
|
|
batch_outputs = []
|
|
|
|
for j in range(batch_cuda.shape[0]):
|
|
input_frame = batch_cuda[j]
|
|
dlpack_out = sr.run(input_frame).image
|
|
output = torch.from_dlpack(dlpack_out).clone()
|
|
batch_outputs.append(output)
|
|
|
|
batch_out_tensor = torch.stack(batch_outputs, dim=0)
|
|
|
|
batch_out_tensor = batch_out_tensor.permute(0, 2, 3, 1).cpu()
|
|
out_batches.append(batch_out_tensor)
|
|
|
|
return torch.cat(out_batches, dim=0)
|
|
|
|
|
|
class RTXVideoSuperResolution(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RTXVideoSuperResolution",
|
|
display_name="RTX Video Super Resolution",
|
|
category="image/upscaling",
|
|
search_aliases=["rtx", "nvidia", "upscale", "super resolution", "vsr", "deblur", "sharpen"],
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
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.Combo.Input("quality", options=["LOW", "MEDIUM", "HIGH", "ULTRA"], default="ULTRA", tooltip="Upscaling quality level."),
|
|
io.Combo.Input("deblur", options=["OFF", "LOW", "MEDIUM", "HIGH", "ULTRA"], default="OFF", tooltip="Deblur/sharpen strength applied to the source before upscaling. OFF disables it."),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("upscaled_images"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, images: torch.Tensor, scale: float, quality: str, deblur: str) -> io.NodeOutput:
|
|
_, h, w, _ = images.shape
|
|
|
|
output_width = int(w * scale)
|
|
output_height = int(h * scale)
|
|
|
|
output_width = max(8, round(output_width / 8) * 8)
|
|
output_height = max(8, round(output_height / 8) * 8)
|
|
|
|
selected_quality = QUALITY_MAPPING.get(quality, nvvfx.effects.QualityLevel.HIGH)
|
|
|
|
current = images
|
|
|
|
if deblur != "OFF":
|
|
deblur_quality = DEBLUR_MAPPING.get(deblur, nvvfx.effects.QualityLevel.DEBLUR_HIGH)
|
|
current = _run_effect(current, deblur_quality, w, h)
|
|
|
|
final_images = _run_effect(current, selected_quality, output_width, output_height)
|
|
|
|
return io.NodeOutput(final_images)
|
|
|
|
|
|
class NVVFXVideoExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
RTXVideoSuperResolution,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> NVVFXVideoExtension:
|
|
return NVVFXVideoExtension()
|