mirror of
https://github.com/vladmandic/automatic
synced 2026-08-31 17:41:06 +02:00
Vladimir Mandic
81 lines
3.6 KiB
Python
81 lines
3.6 KiB
Python
import os
|
|
import time
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from modules.upscaler import Upscaler, UpscalerData
|
|
from modules import devices, shared, errors
|
|
from modules.logger import log
|
|
|
|
|
|
class UpscalerNVVFX(Upscaler):
|
|
def __init__(self, dirname=None): # pylint: disable=unused-argument
|
|
super().__init__(False)
|
|
self.name = "nVidia VFX"
|
|
self.scalers = [
|
|
UpscalerData("nVidia VFX bicubic", None, self, scale=0),
|
|
UpscalerData("nVidia VFX low", None, self, scale=1),
|
|
UpscalerData("nVidia VFX medium", None, self, scale=2),
|
|
UpscalerData("nVidia VFX high", None, self, scale=3),
|
|
UpscalerData("nVidia VFX ultra", None, self, scale=4),
|
|
UpscalerData("nVidia VFX denoise low", None, self, scale=8),
|
|
UpscalerData("nVidia VFX denoise medium", None, self, scale=9),
|
|
UpscalerData("nVidia VFX denoise high", None, self, scale=10),
|
|
UpscalerData("nVidia VFX denoise ultra", None, self, scale=11),
|
|
UpscalerData("nVidia VFX deblur low", None, self, scale=12),
|
|
UpscalerData("nVidia VFX deblur medium", None, self, scale=13),
|
|
UpscalerData("nVidia VFX deblur high", None, self, scale=14),
|
|
UpscalerData("nVidia VFX deblur ultra", None, self, scale=15),
|
|
UpscalerData("nVidia VFX highbitrate low", None, self, scale=16),
|
|
UpscalerData("nVidia VFX highbitrate medium", None, self, scale=17),
|
|
UpscalerData("nVidia VFX highbitrate high", None, self, scale=18),
|
|
UpscalerData("nVidia VFX highbitrate ultra", None, self, scale=19),
|
|
]
|
|
|
|
# nvvfx overrides upscale instead of do_upscale because it handles scale directly
|
|
def upscale(self, img: Image.Image | torch.Tensor, scale, selected_model: str | None = None):
|
|
if selected_model is None:
|
|
return img
|
|
from installer import install
|
|
install('nvidia-vfx')
|
|
os.environ["NV_VFX_LOG_LEVEL"] = "4"
|
|
os.environ["NV_VFX_DEBUG"] = "1"
|
|
try:
|
|
import nvvfx
|
|
except Exception as e:
|
|
log.error(f"Upscaler: nvvfx {e}")
|
|
errors.display(e, "Upscaler: nvvfx error")
|
|
return img
|
|
|
|
jobid = shared.state.begin('Upscale')
|
|
try:
|
|
t0 = time.time()
|
|
upscaler = self.find_model(selected_model)
|
|
|
|
quality = nvvfx.VideoSuperRes.QualityLevel(upscaler.scale)
|
|
vsr = nvvfx.VideoSuperRes(quality=quality)
|
|
vsr.input_width = img.width
|
|
vsr.input_height = img.height
|
|
_scale = 1.0 if 'DEBLUR' in quality.name or 'DENOISE' in quality.name else scale
|
|
vsr.output_width = int(img.width * _scale)
|
|
vsr.output_height = int(img.height * _scale)
|
|
log.debug(f'Upscaler: id={upscaler.scale} scale={_scale} version={nvvfx.__version__} sdk={nvvfx.get_sdk_version()} vsr={vsr}')
|
|
vsr.load()
|
|
|
|
tensor = torch.from_numpy(np.array(img)).permute(2, 0, 1).float().contiguous().to(devices.device) / 255.0
|
|
result = vsr.run(tensor)
|
|
tensor = torch.from_dlpack(result.image).clone()
|
|
tensor = 255.0 * tensor.permute(1, 2, 0).contiguous().cpu()
|
|
upscaled = Image.fromarray(tensor.numpy().astype(np.uint8))
|
|
|
|
vsr.close()
|
|
t1 = time.time()
|
|
log.debug(f'Upscale: name="{selected_model}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}')
|
|
except nvvfx.NvVFXError as e:
|
|
log.error(f"Upscaler: nvvfx {e}")
|
|
errors.display(e, "Upscaler: nvvfx error")
|
|
upscaled = img
|
|
|
|
shared.state.end(jobid)
|
|
return upscaled
|