diff --git a/.gitignore b/.gitignore index c7653ac..bb81153 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__/ .vscode/ example input +*.dll diff --git a/nodes.py b/nodes.py index 0904da2..fba531c 100644 --- a/nodes.py +++ b/nodes.py @@ -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" } diff --git a/r_dlssnr/__init__.py b/r_dlssnr/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/r_dlssnr/dll/README.md b/r_dlssnr/dll/README.md new file mode 100644 index 0000000..5c69621 --- /dev/null +++ b/r_dlssnr/dll/README.md @@ -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 + +* Author [Merserk](https://github.com/Merserk), [LICENSE](https://huggingface.co/datasets/Gourieff/ReActor/blob/main/DLSSNR/LICENSE-Merserk.txt) +
+** Public distribution of this file is prohibited by NVIDIA, [LICENSE](https://huggingface.co/datasets/Gourieff/ReActor/blob/main/DLSSNR/LICENSE-NVIDIA-DLSS.txt) \ No newline at end of file diff --git a/r_dlssnr/dlss5_core.py b/r_dlssnr/dlss5_core.py new file mode 100644 index 0000000..a6a8bed --- /dev/null +++ b/r_dlssnr/dlss5_core.py @@ -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}") diff --git a/r_dlssnr/dlss5_node.py b/r_dlssnr/dlss5_node.py new file mode 100644 index 0000000..ca6026e --- /dev/null +++ b/r_dlssnr/dlss5_node.py @@ -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),)