mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 15:41:00 +02:00
Vladimir Mandic
105 lines
4.7 KiB
Python
105 lines
4.7 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from modules import shared, devices
|
|
from modules.logger import log
|
|
from modules.upscaler import Upscaler, UpscalerData
|
|
|
|
|
|
class UpscalerDiffusion(Upscaler):
|
|
def __init__(self, dirname): # pylint: disable=super-init-not-called
|
|
self.name = "nVidia VFX"
|
|
self.user_path = dirname
|
|
"""
|
|
self.scalers = [
|
|
UpscalerData(name="nVidia VFX 1x Denoise Ultra", path="", upscaler=self, model=None, scale=1),
|
|
UpscalerData(name="nVidia VFX 1x Deblur Ultra", path="", upscaler=self, model=None, scale=1),
|
|
UpscalerData(name="nVidia VFX 1x Denoise High", path="", upscaler=self, model=None, scale=1),
|
|
UpscalerData(name="nVidia VFX 1x Deblur High", path="", upscaler=self, model=None, scale=1),
|
|
UpscalerData(name="nVidia VFX 2x Ultra", path="", upscaler=self, model=None, scale=2),
|
|
UpscalerData(name="nVidia VFX 4x Ultra", path="", upscaler=self, model=None, scale=4),
|
|
UpscalerData(name="nVidia VFX 2x High", path="", upscaler=self, model=None, scale=2),
|
|
UpscalerData(name="nVidia VFX 4x High", path="", upscaler=self, model=None, scale=4),
|
|
]
|
|
"""
|
|
self.scalers = []
|
|
self.models = {}
|
|
|
|
def load_model(self, path: str):
|
|
scaler: UpscalerData = [x for x in self.scalers if x.data_path == path or x.name == path]
|
|
if len(scaler) == 0:
|
|
log.error(f"Upscaler cannot match model: type={self.name} model={path}")
|
|
return None
|
|
scaler = scaler[0]
|
|
if self.models.get(path, None) is not None:
|
|
log.debug(f"Upscaler cached: type={scaler.name} model={path}")
|
|
return self.models[path]
|
|
from installer import install
|
|
install('nvidia-vfx')
|
|
|
|
def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor):
|
|
pass
|
|
|
|
def do_upscale(self, img: Image.Image, selected_model):
|
|
devices.torch_gc()
|
|
self.load_model(selected_model)
|
|
|
|
frame = torch.from_numpy(np.array(img)).permute(2, 0, 1).float().to(devices.device) / 255.0
|
|
frame = frame.to(devices.device)
|
|
|
|
try:
|
|
import nvvfx
|
|
except Exception as e:
|
|
log.error(f"Upscaler: failed to import nvvfx: {e}")
|
|
return img
|
|
|
|
config_map = {
|
|
"nVidia VFX 1x Denoise Ultra": nvvfx.VideoSuperRes.QualityLevel.DENOISE_ULTRA,
|
|
"nVidia VFX 1x Deblur Ultra": nvvfx.VideoSuperRes.QualityLevel.DEBLUR_ULTRA,
|
|
"nVidia VFX 1x Denoise High": nvvfx.VideoSuperRes.QualityLevel.DENOISE_HIGH,
|
|
"nVidia VFX 1x Deblur High": nvvfx.VideoSuperRes.QualityLevel.DEBLUR_HIGH,
|
|
"nVidia VFX 2x Ultra": nvvfx.VideoSuperRes.QualityLevel.ULTRA,
|
|
"nVidia VFX 4x Ultra": nvvfx.VideoSuperRes.QualityLevel.ULTRA,
|
|
"nVidia VFX 2x High": nvvfx.VideoSuperRes.QualityLevel.HIGH,
|
|
"nVidia VFX 4x High": nvvfx.VideoSuperRes.QualityLevel.HIGH,
|
|
}
|
|
quality = config_map.get(selected_model, None)
|
|
log.info(f'Upscaler: type="{self.name}" model="{selected_model}" version={nvvfx.__version__} sdk={nvvfx.get_sdk_version()} quality={quality}')
|
|
if self.models.get(selected_model, None) is not None:
|
|
vsr = self.models[selected_model]
|
|
else:
|
|
vsr = nvvfx.VideoSuperRes(quality=quality)
|
|
self.models[selected_model] = vsr
|
|
if '2x' in selected_model:
|
|
vsr.output_width = img.width * 2
|
|
vsr.output_height = img.height * 2
|
|
elif '4x' in selected_model:
|
|
vsr.output_width = img.width * 4
|
|
vsr.output_height = img.height * 4
|
|
elif 'Denoise' in selected_model or 'Deblur' in selected_model or '1x' in selected_model:
|
|
vsr.output_width = img.width
|
|
vsr.output_height = img.height
|
|
else:
|
|
log.error(f"Upscaler: unknown model: {selected_model}")
|
|
return img
|
|
vsr.input_width = img.width
|
|
vsr.input_height = img.height
|
|
|
|
log.debug(f"Upscaler: {vsr}")
|
|
try:
|
|
vsr.load()
|
|
except Exception as e:
|
|
log.error(f"Upscaler: failed to load model: {selected_model} error={e}")
|
|
return img
|
|
self.models[selected_model] = vsr
|
|
|
|
result = vsr.run(frame)
|
|
result = torch.from_dlpack(result.image).clone()
|
|
image = Image.fromarray((result.permute(1, 2, 0).contiguous().cpu().numpy() * 255).astype(np.uint8))
|
|
|
|
if shared.opts.upscaler_unload and selected_model in self.models:
|
|
del self.models[selected_model]
|
|
log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}")
|
|
devices.torch_gc(force=True)
|
|
return image
|