Merge branch 'dlss5-support' into evolve
This commit is contained in:
Gourieff | 古仁
@@ -3,3 +3,4 @@ __pycache__/
|
||||
.vscode/
|
||||
example
|
||||
input
|
||||
*.dll
|
||||
|
||||
@@ -66,6 +66,8 @@ import scripts.r_masking.segs as masking_segs
|
||||
|
||||
import scripts.reactor_sfw as sfw
|
||||
|
||||
from r_dlssnr.dlss5_node import DLSS5FrameEnhancer
|
||||
|
||||
|
||||
models_dir = folder_paths.models_dir
|
||||
REACTOR_MODELS_PATH = os.path.join(models_dir, "reactor")
|
||||
@@ -1744,6 +1746,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ReActorImageDublicator": ImageDublicator,
|
||||
"ImageRGBA2RGB": ImageRGBA2RGB,
|
||||
"ReActorUnload": ReActorUnload,
|
||||
"DLSS5FrameEnhancer": DLSS5FrameEnhancer,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -1766,4 +1769,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ReActorImageDublicator": "Image Dublicator (List) 🌌 ReActor",
|
||||
"ImageRGBA2RGB": "Convert RGBA to RGB 🌌 ReActor",
|
||||
"ReActorUnload": "Unload ReActor Models 🌌 ReActor",
|
||||
"DLSS5FrameEnhancer": "DLSS5 Frame Enhancer 🌌 ReActor"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
## Third-party DLLs
|
||||
|
||||
1. Download [neuroframe_dlls.zip](https://huggingface.co/datasets/Gourieff/ReActor/blob/main/DLSSNR/neuroframe_dlls.zip) and place `neuroframe_caller.dll`\* and `neuroframe_engine.dll`* here
|
||||
2. Place your `nvngx_dlssnr.dll`** here
|
||||
|
||||
<sub>* Author [Merserk](https://github.com/Merserk), [LICENSE](https://huggingface.co/datasets/Gourieff/ReActor/blob/main/DLSSNR/LICENSE-Merserk.txt)
|
||||
<br>
|
||||
** Public distribution of this file is prohibited by NVIDIA, [LICENSE](https://huggingface.co/datasets/Gourieff/ReActor/blob/main/DLSSNR/LICENSE-NVIDIA-DLSS.txt)
|
||||
@@ -0,0 +1,147 @@
|
||||
import ctypes
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
import numpy as np
|
||||
|
||||
# --- КОНСТАНТЫ ---
|
||||
BRIDGE_ABI_VERSION = 6
|
||||
MEMORY_HOST = 0
|
||||
MEMORY_CUDA = 1
|
||||
MEMORY_NONE = 2
|
||||
|
||||
class NeuralBridgeError(Exception):
|
||||
pass
|
||||
|
||||
# --- C-СТРУКТУРЫ ДЛЯ ВЗАИМОДЕЙСТВИЯ С DLL ---
|
||||
|
||||
class RenderParameters(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("struct_size", ctypes.c_uint32),
|
||||
("abi_version", ctypes.c_uint32),
|
||||
("style", ctypes.c_int32),
|
||||
("intensity", ctypes.c_float),
|
||||
("tone", ctypes.c_float),
|
||||
("structure", ctypes.c_float),
|
||||
("skin", ctypes.c_float),
|
||||
("automask", ctypes.c_int32),
|
||||
("reset", ctypes.c_int32),
|
||||
("color_strength", ctypes.c_float),
|
||||
("tone_preservation", ctypes.c_float),
|
||||
("mask_memory_type", ctypes.c_uint32),
|
||||
("mask_width", ctypes.c_uint32),
|
||||
("mask_height", ctypes.c_uint32),
|
||||
("mask_stride", ctypes.c_uint32),
|
||||
("mask_plane", ctypes.c_uint64), # Из-за uint64 здесь будет 4 байта системного отступа
|
||||
("face_skin_protection", ctypes.c_float),
|
||||
("grain_preservation", ctypes.c_float),
|
||||
("nr_passes", ctypes.c_int32),
|
||||
("shimmer_suppression", ctypes.c_float),
|
||||
("prefer_nvof", ctypes.c_int32),
|
||||
]
|
||||
|
||||
|
||||
class DLSSStandaloneManager:
|
||||
def __init__(self, dll_dir: str):
|
||||
self._lock = threading.RLock()
|
||||
self._library = None
|
||||
self.dll_dir = dll_dir
|
||||
|
||||
def initialize(self, ordinal: int):
|
||||
with self._lock:
|
||||
if self._library is not None:
|
||||
return True
|
||||
|
||||
import os
|
||||
|
||||
if hasattr(os, 'add_dll_directory'):
|
||||
os.add_dll_directory(self.dll_dir)
|
||||
|
||||
engine_path = os.path.join(self.dll_dir, "neuroframe_engine.dll")
|
||||
if not os.path.exists(engine_path):
|
||||
raise NeuralBridgeError(f"Missing DLL: {engine_path}")
|
||||
|
||||
loader = getattr(ctypes, "WinDLL", ctypes.CDLL)
|
||||
try:
|
||||
self._library = loader(engine_path)
|
||||
except OSError as exc:
|
||||
raise NeuralBridgeError(f"DLL load failed: {exc}")
|
||||
|
||||
# Сигнатура инициализации
|
||||
self._library.dlss5nr_init.argtypes = [
|
||||
ctypes.c_int, ctypes.c_wchar_p, ctypes.c_char_p, ctypes.c_int
|
||||
]
|
||||
self._library.dlss5nr_init.restype = ctypes.c_int
|
||||
|
||||
# Сигнатура HOST-рендера (process_v6 вместо process_cuda_v6)
|
||||
c_float_p = ctypes.POINTER(ctypes.c_float)
|
||||
self._library.dlss5nr_process_v6.argtypes = [
|
||||
c_float_p, c_float_p, ctypes.c_int, ctypes.c_int,
|
||||
ctypes.POINTER(RenderParameters), ctypes.c_char_p, ctypes.c_int
|
||||
]
|
||||
self._library.dlss5nr_process_v6.restype = ctypes.c_int
|
||||
|
||||
try:
|
||||
self._library.dlss5nr_frame_abi_version.argtypes = []
|
||||
self._library.dlss5nr_frame_abi_version.restype = ctypes.c_uint32
|
||||
self.actual_abi = self._library.dlss5nr_frame_abi_version()
|
||||
except Exception:
|
||||
self.actual_abi = BRIDGE_ABI_VERSION
|
||||
|
||||
error = ctypes.create_string_buffer(4096)
|
||||
ok = self._library.dlss5nr_init(ordinal, self.dll_dir, error, len(error))
|
||||
|
||||
if not ok:
|
||||
err_msg = error.value.decode('utf-8', errors='ignore')
|
||||
raise NeuralBridgeError(f"Bridge initialization failed: {err_msg}")
|
||||
|
||||
return True
|
||||
|
||||
def process_host(self, source: np.ndarray, destination: np.ndarray, settings: dict, reset: bool, mask: np.ndarray = None):
|
||||
with self._lock:
|
||||
error = ctypes.create_string_buffer(4096)
|
||||
|
||||
params = RenderParameters()
|
||||
params.struct_size = ctypes.sizeof(RenderParameters)
|
||||
params.abi_version = getattr(self, "actual_abi", BRIDGE_ABI_VERSION)
|
||||
|
||||
params.style = int(settings.get("style"))
|
||||
params.intensity = float(settings.get("intensity"))
|
||||
params.tone = float(settings.get("local_tone"))
|
||||
params.structure = float(settings.get("local_structure"))
|
||||
params.skin = float(settings.get("skin_structure"))
|
||||
params.automask = int(bool(settings.get("auto_mask")))
|
||||
params.reset = int(bool(reset))
|
||||
params.color_strength = float(settings.get("color_strength"))
|
||||
params.tone_preservation = float(settings.get("tone_preservation"))
|
||||
params.face_skin_protection = float(settings.get("face_skin_protection"))
|
||||
params.grain_preservation = float(settings.get("grain_preservation"))
|
||||
params.nr_passes = int(settings.get("nr_passes"))
|
||||
params.shimmer_suppression = float(settings.get("shimmer_suppression", 0.0))
|
||||
params.prefer_nvof = int(bool(settings.get("prefer_nvof", False)))
|
||||
|
||||
# Обработка маски через HOST память
|
||||
params.mask_memory_type = MEMORY_NONE
|
||||
if mask is not None:
|
||||
params.mask_memory_type = MEMORY_HOST
|
||||
params.mask_width = int(mask.shape[1])
|
||||
params.mask_height = int(mask.shape[0])
|
||||
params.mask_stride = int(mask.strides[0])
|
||||
params.mask_plane = int(mask.ctypes.data) # Передаем указатель RAM
|
||||
|
||||
c_float_p = ctypes.POINTER(ctypes.c_float)
|
||||
|
||||
# Вызываем HOST функцию (DLL сама разберется с видеокартой)
|
||||
ok = self._library.dlss5nr_process_v6(
|
||||
source.ctypes.data_as(c_float_p),
|
||||
destination.ctypes.data_as(c_float_p),
|
||||
source.shape[1],
|
||||
source.shape[0],
|
||||
ctypes.byref(params),
|
||||
error,
|
||||
len(error)
|
||||
)
|
||||
|
||||
if not ok:
|
||||
err_msg = error.value.decode('utf-8', errors='ignore')
|
||||
raise NeuralBridgeError(f"DLSS-5 process failed: {err_msg}")
|
||||
@@ -0,0 +1,105 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import os
|
||||
import comfy.model_management as model_management
|
||||
from scripts.reactor_logger import logger
|
||||
from .dlss5_core import DLSSStandaloneManager
|
||||
from r_modules.shared import state
|
||||
from reactor_utils import (
|
||||
batch_tensor_to_pil,
|
||||
progress_bar,
|
||||
progress_bar_reset
|
||||
)
|
||||
|
||||
class DLSS5FrameEnhancer:
|
||||
def __init__(self):
|
||||
self.device = model_management.get_torch_device()
|
||||
self.manager = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"style": (["Default", "Nature", "Cinematic"],),
|
||||
"intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "0..2, def: 1.0"}),
|
||||
"local_tone": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "0..2, def: 0.0"}),
|
||||
"local_structure": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "0..2, def: 1.0"}),
|
||||
"skin_structure": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 2.0, "step": 0.05, "tooltip": "-1..2, def: 0.5"}),
|
||||
"color_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "0..1, def: 0.5"}),
|
||||
"tone_preservation": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "0..1, def: 0.5"}),
|
||||
"face_skin_protection": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "0..1, def: 0.0"}),
|
||||
"grain_preservation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "0..1, def: 0.0"}),
|
||||
"nr_passes": ("INT", {"default": 1, "min": 1, "max": 4, "tooltip": "0..4, def: 1"}),
|
||||
"auto_mask": ("BOOLEAN", {"default": False, "label_off": "OFF", "label_on": "ON", "tooltip": "Smart Protection Mask"}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("enhanced_image",)
|
||||
FUNCTION = "enhance"
|
||||
CATEGORY = "🌌 ReActor"
|
||||
|
||||
def load_bridge(self):
|
||||
if self.manager is None:
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
dll_dir = os.path.join(current_dir, "dll")
|
||||
self.manager = DLSSStandaloneManager(dll_dir)
|
||||
ordinal = getattr(self.device, 'index', 0) if self.device.index is not None else 0
|
||||
self.manager.initialize(ordinal)
|
||||
logger.status(f"DLSS-5 Bridge initialized on GPU {ordinal}")
|
||||
|
||||
def enhance(self, image, style, intensity, local_tone, local_structure,
|
||||
skin_structure, color_strength, tone_preservation,
|
||||
face_skin_protection, grain_preservation, nr_passes, auto_mask, mask=None):
|
||||
|
||||
self.load_bridge()
|
||||
|
||||
style_map = {"Default": 0, "Nature": 1, "Cinematic": 2}
|
||||
settings = {
|
||||
"style": style_map[style], "intensity": intensity, "local_tone": local_tone,
|
||||
"local_structure": local_structure, "skin_structure": skin_structure,
|
||||
"color_strength": color_strength, "tone_preservation": tone_preservation,
|
||||
"face_skin_protection": face_skin_protection, "grain_preservation": grain_preservation,
|
||||
"nr_passes": nr_passes, "auto_mask": auto_mask,
|
||||
"shimmer_suppression": 0.0, "prefer_nvof": False
|
||||
}
|
||||
|
||||
enhanced_batch = []
|
||||
|
||||
pil_images = batch_tensor_to_pil(image)
|
||||
pbar = progress_bar(len(pil_images))
|
||||
|
||||
for i in range(len(image)):
|
||||
|
||||
if state.interrupted or model_management.processing_interrupted():
|
||||
logger.status("Interrupted by User")
|
||||
break
|
||||
|
||||
img_np = np.ascontiguousarray(image[i].cpu().numpy().astype(np.float32))
|
||||
dest_np = np.ascontiguousarray(np.zeros_like(img_np))
|
||||
|
||||
mask_np = None
|
||||
if mask is not None:
|
||||
mask_np = np.ascontiguousarray(mask[i].cpu().numpy().astype(np.float32))
|
||||
|
||||
# Вызываем Host-обработчик (без CUDA конфликтов)
|
||||
self.manager.process_host(
|
||||
source=img_np,
|
||||
destination=dest_np,
|
||||
settings=settings,
|
||||
reset=True,
|
||||
mask=mask_np
|
||||
)
|
||||
|
||||
out_tensor = torch.from_numpy(dest_np).to(self.device)
|
||||
enhanced_batch.append(out_tensor)
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
progress_bar_reset(pbar)
|
||||
|
||||
return (torch.stack(enhanced_batch),)
|
||||
Reference in New Issue
Block a user