From 1cdcdb21b548bc66df571bd5d89f7d444af8b482 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 9 Sep 2026 12:47:34 +0200 Subject: [PATCH] dlss initial stable Signed-off-by: Vladimir Mandic --- extensions-builtin/sdnext-modernui | 2 +- modules/processing.py | 14 +- modules/scripts_manager.py | 6 +- modules/ui_definitions.py | 1 + modules/ui_symbols.py | 2 +- scripts/dlss/__init__.py | 0 scripts/dlss/controller_cli.py | 229 +++++++++++++++++ scripts/dlss/controller_srv.py | 400 +++++++++++++++++++++++++++++ scripts/dlss/framegen.py | 209 +++++++++++++++ scripts/dlss/render.py | 162 ++++++++++++ scripts/dlss/supersample.py | 117 +++++++++ scripts/dlss/utils.py | 118 +++++++++ scripts/dlss/verify.py | 201 +++++++++++++++ scripts/dlss_ext.py | 332 ++++++++++++++++++++++++ 14 files changed, 1787 insertions(+), 6 deletions(-) create mode 100644 scripts/dlss/__init__.py create mode 100644 scripts/dlss/controller_cli.py create mode 100644 scripts/dlss/controller_srv.py create mode 100644 scripts/dlss/framegen.py create mode 100644 scripts/dlss/render.py create mode 100644 scripts/dlss/supersample.py create mode 100644 scripts/dlss/utils.py create mode 100644 scripts/dlss/verify.py create mode 100644 scripts/dlss_ext.py diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 8dec87411..2840cb952 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 8dec87411813d6cb413efb4f114c4d96388223ae +Subproject commit 2840cb952665d416768917c780d1acd9175d09cf diff --git a/modules/processing.py b/modules/processing.py index 010b28d72..e6b3b195f 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -356,8 +356,14 @@ def process_samples(p: StableDiffusionProcessing, samples): pp = scripts_manager.PostprocessImageArgs(image) p.scripts.postprocess_image(p, pp) if pp.image is not None: - image = pp.image - + if isinstance(pp.image, list) and len(pp.image) > 0: # post process image can return original+processed + for i, img in enumerate(pp.image): + if i+1 < len(pp.image): + out_images.append(img) + out_infotexts.append(f"Postprocess image {i+1}") + image = pp.image[-1] + else: + image = pp.image grading_params = processing_grading.GradingParams( brightness=getattr(p, 'grading_brightness', 0.0), contrast=getattr(p, 'grading_contrast', 0.0), @@ -609,7 +615,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: audio=audio, ) if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner) and not (shared.state.interrupted or shared.state.skipped): - p.scripts.postprocess(p, results) + _results = p.scripts.postprocess(p, results) + if _results is not None: + results = _results timer.process.record('post') p.ops = list(set(p.ops)) t3 = time.time() diff --git a/modules/scripts_manager.py b/modules/scripts_manager.py index f66d4516a..8598fd839 100644 --- a/modules/scripts_manager.py +++ b/modules/scripts_manager.py @@ -743,15 +743,19 @@ class ScriptRunner: def postprocess(self, p: StableDiffusionProcessing, processed): s = ScriptSummary('postprocess') + _processed = processed for script in self.alwayson_scripts: try: args = resolve_script_args(script, p.script_args, p.per_script_args) if args is not None: - script.postprocess(p, processed, *args) + result = script.postprocess(p, _processed, *args) + if result is not None: # allow postprocessing script to optionally modify results + _processed = result except Exception as e: errors.display(e, f'Running script postprocess: {script.filename}') s.record(script.title()) s.report() + return _processed def postprocess_batch(self, p: StableDiffusionProcessing, images, **kwargs): s = ScriptSummary('postprocess-batch') diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 825722aa3..f8f71554a 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -499,6 +499,7 @@ def create_settings(cmd_opts): "openvino_cache_path": OptionInfo('cache', "Folder for OpenVINO cache", folder=True), "onnx_cached_models_path": OptionInfo(os.path.join(paths.models_path, 'ONNX', 'cache'), "Folder for ONNX cached models", folder=True), "onnx_temp_dir": OptionInfo(os.path.join(paths.models_path, 'ONNX', 'temp'), "Folder for ONNX conversion", folder=True), + "dlss_pkg_path": OptionInfo('', "Folder with DLSS package", gr.Textbox, { "visible": False}), })) # --- Image Options --- diff --git a/modules/ui_symbols.py b/modules/ui_symbols.py index 1a82cfc31..01d21a454 100644 --- a/modules/ui_symbols.py +++ b/modules/ui_symbols.py @@ -23,13 +23,13 @@ sort = '⇕' detect = '📐' folder = '📂' random = '🎲️' -reuse = '♻️' info = 'ℹ' # noqa reset = '🔄' upload = '⬆️' loading = '↺' reuse = '⬅️' search = '🔍' +tools = '🛠' preview = '🖼️' image = '🖌️' resize = '⁜' diff --git a/scripts/dlss/__init__.py b/scripts/dlss/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/dlss/controller_cli.py b/scripts/dlss/controller_cli.py new file mode 100644 index 000000000..f50a97696 --- /dev/null +++ b/scripts/dlss/controller_cli.py @@ -0,0 +1,229 @@ +import os +import json +import shutil +import base64 +import select +import threading +import subprocess +import time +import uuid +import numpy as np +from PIL import Image +from modules.logger import log +from modules.errors import display + + +debug = os.environ.get('SD_DLSS_DEBUG', None) is not None + + +def image_to_nchw(image: Image.Image) -> np.ndarray: + arr = np.array(image.convert('RGB'), dtype=np.uint8) # HWC + return np.ascontiguousarray(arr.transpose(2, 0, 1))[np.newaxis, ...] # 1CHW + + +def images_to_nchw(images: list) -> np.ndarray: + # last image is reference, skip all images that do not have same dimensions as the last image + if len(images) > 1: + ref_size = images[-1].size + images = [image for image in images if image.size == ref_size] + return np.concatenate([image_to_nchw(image) for image in images], axis=0) + + +def nchw_to_images(arr) -> list: + arr = np.asarray(arr) + return [Image.fromarray(arr[i].transpose(1, 2, 0), 'RGB') for i in range(arr.shape[0])] + + +def _encode_value(value): + if isinstance(value, np.ndarray): + arr = np.ascontiguousarray(value) + return { '__ndarray__': True, 'dtype': str(arr.dtype), 'shape': list(arr.shape), 'data': base64.b64encode(arr.tobytes()).decode('ascii') } + return value + + +def _decode_value(value): + if isinstance(value, dict) and value.get('__ndarray__'): + data = base64.b64decode(value['data']) + return np.frombuffer(data, dtype=value['dtype']).reshape(value['shape']) + if isinstance(value, dict): + return { k: _decode_value(v) for k, v in value.items() } + if isinstance(value, list): + return [_decode_value(v) for v in value] + return value + + +class DLSSController: + """Persistent stdio bridge to the DLSS package's long-lived controller worker (app/controller.py).""" + + def __init__(self): + self.process: subprocess.Popen | None = None + self.pkg_path: str | None = None + self.lock = threading.Lock() + + def get_python(self, pkg_path: str): + python_exe = os.path.join(pkg_path, 'bin', 'python-3.13.15-embed-amd64', 'python.exe') + if not os.path.exists(python_exe): + log.error(f'DLSS: path={pkg_path} python={python_exe} not found') + return None + return python_exe + + def is_alive(self) -> bool: + return self.process is not None and self.process.poll() is None + + def stop(self): + process = self.process + self.process = None + if process is None: + return + try: + if process.poll() is None and process.stdin is not None: + line = json.dumps({ 'request_id': str(uuid.uuid4()), 'command': 'shutdown', 'args': [], 'kwargs': {} }) + '\n' + process.stdin.write(line.encode('utf-8')) + process.stdin.flush() + except Exception: + pass + try: + if process.poll() is None: + process.terminate() + process.wait(timeout=5.0) + except Exception: + pass + + def ensure_installed(self, pkg_path: str) -> bool: + if self.is_alive() and self.pkg_path == pkg_path: + return True + self.stop() + python_exe = self.get_python(pkg_path) + if not python_exe: + return False + # create _sdnext directory if it doesn't exist + sdnext_path = os.path.join(pkg_path, '_sdnext') + if not os.path.exists(sdnext_path): + if debug: + log.trace(f'DLSS install: create folder="{sdnext_path}"') + try: + os.makedirs(sdnext_path, exist_ok=True) + except Exception as e: + log.error(f'DLSS install: failed to create folder: {e}') + display(e, 'DLSS') + return False + files_to_copy = ['__init__.py', 'controller_srv.py', 'utils.py', 'verify.py', 'render.py', 'supersample.py', 'framegen.py'] + for file_name in files_to_copy: + # src path is current path of this file + src = os.path.join(os.path.dirname(__file__), file_name) + dst = os.path.join(sdnext_path, file_name) + # not exist or newer + if not os.path.exists(dst) or os.path.getmtime(src) > os.path.getmtime(dst): + if debug: + log.trace(f'DLSS install: copy src="{src}" "{dst}"') + try: + shutil.copy2(src, dst) + except Exception as e: + log.error(f'DLSS install: failed to copy {file_name}: {e}') + display(e, 'DLSS') + return False + return True + + def ensure_started(self, pkg_path: str) -> bool: + if self.is_alive() and self.pkg_path == pkg_path: + return True + self.stop() + python_exe = self.get_python(pkg_path) + if not python_exe: + return False + env = { + 'GRADIO_ANALYTICS_ENABLED': 'False', + 'PYTHONNOUSERSITE': '1', + 'PYTHONIOENCODING': 'utf-8' + } + if debug: + env['SD_DLSS_DEBUG'] = 'True' + log.trace(f'DLSS controller start: env={env}') + try: + self.process = subprocess.Popen( # pylint: disable=consider-using-with + [python_exe, '-m', '_sdnext.controller_srv'], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd=pkg_path, + env=env, + bufsize=0, + ) + except Exception as e: + log.error(f'DLSS controller start: {e}') + display(e, 'DLSS') + self.process = None + return False + self.pkg_path = pkg_path + response = self._send({ 'request_id': str(uuid.uuid4()), 'command': 'status', 'args': [], 'kwargs': {} }, timeout=30.0) + if response is None or response.get('status') != 'ok': + log.error(f'DLSS controller start: response={response}') + self.stop() + return False + if debug: + log.trace(f'DLSS controller start: result={response.get("result")}') + return True + + def _send(self, request: dict, timeout: float = 60.0): + process = self.process + if process is None or process.stdin is None or process.stdout is None: + return None + try: + process.stdin.write((json.dumps(request) + '\n').encode('utf-8')) + process.stdin.flush() + except Exception as e: + log.error(f'DLSS: failed to send request: {e}') + display(e, 'DLSS') + self.process = None + return None + deadline = time.time() + timeout + while True: + remaining = deadline - time.time() + if remaining <= 0: + log.error(f'DLSS controller: timeout={timeout}') + return None + try: + ready, _, _ = select.select([process.stdout], [], [], remaining) + except Exception: + ready = [process.stdout] # select() is not supported on pipes on some platforms: fall back to a blocking read + if not ready: + log.error(f'DLSS controller: timeout={timeout}') + return None + try: + raw = process.stdout.readline() + except Exception as e: + log.error(f'DLSS controller read: {e}') + display(e, 'DLSS') + self.process = None + return None + if not raw: + stderr = process.stderr.read().decode('utf-8', errors='ignore') if process.stderr else '' + log.error(f'DLSS controller process: stderr="{stderr.strip()}"') + self.process = None + return None + try: + return json.loads(raw.decode('utf-8')) + except Exception: + if debug: + log.trace(f'DLSS controller stray output: {raw!r}') + continue # skip any non-JSON noise emitted before the JSON response line + + def call(self, pkg_path: str, command: str, kwargs: dict, timeout: float = 60.0) -> dict: + with self.lock: + if not self.ensure_installed(pkg_path): + return { 'status': 'error', 'result': None, 'error': { 'code': 'not_installed', 'message': 'controller is not installed' } } + if not self.ensure_started(pkg_path): + return { 'status': 'error', 'result': None, 'error': { 'code': 'not_ready', 'message': 'controller failed to start' } } + encoded_kwargs = { key: _encode_value(value) for key, value in kwargs.items() } + request = { 'request_id': str(uuid.uuid4()), 'command': command, 'args': [], 'kwargs': encoded_kwargs } + if debug: + log.trace(f'DLSS controller request: command={command}') + response = self._send(request, timeout=timeout) + if response is None: + return { 'status': 'error', 'result': None, 'error': { 'code': 'not_ready', 'message': 'controller is not responding' } } + if isinstance(response.get('result'), (dict, list)): + response['result'] = _decode_value(response['result']) + return response + + +controller = DLSSController() diff --git a/scripts/dlss/controller_srv.py b/scripts/dlss/controller_srv.py new file mode 100644 index 000000000..44117c248 --- /dev/null +++ b/scripts/dlss/controller_srv.py @@ -0,0 +1,400 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +# when launched directly as the stdio bridge subprocess, keep the real stdout clean of stray +# prints from native imports/libraries below so only explicit JSON response lines reach the pipe +_stdio_stdout = None +if __name__ == "__main__": + _stdio_stdout = sys.stdout + sys.stdout = sys.stderr + +import base64 +import ctypes +import dataclasses +import json +import multiprocessing as mp +import os +import threading +import uuid +from queue import Empty +from typing import Any + +import numpy as np + +from .utils import StandaloneError, log +from .framegen import DLSSFrameGen, InterpolationOptions +from .render import DLSSNeuralRenderer, RenderOptions +from .supersample import DLSSSuperSample, UpscaleOptions +from .verify import DLSSVerify, VerifyOptions + +_SUPPORTED_COMMANDS = { + "status", + "verify", + "render", + "upscale", + "framegen", + "cancel", + "reset", + "shutdown", +} + + +def _response(request_id: str, *, status: str, result: Any = None, error: dict[str, str] | None = None, diagnostics: dict[str, Any] | None = None) -> dict[str, Any]: + return { + "request_id": request_id, + "status": status, + "result": result, + "error": error, + "diagnostics": diagnostics or {}, + } + + +def _coerce_options(options: Any, *, default: Any, option_type: type[Any]) -> Any: + if options is None: + return default + if isinstance(options, option_type): + return options + if isinstance(options, dict): + return option_type(**options) + raise TypeError(f"Expected {option_type.__name__} or dict, got {type(options).__name__}") + + +def _set_shared_text(buffer: Any, value: str) -> None: + with buffer.get_lock(): + for index in range(len(buffer)): + buffer[index] = "\0" + for index, character in enumerate(value[: len(buffer) - 1]): + buffer[index] = character + + +def _get_shared_text(buffer: Any) -> str: + with buffer.get_lock(): + return "".join(buffer).split("\0", 1)[0] + + +def _dispatch_command(command: str, request_id: str, args: tuple[Any, ...], kwargs: dict[str, Any], *, busy: Any = None, current_request_id: Any = None, current_command: Any = None) -> dict[str, Any]: # pylint: disable=unused-argument + log.debug(f'DLSSController dispatch: command={command} id={request_id}') + if command == "status": + is_busy = bool(busy.value) if busy is not None else False + active_request_id = _get_shared_text(current_request_id) if current_request_id is not None else "" + active_command = _get_shared_text(current_command) if current_command is not None else "" + return _response( + request_id, + status="ok", + result={ + "ready": True, + "busy": is_busy, + "pid": os.getpid(), + "id": active_request_id if is_busy else "", + "job": active_command if is_busy else "idle", + }, + diagnostics={"controller": "ready"}, + ) + + if command == "reset": + return _response(request_id, status="ok", result={"reset": True}, diagnostics={"controller": "reset"}) + + if command == "cancel": + target_id = str(kwargs.get("request_id") or "") + return _response(request_id, status="ok", result={"cancelled": bool(target_id), "target_request_id": target_id}, diagnostics={"controller": "cancelled"}) + + if command == "verify": + gpu_uuid = str(kwargs.get("gpu_uuid", "auto")) + options = _coerce_options(kwargs.get("options"), default=VerifyOptions(), option_type=VerifyOptions) + result = DLSSVerify()(gpu_uuid, options) + return _response(request_id, status="ok", result={"ok": result.ok, "report": result.to_dict()}, diagnostics={"gpu": result.gpu or {}}) + + if command == "render": + images = kwargs.get("images") + if images is None: + raise StandaloneError("invalid_arguments", "Missing required 'images' argument for render command.") + options = _coerce_options(kwargs.get("options"), default=RenderOptions(), option_type=RenderOptions) + result = DLSSNeuralRenderer()(np.asarray(images), options) + return _response(request_id, status="ok", result=result, diagnostics={"shape": list(result.shape)}) + + if command == "upscale": + images = kwargs.get("images") + if images is None: + raise StandaloneError("invalid_arguments", "Missing required 'images' argument for upscale command.") + options = _coerce_options(kwargs.get("options"), default=UpscaleOptions(), option_type=UpscaleOptions) + result = DLSSSuperSample()(np.asarray(images), options) + return _response(request_id, status="ok", result=result, diagnostics={"shape": list(result.shape)}) + + if command == "framegen": + frames = kwargs.get("frames") + if frames is None: + raise StandaloneError("invalid_arguments", "Missing required 'frames' argument for framegen command.") + source_fps = kwargs.get("source_fps") + target_fps = kwargs.get("target_fps") + if source_fps is None or target_fps is None: + raise StandaloneError("invalid_arguments", "framegen requires both 'source_fps' and 'target_fps'.") + options = _coerce_options(kwargs.get("options"), default=InterpolationOptions(), option_type=InterpolationOptions) + result = DLSSFrameGen()(np.asarray(frames), source_fps, target_fps, options) + return _response(request_id, status="ok", result=result, diagnostics={"shape": list(result.shape)}) + + raise StandaloneError("invalid_arguments", f"Unsupported controller command: {command!r}") + + +def _controller_worker(request_queue: mp.Queue, response_queue: mp.Queue, busy: Any, current_request_id: Any, current_command: Any) -> None: + worker_lock = threading.Lock() + while True: + try: + request = request_queue.get(timeout=0.25) + except Empty: + continue + + if not isinstance(request, dict): + response_queue.put(_response(str(uuid.uuid4()), status="error", error={"code": "invalid_arguments", "message": "Controller request must be a dict."})) + continue + + request_id = str(request.get("request_id") or uuid.uuid4()) + command = str(request.get("command") or "").strip().lower() + args = tuple(request.get("args", ())) + kwargs = dict(request.get("kwargs", {})) + + if command == "shutdown": + log.debug(f'DLSSController shutdown: id={request_id}') + response_queue.put(_response(request_id, status="ok", result={"shutdown": True}, diagnostics={"controller": "shutdown"})) + return + + if command not in _SUPPORTED_COMMANDS: + log.warning(f'DLSSController: command={command} id={request_id} unsupported') + response_queue.put(_response(request_id, status="error", error={"code": "invalid_arguments", "message": f"Unsupported command: {command!r}"}, diagnostics={"controller": "invalid_command"})) + continue + + try: + if command != "status": + busy.value = True + _set_shared_text(current_request_id, request_id) + _set_shared_text(current_command, command) + with worker_lock: + response = _dispatch_command(command, request_id, args, kwargs, busy=busy, current_request_id=current_request_id, current_command=current_command) + response_queue.put(response) + except StandaloneError as exc: + log.error(f'DLSSController: StandaloneError command={command} id={request_id} code={exc.code} message={exc.message}') + response_queue.put(_response(request_id, status="error", error={"code": exc.code, "message": exc.message}, diagnostics={"controller": "error"})) + except Exception as exc: # pragma: no cover - defensive catch for controller safety + log.error(f'DLSSController: unexpected exception command={command} id={request_id} error={exc}') + response_queue.put(_response(request_id, status="error", error={"code": "processing_failed", "message": str(exc)}, diagnostics={"controller": "error"})) + finally: + if command != "status": + busy.value = False + _set_shared_text(current_request_id, "") + _set_shared_text(current_command, "") + + +class ControllerProcess(mp.Process): + def __init__(self, request_queue: mp.Queue | None = None, response_queue: mp.Queue | None = None, *, busy: Any = None, current_request_id: Any = None, current_command: Any = None, ctx: mp.context.BaseContext | None = None) -> None: + self.ctx = ctx or mp.get_context("spawn") + self.request_queue = request_queue or self.ctx.Queue() + self.response_queue = response_queue or self.ctx.Queue() + self.busy = busy or self.ctx.Value("b", False) + self.current_request_id = current_request_id or self.ctx.Array(ctypes.c_wchar, 256) + self.current_command = current_command or self.ctx.Array(ctypes.c_wchar, 64) + super().__init__(target=_controller_worker, args=(self.request_queue, self.response_queue, self.busy, self.current_request_id, self.current_command)) + + +class ControllerClient: + """Simple client wrapper for callers that want a long-lived controller process.""" + + def __init__(self, request_queue: mp.Queue | None = None, response_queue: mp.Queue | None = None, *, process: ControllerProcess | None = None, timeout: float = 30.0, ctx: mp.context.BaseContext | None = None) -> None: + self.ctx = ctx or mp.get_context("spawn") + self.request_queue = request_queue or self.ctx.Queue() + self.response_queue = response_queue or self.ctx.Queue() + self.timeout = timeout + self.process = process + self.busy = process.busy if process is not None else self.ctx.Value("b", False) + self.current_request_id = process.current_request_id if process is not None else self.ctx.Array(ctypes.c_wchar, 256) + self.current_command = process.current_command if process is not None else self.ctx.Array(ctypes.c_wchar, 64) + self._pending: dict[str, dict[str, Any]] = {} + + def start(self) -> "ControllerClient": + if self.process is None or not self.process.is_alive(): + self.process = ControllerProcess(self.request_queue, self.response_queue, busy=self.busy, current_request_id=self.current_request_id, current_command=self.current_command, ctx=self.ctx) + self.process.start() + log.info(f'DLSSController: pid={self.process.pid} started') + return self + + def _send_and_wait(self, command: str, *args: Any, **kwargs: Any) -> dict[str, Any]: + if self.process is None or not self.process.is_alive(): + self.start() + request_id = str(uuid.uuid4()) + request = {"request_id": request_id, "command": command, "args": list(args), "kwargs": kwargs} + self.request_queue.put(request) + + while True: + try: + response = self.response_queue.get(timeout=self.timeout) + except Empty as exc: + raise TimeoutError(f"Controller request timed out for command {command!r}.") from exc + if response.get("request_id") == request_id: + return response + self._pending[response.get("request_id", str(uuid.uuid4()))] = response + + def status(self) -> dict[str, Any]: + if self.process is None or not self.process.is_alive(): + self.start() + is_busy = bool(self.busy.value) + active_request_id = _get_shared_text(self.current_request_id) + active_command = _get_shared_text(self.current_command) + return _response( + str(uuid.uuid4()), + status="ok", + result={ + "ready": self.process.is_alive(), + "busy": is_busy, + "pid": self.process.pid, + "id": active_request_id if is_busy else "", + "job": active_command if is_busy else "idle", + }, + diagnostics={"controller": "busy" if is_busy else "ready"}, + ) + + def verify(self, *, gpu_uuid: str = "auto", options: VerifyOptions | dict[str, Any] | None = None) -> dict[str, Any]: + return self._send_and_wait("verify", gpu_uuid=gpu_uuid, options=options) + + def render(self, *, images: np.ndarray, options: RenderOptions | dict[str, Any] | None = None) -> dict[str, Any]: + return self._send_and_wait("render", images=images, options=options) + + def upscale(self, *, images: np.ndarray, options: UpscaleOptions | dict[str, Any] | None = None) -> dict[str, Any]: + return self._send_and_wait("upscale", images=images, options=options) + + def framegen(self, *, frames: np.ndarray, source_fps: float | str, target_fps: float | str, options: InterpolationOptions | dict[str, Any] | None = None) -> dict[str, Any]: + return self._send_and_wait("framegen", frames=frames, source_fps=source_fps, target_fps=target_fps, options=options) + + def cancel(self, request_id: str) -> dict[str, Any]: + return self._send_and_wait("cancel", request_id=request_id) + + def reset(self) -> dict[str, Any]: + return self._send_and_wait("reset") + + def shutdown(self) -> dict[str, Any]: + if self.process is None or not self.process.is_alive(): + return {"request_id": "shutdown", "status": "ok", "result": {"shutdown": True}, "error": None, "diagnostics": {}} + response = self._send_and_wait("shutdown") + if self.process.is_alive(): + self.process.join(timeout=5.0) + return response + + def close(self) -> None: + try: + self.shutdown() + except Exception: + pass + if self.process is not None and self.process.is_alive(): + self.process.terminate() + self.process.join(timeout=5.0) + + def __enter__(self) -> "ControllerClient": + return self.start() + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.close() + + +def start_controller(*, request_queue: mp.Queue | None = None, response_queue: mp.Queue | None = None, timeout: float = 30.0, ctx: mp.context.BaseContext | None = None) -> ControllerClient: + client = ControllerClient(request_queue=request_queue, response_queue=response_queue, timeout=timeout, ctx=ctx) + return client.start() + + +# ---- stdio bridge ----------------------------------------------------------------- +# Used only when this module is launched directly as a subprocess, e.g. +# `python.exe app/controller.py`, to drive the same dispatch logic over a single +# stdin/stdout JSON-lines protocol instead of multiprocessing queues. This allows an +# external caller running a different Python interpreter (for example WSL/Linux Python +# invoking the packaged Windows embedded python.exe) to reuse one long-lived worker. + + +def _stdio_encode(value: Any) -> Any: + if isinstance(value, np.ndarray): + arr = np.ascontiguousarray(value) + return {"__ndarray__": True, "dtype": str(arr.dtype), "shape": list(arr.shape), "data": base64.b64encode(arr.tobytes()).decode("ascii")} + if hasattr(value, "to_dict") and callable(value.to_dict): + return _stdio_encode(value.to_dict()) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return {k: _stdio_encode(v) for k, v in dataclasses.asdict(value).items()} + if isinstance(value, dict): + return {k: _stdio_encode(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_stdio_encode(v) for v in value] + return value + + +def _stdio_decode(value: Any) -> Any: + if isinstance(value, dict) and value.get("__ndarray__"): + data = base64.b64decode(value["data"]) + return np.frombuffer(data, dtype=value["dtype"]).reshape(value["shape"]) + if isinstance(value, dict): + return {k: _stdio_decode(v) for k, v in value.items()} + if isinstance(value, list): + return [_stdio_decode(v) for v in value] + return value + + +def _stdio_write(payload: dict[str, Any]) -> None: + _stdio_stdout.write(json.dumps(payload) + "\n") + _stdio_stdout.flush() + + +def _stdio_error(request_id: str, code: str, message: str) -> dict[str, Any]: + return {"request_id": request_id, "status": "error", "result": None, "error": {"code": code, "message": message}, "diagnostics": {}} + + +def _stdio_main() -> None: + # long-lived worker: one JSON request per line on stdin, one JSON response per line on the real stdout + for raw_line in sys.stdin: + line = raw_line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError as exc: + _stdio_write(_stdio_error("", "invalid_arguments", f"Malformed request: {exc}")) + continue + + request_id = str(request.get("request_id") or "") + command = str(request.get("command") or "").strip().lower() + args = tuple(request.get("args", ())) + try: + kwargs = _stdio_decode(dict(request.get("kwargs", {}))) + except Exception as exc: + _stdio_write(_stdio_error(request_id, "invalid_arguments", f"Failed to decode request payload: {exc}")) + continue + + if command == "shutdown": + _stdio_write({"request_id": request_id, "status": "ok", "result": {"shutdown": True}, "error": None, "diagnostics": {}}) + return + if command not in _SUPPORTED_COMMANDS: + _stdio_write(_stdio_error(request_id, "invalid_arguments", f"Unsupported command: {command!r}")) + continue + + try: + response = _dispatch_command(command, request_id, args, kwargs) + except StandaloneError as exc: + response = _stdio_error(request_id, exc.code, exc.message) + except Exception as exc: # pragma: no cover - defensive catch for controller safety + response = _stdio_error(request_id, "processing_failed", str(exc)) + _stdio_write(_stdio_encode(response)) + + +if __name__ == "__main__": + _stdio_main() + + def __enter__(self) -> "ControllerClient": + return self.start() + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: # pylint: disable=unused-argument + self.close() + + +__all__ = [ + "ControllerClient", + "ControllerProcess", + "start_controller", +] diff --git a/scripts/dlss/framegen.py b/scripts/dlss/framegen.py new file mode 100644 index 000000000..4dda189b0 --- /dev/null +++ b/scripts/dlss/framegen.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from dataclasses import dataclass +from fractions import Fraction +from typing import Any + +import numpy as np + +from src.core.jobs import JobController, active_job +from src.frame_interpolation.capabilities import probe_frame_interpolation_capabilities +from src.frame_interpolation.guides import DLSSGGuideGenerator +from src.frame_interpolation.models import ENGINE_CHOICES, resolve_target_rate +from src.frame_interpolation.native import DirectDLSSGSession +from src.frame_interpolation.scheduler import choose_interpolation_plan, output_frame_count + +from .utils import StandaloneError, nchw_image_to_hwc, rgba_to_rgb_nchw, validate_nchw, rgb_to_rgba, log + + +@dataclass(frozen=True, slots=True) +class InterpolationOptions: + ai_gpu_uuid: str = "auto" + engine: str = "Auto" + + def validate(self) -> None: + if self.engine not in ENGINE_CHOICES: + raise ValueError(f"Unknown frame interpolation engine: {self.engine!r}.") + + +@dataclass(frozen=True, slots=True) +class _TimedFrame: + rgba: np.ndarray + timestamp: Fraction + + +class _Stage: + def __init__(self, session: DirectDLSSGSession, width: int, height: int) -> None: + self.session = session + self.guides = DLSSGGuideGenerator(width, height) + self.previous: _TimedFrame | None = None + + def push(self, frame: _TimedFrame) -> list[_TimedFrame]: + previous = self.previous + guide = self.guides.process(frame.rgba, force_reset=previous is not None and frame.timestamp <= previous.timestamp) + self.previous = frame + generated = self.session.process_frame( + frame.rgba, + guide.motion, + frame.timestamp, + reset=previous is None or guide.reset, + ) + result: list[_TimedFrame] = [] + if previous is not None and not guide.reset: + interval = frame.timestamp - previous.timestamp + count = len(generated) + for index, rgba in enumerate(generated, start=1): + result.append(_TimedFrame( + rgba, + previous.timestamp + interval * Fraction(index, count + 1), + )) + result.append(frame) + return result + + +class DLSSFrameGen: + """In-memory RGB NCHW constant-frame-rate DLSS frame interpolation.""" + + def __init__(self) -> None: + log.info('DLSSFrameGen: init') + self.diagnostics: dict[str, Any] = {} + self.last_report: dict[str, Any] = {} + + def __call__( + self, + frames: np.ndarray, + source_fps: str | int | float | Fraction, + target_fps: str | int | float | Fraction, + options: InterpolationOptions | None = None, + *, + controller: JobController | None = None, + ) -> np.ndarray: + log.info('DLSSFrameGen: call') + options = options or InterpolationOptions() + options.validate() + batch, _, height, width = validate_nchw(frames, name="frames") + if width < 64 or height < 64: + raise StandaloneError("invalid_dimensions", "Frame interpolation requires frames at least 64x64 pixels.") + source_rate = resolve_target_rate(source_fps) + target_rate = resolve_target_rate(target_fps) + own_controller = controller or JobController() + log.debug(f'DLSSFrameGen: controller={own_controller}') + try: + with active_job(own_controller) as active_controller: + capabilities = probe_frame_interpolation_capabilities(options.ai_gpu_uuid) + log.debug(f'DLSSFrameGen: capabilities={capabilities}') + if not capabilities.available: + raise StandaloneError( + "feature_unavailable", + "DLSS Frame Generation is unavailable. " + capabilities.detail, + ) + plan = choose_interpolation_plan( + source_rate, + target_rate, + options.engine, + capabilities.native_multiplier, + cfr=True, + ) + source_frames = [ + _TimedFrame(rgb_to_rgba(nchw_image_to_hwc(frames, index, name="frames")), Fraction(index, 1) / source_rate) + for index in range(batch) + ] + if plan.generated_per_interval == 0: + result = self._resample_source(source_frames, target_rate, source_rate) + else: + result = self._generate(source_frames, plan, active_controller, width, height) + expected = output_frame_count(Fraction(batch, 1) / source_rate, target_rate) + if len(result) != expected: + log.error(f'DLSSFrameGen: result length={len(result)} expected={expected}') + raise StandaloneError( + "invalid_native_output", + f"Interpolation produced {len(result)} frames; expected {expected}.", + ) + output = np.stack([rgba_to_rgb_nchw(item.rgba)[0] for item in result], axis=0) + log.debug(f'DLSSFrameGen: output={output.shape}') + self.diagnostics = { + "gpu": capabilities.gpu, + "driver": capabilities.driver, + "runtime_version": capabilities.runtime_version, + "worker_version": capabilities.worker_version, + "selected_path": plan.path, + "native_multiplier": plan.native_multiplier, + "cascade_stages": plan.cascade_stages, + } + except StandaloneError: + raise + except Exception as exc: + log.error(f'DLSSFrameGen: unexpected exception {exc}') + raise StandaloneError("processing_failed", f"DLSS frame interpolation failed: {exc}") from exc + self.last_report = { + "input_shape": tuple(frames.shape), + "output_shape": tuple(output.shape), + "source_fps": str(source_rate), + "target_fps": str(target_rate), + } + return np.ascontiguousarray(output) + + @staticmethod + def _resample_source(frames: list[_TimedFrame], target_rate: Fraction, source_rate: Fraction) -> list[_TimedFrame]: + count = output_frame_count(Fraction(len(frames), 1) / source_rate, target_rate) + result: list[_TimedFrame] = [] + for index in range(count): + ideal = Fraction(index, 1) / target_rate + selected = min(frames, key=lambda frame, target=ideal: abs(frame.timestamp - target)) + result.append(_TimedFrame(selected.rgba.copy(), ideal)) + return result + + @staticmethod + def _generate(source_frames, plan, controller, width, height) -> list[_TimedFrame]: + sessions: list[DirectDLSSGSession] = [] + stages: list[_Stage] = [] + try: + stage_count = plan.cascade_stages or 1 + for stage_index in range(stage_count): + generated_count = ( + plan.generated_per_interval + if plan.path == "Native DLSSG" + else 1 + ) + expected_frames = ( + len(source_frames) + if stage_index == 0 + else max(1, (len(source_frames) - 1) * (1 << stage_index) + 1) + ) + session = DirectDLSSGSession( + width, + height, + expected_frames, + generated_count, + controller, + ) + sessions.append(session) + stages.append(_Stage(session, width, height)) + candidates: list[_TimedFrame] = [] + for source in source_frames: + if controller.cancel.is_set(): + raise StandaloneError("cancelled", "Frame interpolation was cancelled.") + items = [source] + for stage in stages: + next_items: list[_TimedFrame] = [] + for item in items: + next_items.extend(stage.push(item)) + items = next_items + candidates.extend(items) + duration = Fraction(len(source_frames), 1) / plan.source_rate + count = output_frame_count(duration, plan.target_rate) + result: list[_TimedFrame] = [] + for index in range(count): + ideal = Fraction(index, 1) / plan.target_rate + selected = min(candidates, key=lambda frame, target=ideal: abs(frame.timestamp - target)) + result.append(_TimedFrame(selected.rgba.copy(), ideal)) + return result + finally: + for session in reversed(sessions): + try: + session.close() + except (OSError, RuntimeError, ValueError): + session.abort() + + +__all__ = ["DLSSFrameGen", "InterpolationOptions"] diff --git a/scripts/dlss/render.py b/scripts/dlss/render.py new file mode 100644 index 000000000..4e34b6e65 --- /dev/null +++ b/scripts/dlss/render.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from src.core.jobs import JobController, active_job +from src.core.runtime import ( + DLSSFrameSession, + prepare_runtime, + resolve_native_settings, + resolve_output_size, + resolve_runtime_ai_gpu, + resolve_upscaling_mode, + resize_fit, +) +from src.neural_rendering.image.models import ImageConversionOptions + +from .utils import StandaloneError, nchw_image_to_hwc, rgba_to_rgb_nchw, validate_nchw, rgb_to_rgba, log + + +@dataclass(frozen=True, slots=True) +class RenderOptions: + ai_gpu_uuid: str = "auto" + nr_style: str = "Default" + nr_intensity: float = 1.0 + local_tone_strength: float = 1.0 + local_structure_strength: float = 1.0 + skin_structure_strength: float = -1.0 + upscaling_factor: float = 1.0 + warmup_frames: int = 0 + nr_preset: str = "Default" + automatic_mask: bool = False + dlss_model_preset: str = "Default" + + def source_options(self) -> ImageConversionOptions: + return ImageConversionOptions( + ai_gpu_uuid=self.ai_gpu_uuid, + nr_style=self.nr_style, + nr_intensity=self.nr_intensity, + local_tone_strength=self.local_tone_strength, + local_structure_strength=self.local_structure_strength, + skin_structure_strength=self.skin_structure_strength, + upscaling_factor=self.upscaling_factor, + warmup_frames=self.warmup_frames, + nr_preset=self.nr_preset, + automatic_mask=self.automatic_mask, + dlss_model_preset=self.dlss_model_preset, + ) + + def validate(self) -> None: + if isinstance(self.warmup_frames, bool) or not isinstance(self.warmup_frames, int) or self.warmup_frames < 0: + raise ValueError("warmup_frames must be a non-negative integer.") + if not isinstance(self.automatic_mask, bool): + raise ValueError("automatic_mask must be a boolean.") + options = self.source_options() + resolve_upscaling_mode(options.upscaling_factor) + resolve_native_settings(options) + + +class DLSSNeuralRenderer: + """RGB NCHW adapter for still-image DLSS Neural Rendering.""" + + def __init__(self) -> None: + log.info('DLSSNeuralRenderer: init') + self.diagnostics: dict[str, Any] = {} + self.last_report: dict[str, Any] = {} + + def __call__( + self, + images: np.ndarray, + options: RenderOptions | None = None, + *, + controller: JobController | None = None, + ) -> np.ndarray: + log.info('DLSSNeuralRenderer: call') + options = options or RenderOptions() + options.validate() + batch, _, height, width = validate_nchw(images, name="images") + if width < 64 or height < 64: + raise StandaloneError("invalid_dimensions", "DLSS Neural Rendering requires images at least 64x64 pixels.") + output_width, output_height = resolve_output_size(width, height, options.upscaling_factor) + own_controller = controller or JobController() + log.debug(f'DLSSNeuralRenderer: controller={own_controller}') + outputs: list[np.ndarray] = [] + try: + with active_job(own_controller) as active_controller: + prepared = prepare_runtime() + log.debug(f'DLSSNeuralRenderer: runtime={prepared}') + gpu = resolve_runtime_ai_gpu(prepared.gpus, prepared.runtime_bundle, options.ai_gpu_uuid) + log.debug(f'DLSSNeuralRenderer: gpu={gpu}') + factor, mode = resolve_upscaling_mode(options.upscaling_factor) + native_settings = resolve_native_settings(options.source_options()) + session_diagnostics: list[dict[str, Any]] = [] + for index in range(batch): + if active_controller.cancel.is_set(): + raise StandaloneError("cancelled", "Neural rendering was cancelled.") + try: + session = DLSSFrameSession( + input_width=width, + input_height=height, + output_width=output_width, + output_height=output_height, + frame_count=1, + warmup_frames=options.warmup_frames, + factor=factor, + mode=mode, + native_settings=native_settings, + gpu=gpu, + runtime_bundle=prepared.runtime_bundle, + controller=active_controller, + ) + log.debug(f'DLSSNeuralRenderer: session={session}') + rgb = nchw_image_to_hwc(images, index, name="images") + rgba = rgb_to_rgba(rgb) + render_input = resize_fit(rgba, session.render_width, session.render_height) + motion = np.zeros((session.render_height, session.render_width, 2), dtype=np.float16) + processed, _ = session.process( + index=0, + rgba=render_input, + motion=motion, + reset=True, + pts=0, + ) + log.debug(f'DLSSNeuralRenderer: processed={processed.shape}') + outputs.append(rgba_to_rgb_nchw(processed)[0]) + session_diagnostics.append({ + "render_width": session.render_width, + "render_height": session.render_height, + "applied_dlss_model_preset": session.applied_dlss_model_preset, + "worker_logs": session.worker_logs, + "completed_frames": session.completed_frames, + }) + for l in session.worker_logs or []: + log.debug(f'DLSSNeuralRenderer worker: {l}') + session.close() + except Exception as e: + log.error(f'DLSSNeuralRenderer: exception {e}') + if session is not None and not session.closed: + session.abort() + raise + self.diagnostics = { + "gpu": dict(gpu), + "runtime_bundle": prepared.runtime_bundle, + "sessions": session_diagnostics, + } + except StandaloneError: + raise + except Exception as exc: + log.error(f'DLSSNeuralRenderer: unexpected exception {exc}') + raise StandaloneError("processing_failed", f"DLSS Neural Rendering failed: {exc}") from exc + result = np.ascontiguousarray(np.stack(outputs, axis=0)) + self.last_report = { + "input_shape": tuple(images.shape), + "output_shape": tuple(result.shape), + "completed_images": batch, + } + return result + + +__all__ = ["DLSSNeuralRenderer", "RenderOptions"] diff --git a/scripts/dlss/supersample.py b/scripts/dlss/supersample.py new file mode 100644 index 000000000..cf874ed29 --- /dev/null +++ b/scripts/dlss/supersample.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from src.core.jobs import JobController, active_job +from src.upscale.image.models import ImageUpscaleOptions, output_size as source_output_size +from src.upscale.video.models import UpscaleOptions as NativeUpscaleOptions +from src.upscale.video.native import RTXVideoSession, probe_capabilities + +from .utils import StandaloneError, nchw_image_to_hwc, hwc_to_nchw, validate_nchw, srgb_to_worker, worker_to_srgb_rgb, log + + +@dataclass(frozen=True, slots=True) +class UpscaleOptions: + vsr_quality: int = 4 + size_mode: str = "Scale factor" + scale_factor: float = 2.0 + width: int = 3840 + height: int = 2160 + aspect_lock: bool = True + ai_gpu_uuid: str = "auto" + + def source_options(self) -> ImageUpscaleOptions: + return ImageUpscaleOptions( + vsr_quality=self.vsr_quality, + size_mode=self.size_mode, + scale_factor=self.scale_factor, + width=self.width, + height=self.height, + aspect_lock=self.aspect_lock, + ai_gpu_uuid=self.ai_gpu_uuid, + ) + + def validate(self) -> None: + source = self.source_options() + source.validate() + + +class DLSSSuperSample: + """RGB NCHW adapter for the native RTX Video Super Resolution worker.""" + + def __init__(self) -> None: + log.info('DLSSSuperSample: init') + self.diagnostics: dict[str, Any] = {} + self.last_report: dict[str, Any] = {} + + def __call__( + self, + images: np.ndarray, + options: UpscaleOptions | None = None, + *, + controller: JobController | None = None, + ) -> np.ndarray: + log.info('DLSSSuperSample: call') + options = options or UpscaleOptions() + options.validate() + batch, _, height, width = validate_nchw(images, name="images") + source = options.source_options() + output_width, output_height = source_output_size(width, height, source) + own_controller = controller or JobController() + log.debug(f'DLSSSuperSample: controller={own_controller}') + outputs: list[np.ndarray] = [] + try: + with active_job(own_controller) as active_controller: + capabilities = probe_capabilities(options.ai_gpu_uuid, controller=active_controller) + log.debug(f'DLSSSuperSample: capabilities={capabilities}') + native_options = NativeUpscaleOptions( + vsr_enabled=True, + vsr_quality=int(options.vsr_quality), + ai_gpu_uuid=options.ai_gpu_uuid, + ) + native_options.validate(for_render=False) + with RTXVideoSession( + width, + height, + output_width, + output_height, + native_options, + 1, + capabilities, + active_controller, + ) as session: + log.debug(f'DLSSSuperSample: session={session}') + for index in range(batch): + if active_controller.cancel.is_set(): + raise StandaloneError("cancelled", "Upscale was cancelled.") + frame = nchw_image_to_hwc(images, index, name="images") + worker_input = srgb_to_worker(frame) + worker_output = session.process_frame(worker_input) + rgb = worker_to_srgb_rgb(worker_output, output_width, output_height) + log.debug(f'DLSSSuperSample: processed={rgb.shape}') + outputs.append(hwc_to_nchw(rgb, name="upscaled RGB output")[0]) + self.diagnostics = { + "gpu": dict(capabilities.gpu), + "sdk_version": capabilities.sdk_version, + "worker_version": capabilities.worker_version, + "completed_frames": session.completed_frames, + "last_results": session.last_results, + } + except StandaloneError: + raise + except Exception as exc: + log.error(f'DLSSSuperSample: unexpected exception {exc}') + raise StandaloneError("processing_failed", f"RTX Video upscaling failed: {exc}") from exc + result = np.ascontiguousarray(np.stack(outputs, axis=0)) + self.last_report = { + "input_shape": tuple(images.shape), + "output_shape": tuple(result.shape), + "completed_images": batch, + } + return result + + +__all__ = ["DLSSSuperSample", "UpscaleOptions"] diff --git a/scripts/dlss/utils.py b/scripts/dlss/utils.py new file mode 100644 index 000000000..19c749048 --- /dev/null +++ b/scripts/dlss/utils.py @@ -0,0 +1,118 @@ +from __future__ import annotations +import logging +import numpy as np + + +MAX_DIMENSION = 16_384 + +logging.getLogger().handlers.clear() +logging.basicConfig( + level=logging.DEBUG, + filename='dlss.log', + encoding='utf-8', + filemode='a', + format='%(asctime)s %(levelname)s %(message)s', + # datefmt='%Y-%m-%d %H:%M:%S-%f', + force=True, +) +log = logging.getLogger(__name__) +log.debug('DLSSInit') + + +class StandaloneError(RuntimeError): + """Base error with a stable machine-readable code.""" + + def __init__(self, code: str, message: str) -> None: + super().__init__(message) + self.code = code + self.message = message + + +class InvalidArrayError(StandaloneError): + def __init__(self, message: str) -> None: + super().__init__("invalid_array", message) + + +class VerificationError(StandaloneError): + def __init__(self, message: str) -> None: + super().__init__("verification_failed", message) + + +class ProcessingError(StandaloneError): + def __init__(self, message: str, *, code: str = "processing_failed") -> None: + super().__init__(code, message) + + +def validate_nchw(array: np.ndarray, *, name: str = "array") -> tuple[int, int, int, int]: + if not isinstance(array, np.ndarray): + raise InvalidArrayError(f"{name} must be a NumPy array.") + if array.ndim != 4: + raise InvalidArrayError(f"{name} must have shape (N, 3, H, W); got {array.shape}.") + batch, channels, height, width = array.shape + if batch < 1: + raise InvalidArrayError(f"{name} must contain at least one image.") + if channels != 3: + raise InvalidArrayError(f"{name} must contain RGB data with C=3; got C={channels}.") + if not 1 <= height <= MAX_DIMENSION or not 1 <= width <= MAX_DIMENSION: + raise InvalidArrayError( + f"{name} dimensions must be between 1 and {MAX_DIMENSION}; got {width}x{height}." + ) + if array.dtype != np.uint8: + raise InvalidArrayError(f"{name} must use dtype uint8; got {array.dtype}.") + return batch, channels, height, width + + +def copy_nchw(array: np.ndarray, *, name: str = "array") -> np.ndarray: + validate_nchw(array, name=name) + return np.ascontiguousarray(array.copy()) + + +def nchw_image_to_hwc(array: np.ndarray, index: int = 0, *, name: str = "array") -> np.ndarray: + batch, _, _, _ = validate_nchw(array, name=name) + if not 0 <= index < batch: + raise InvalidArrayError(f"{name} image index {index} is outside batch size {batch}.") + return np.ascontiguousarray(array[index].transpose(1, 2, 0)) + + +def hwc_to_nchw(array: np.ndarray, *, name: str = "image") -> np.ndarray: + if not isinstance(array, np.ndarray) or array.ndim != 3 or array.shape[2] != 3: + raise InvalidArrayError(f"{name} must have HWC RGB shape (H, W, 3); got {getattr(array, 'shape', None)}.") + if array.dtype != np.uint8: + raise InvalidArrayError(f"{name} must use dtype uint8; got {array.dtype}.") + return np.ascontiguousarray(array.transpose(2, 0, 1)[None, ...]) + + +def rgb_to_rgba(array: np.ndarray) -> np.ndarray: + """Add opaque alpha only at the private native-worker boundary.""" + if array.ndim != 3 or array.shape[2] != 3 or array.dtype != np.uint8: + raise InvalidArrayError("Native RGB input must have HWC uint8 shape with three channels.") + result = np.empty((*array.shape[:2], 4), dtype=np.uint8) + result[..., :3] = array + result[..., 3] = 255 + return np.ascontiguousarray(result) + + +def rgba_to_rgb_nchw(array: np.ndarray) -> np.ndarray: + if array.ndim != 3 or array.shape[2] != 4 or array.dtype != np.uint8: + raise InvalidArrayError("Native RGBA output must have HWC uint8 shape with four channels.") + return hwc_to_nchw(np.ascontiguousarray(array[..., :3]), name="native RGB output") + + +def srgb_to_worker(rgb: np.ndarray) -> np.ndarray: + """Convert HWC sRGB RGB data to the RTX Video worker's gamma-2.2 RGBA data.""" + if rgb.ndim != 3 or rgb.shape[2] != 3 or rgb.dtype != np.uint8: + raise InvalidArrayError("RGB input must have HWC uint8 shape with three channels.") + rgba = rgb_to_rgba(rgb) + values = rgba[..., :3].astype(np.float32) / 255.0 + linear = np.where(values <= 0.04045, values / 12.92, ((values + 0.055) / 1.055) ** 2.4) + rgba[..., :3] = np.rint(np.clip(linear, 0.0, 1.0) ** (1.0 / 2.2) * 255.0).astype(np.uint8) # pylint: disable=unsupported-assignment-operation + return np.ascontiguousarray(rgba) + + +def worker_to_srgb_rgb(data: bytes | bytearray | memoryview, width: int, height: int) -> np.ndarray: + """Convert packed worker RGBA output to an HWC sRGB RGB array.""" + rgba = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 4).copy() + values = (rgba[..., :3].astype(np.float32) / 255.0) ** 2.2 + rgb = np.where(values <= 0.0031308, values * 12.92, 1.055 * values ** (1.0 / 2.4) - 0.055) + rgba[..., :3] = np.rint(np.clip(rgb, 0.0, 1.0) * 255.0).astype(np.uint8) + return np.ascontiguousarray(rgba[..., :3]) diff --git a/scripts/dlss/verify.py b/scripts/dlss/verify.py new file mode 100644 index 000000000..22ee46b12 --- /dev/null +++ b/scripts/dlss/verify.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import importlib.util +import platform +import sys +from dataclasses import dataclass, field +from typing import Any + +from src.core.gpu_selection import resolve_ai_gpu +from src.core.gpu_detection import detect_gpus +from src.core.paths import ADDON, DLSS_SUPERRES, FFMPEG, FFPROBE, HOST_DXGI, NEURAL_RUNTIME, WORKER +from src.core.runtime import validate_runtime_files + +from .utils import log + + +@dataclass(frozen=True, slots=True) +class VerifyOptions: + level: str = "basic" + check_neural: bool = True + check_upscale: bool = True + check_interpolation: bool = True + + def validate(self) -> None: + if self.level not in {"basic", "deep"}: + raise ValueError("Verification level must be 'basic' or 'deep'.") + for name in ("check_neural", "check_upscale", "check_interpolation"): + if not isinstance(getattr(self, name), bool): + raise ValueError(f"{name} must be a boolean.") + + +@dataclass(frozen=True, slots=True) +class VerificationCheck: + name: str + passed: bool + detail: str = "" + + def to_dict(self) -> dict[str, Any]: + log.debug(f'DLSSVerify: check="{self.name}" passed={self.passed} detail="{self.detail}"') + result = {"name": self.name, "passed": self.passed, "detail": self.detail} + return result + + +@dataclass(frozen=True, slots=True) +class VerificationReport: + ok: bool + level: str + python: str + platform: str + gpu: dict[str, Any] | None + paths: dict[str, str] + checks: tuple[VerificationCheck, ...] + diagnostics: tuple[str, ...] = field(default_factory=tuple) + + @property + def failed(self) -> tuple[VerificationCheck, ...]: + return tuple(check for check in self.checks if not check.passed) + + def to_dict(self) -> dict[str, Any]: + log.info(f'DLSSVerify: level="{self.level}" python="{self.python}" platform="{self.platform}"') + log.info(f'DLSSVerify: paths={self.paths}') + log.info(f'DLSSVerify: gpu={self.gpu}') + log.debug(f'DLSSVerify: diagnostics={self.diagnostics}') + return { + "ok": self.ok, + "python": self.python, + "platform": self.platform, + "gpu": self.gpu, + "paths": self.paths, + "checks": [check.to_dict() for check in self.checks], + "diagnostics": self.diagnostics, + } + + +_RUNTIME_PATHS = { + "ffmpeg": FFMPEG, + "ffprobe": FFPROBE, + "worker": WORKER, + "host_dxgi": HOST_DXGI, + "dlss_addon": ADDON, + "dlss_superres": DLSS_SUPERRES, + "dlss_neural": NEURAL_RUNTIME, +} + + +class DLSSVerify: + """Perform side-effect-free runtime preflight checks by default.""" + + def __init__(self) -> None: + self.last_report: VerificationReport | None = None + + def __call__(self, gpu_uuid: str = "auto", options: VerifyOptions | None = None) -> VerificationReport: + options = options or VerifyOptions() + options.validate() + checks: list[VerificationCheck] = [] + diagnostics: list[str] = [] + selected_gpu: dict[str, Any] | None = None + + checks.extend(self._check_files()) + checks.append(self._check_import("numpy")) + checks.append(self._check_import("PIL")) + checks.append(self._check_import("cv2")) + checks.append(self._check_import("av")) + + try: + gpus = detect_gpus() + selected_gpu = resolve_ai_gpu(gpus, gpu_uuid) + checks.append(VerificationCheck("gpu", True, self._gpu_detail(selected_gpu))) + except (OSError, RuntimeError, ValueError) as exc: + checks.append(VerificationCheck("gpu", False, str(exc))) + diagnostics.append(str(exc)) + + try: + validate_runtime_files() + checks.append(VerificationCheck("runtime", True, "Required runtime files are present.")) + except (OSError, RuntimeError, ValueError) as exc: + checks.append(VerificationCheck("runtime", False, str(exc))) + diagnostics.append(str(exc)) + + if options.level == "deep": + checks.extend(self._deep_checks(options, gpu_uuid, selected_gpu)) + else: + checks.append(VerificationCheck("deep_capabilities", True, "Deep capability checks were not requested.")) + + filtered = tuple(check for check in checks if self._feature_enabled(check.name, options)) + report = VerificationReport( + ok=all(check.passed or check.status == "not_run" for check in filtered), + level=options.level, + python=platform.python_version(), + platform=sys.platform, + gpu=selected_gpu, + paths={name: str(path) for name, path in _RUNTIME_PATHS.items()}, + checks=filtered, + diagnostics=tuple(diagnostics), + ) + self.last_report = report + return report + + @staticmethod + def _check_files() -> list[VerificationCheck]: + return [ + VerificationCheck( + name=f"file:{name}", + passed=path.is_file(), + detail=str(path), + ) + for name, path in _RUNTIME_PATHS.items() + ] + + @staticmethod + def _check_import(name: str) -> VerificationCheck: + available = importlib.util.find_spec(name) is not None + return VerificationCheck( + name=f"dependency:{name}", + passed=available, + detail="available" if available else "not installed", + ) + + @staticmethod + def _gpu_detail(gpu: dict[str, Any]) -> str: + return f"{gpu.get('name', 'NVIDIA GPU')} driver={gpu.get('driver', 'unknown')} uuid={gpu.get('uuid', 'unknown')}" + + @staticmethod + def _feature_enabled(name: str, options: VerifyOptions) -> bool: + if name.startswith("neural:"): + return options.check_neural + if name.startswith("upscale:"): + return options.check_upscale + if name.startswith("interpolation:"): + return options.check_interpolation + return True + + @staticmethod + def _deep_checks(options: VerifyOptions, gpu_uuid: str, selected_gpu: dict[str, Any] | None) -> list[VerificationCheck]: + del selected_gpu + checks: list[VerificationCheck] = [] + if options.check_neural: + try: + from src.core.runtime import prepare_runtime + + prepared = prepare_runtime() + checks.append(VerificationCheck("neural:runtime", True, f"Prepared {len(prepared.warmed_files)} runtime components.")) + except (ImportError, OSError, RuntimeError, ValueError) as exc: + checks.append(VerificationCheck("neural:runtime", False, str(exc))) + if options.check_upscale: + try: + from src.upscale.video.native import probe_capabilities + + capabilities = probe_capabilities(gpu_uuid) + checks.append(VerificationCheck("upscale:capability", bool(capabilities.vsr.get("available")), str(capabilities.vsr))) + except (ImportError, OSError, RuntimeError, ValueError) as exc: + checks.append(VerificationCheck("upscale:capability", False, str(exc))) + if options.check_interpolation: + try: + from src.frame_interpolation.capabilities import probe_frame_interpolation_capabilities + + capabilities = probe_frame_interpolation_capabilities(gpu_uuid) + checks.append(VerificationCheck("interpolation:capability", capabilities.available, capabilities.detail or f"native_multiplier={capabilities.native_multiplier}")) + except (ImportError, OSError, RuntimeError, ValueError) as exc: + checks.append(VerificationCheck("interpolation:capability", False, str(exc))) + return checks diff --git a/scripts/dlss_ext.py b/scripts/dlss_ext.py new file mode 100644 index 000000000..13c27fed5 --- /dev/null +++ b/scripts/dlss_ext.py @@ -0,0 +1,332 @@ +import os +import time +import textwrap +import gradio as gr +from modules.logger import log +from modules import scripts_manager, shared, devices, processing, timer, errors +from scripts.dlss import controller_cli as c + + +debug = os.environ.get('SD_DLSS_DEBUG', None) is not None +FPS_CHOICES = ['23.976', '24', '25', '29.97', '30', '50', '59.94', '60', '90', '119.88', '120', '144', '165', '180', '240', '360', '480'] + + +class DLSSScript(scripts_manager.Script): + def title(self): + return 'nVidia DLSS' + + def show(self, _is_img2img): + if devices.backend != 'cuda': + return False + return scripts_manager.AlwaysVisible + + def ui(self, _is_img2img): + with gr.Accordion('nVidia DLSS', open=False, elem_id='dlss'): + with gr.Row(): + pkg_path = gr.Textbox(label='DLSS package path', value=shared.opts.dlss_pkg_path, placeholder='path to dlss 5 visual enhancer', elem_id='dlss_pkg_path') + with gr.Row(): + """ + btn_verify = ui_components.ToolButton(value=ui_symbols.tools, elem_id='dlss_verify') + btn_status = ui_components.ToolButton(value=ui_symbols.info, elem_id='dlss_status_btn') + btn_reset = ui_components.ToolButton(value=ui_symbols.reset, elem_id='dlss_reset') + btn_shutdown = ui_components.ToolButton(value=ui_symbols.close, elem_id='dlss_shutdown') + """ + btn_install = gr.Button(value="Install", elem_id='dlss_install') + btn_verify = gr.Button(value="Verify", elem_id='dlss_verify') + btn_status = gr.Button(value="Status", elem_id='dlss_status_btn') + btn_reset = gr.Button(value="Reset", elem_id='dlss_reset') + btn_shutdown = gr.Button(value="Shutdown", elem_id='dlss_shutdown') + with gr.Row(): + install_note = gr.Markdown("", elem_id='dlss_install_note', visible=False) + with gr.Accordion('DLSS Status', open=True, elem_id='dlss_status'): + ss_status = gr.JSON({ 'Status': 'unknown' if len(shared.opts.dlss_pkg_path) < 4 else 'stored'}) + btn_install.click(self.install, inputs=[], outputs=[install_note]) + btn_verify.click(self.verify, inputs=[pkg_path], outputs=[ss_status]) + btn_status.click(self.status, inputs=[pkg_path], outputs=[ss_status]) + btn_reset.click(self.reset, inputs=[pkg_path], outputs=[ss_status]) + btn_shutdown.click(self.shutdown, inputs=[pkg_path], outputs=[ss_status]) + + with gr.Accordion('DLSS NeuralRender', open=True, elem_id='dlss_nn'): + with gr.Row(): + nr_enabled = gr.Checkbox(label='NR enable', value=False, elem_id='dlss_nr_enabled') + nr_append = gr.Checkbox(label='Append result', value=True, elem_id='dlss_nr_append') + with gr.Row(): + nr_style = gr.Dropdown(label='NR style', choices=['Default', 'Natural', 'Cinematic'], value='Default', elem_id='dlss_nr_style') + nr_intensity = gr.Slider(label='NR intensity', minimum=0.0, maximum=2.0, step=0.05, value=1.0, elem_id='dlss_nr_intensity') + nr_upscaling_factor = gr.Dropdown(label='NR upscaling factor', choices=["1.0", "1.5", "1.724", "2.0", "3.0"], value="1.0", elem_id='dlss_nr_upscaling_factor') + with gr.Row(): + nr_preset = gr.Dropdown(label='NR preset', choices=['Default', 'Preset #1', 'Preset #2', 'Preset #3'], value='Default', elem_id='dlss_nr_preset') + nr_model_preset = gr.Dropdown(label='NR model', choices=['Default', 'J', 'K', 'L', 'M'], value='Default', elem_id='dlss_nr_model_preset') + with gr.Row(): + nr_local_tone = gr.Slider(label='NR tone strength', minimum=0.0, maximum=2.0, step=0.05, value=1.0, elem_id='dlss_nr_local_tone') + nr_local_structure = gr.Slider(label='NR local structure', minimum=0.0, maximum=2.0, step=0.05, value=1.0, elem_id='dlss_nr_local_structure') + nr_skin_structure = gr.Slider(label='NR skin structure', minimum=-1.0, maximum=2.0, step=0.05, value=-1.0, elem_id='dlss_nr_skin_structure') + nr_automatic_mask = gr.Checkbox(label='Automatic mask', value=False, elem_id='dlss_nr_automatic_mask') + + with gr.Accordion('DLSS SuperSample', open=True, elem_id='dlss_ss'): + with gr.Row(): + ss_enabled = gr.Checkbox(label='SR enable', value=False, elem_id='dlss_ss_enabled') + ss_append = gr.Checkbox(label='Append result', value=True, elem_id='dlss_ss_append') + with gr.Row(): + ss_vsr_quality = gr.Dropdown(label='SS VSR quality', choices=["1: Low", "2: Medium", "3: High", "4: Ultra"], value="4: Ultra", type='value', elem_id='dlss_ss_vsr_quality') + with gr.Row(): + ss_size_mode = gr.Dropdown(label='SS size mode', choices=['Scale factor', 'Target size'], value='Scale factor', elem_id='dlss_ss_size_mode') + with gr.Row(): + ss_scale_factor = gr.Slider(label='SS scale factor', minimum=1.0, maximum=8.0, step=0.05, value=2.0, elem_id='dlss_ss_scale_factor') + with gr.Row(): + ss_width = gr.Number(label='SS width', minimum=64, maximum=16384, step=8, value=3840, elem_id='dlss_ss_width') + ss_height = gr.Number(label='SS height', minimum=64, maximum=16384, step=8, value=2160, elem_id='dlss_ss_height') + + with gr.Accordion('DLSS FrameGen', open=True, elem_id='dlss_fg'): + with gr.Row(): + fg_enabled = gr.Checkbox(label='FG enable', value=False, elem_id='dlss_fg_enabled') + with gr.Row(): + fg_source_fps = gr.Dropdown(label='Source FPS', choices=FPS_CHOICES, value='24', elem_id='dlss_fg_source_fps') + fg_target_fps = gr.Dropdown(label='Target FPS', choices=FPS_CHOICES, value='60', elem_id='dlss_fg_target_fps') + with gr.Row(): + fg_engine = gr.Dropdown(label='Engine', choices=['Auto', 'Native DLSSG', 'Cascade'], value='Auto', elem_id='dlss_fg_engine') + + return [nr_enabled, nr_append, nr_style, nr_intensity, nr_local_tone, nr_local_structure, nr_skin_structure, nr_upscaling_factor, nr_preset, nr_automatic_mask, nr_model_preset, ss_enabled, ss_append, ss_vsr_quality, ss_size_mode, ss_scale_factor, ss_width, ss_height, fg_enabled, fg_source_fps, fg_target_fps, fg_engine] + + def install(self): + note = textwrap.dedent("""\ + ### Install + 1. Download and unpack: [DLSS 5 Visual Enhancer](https://github.com/Merserk/dlss5-visual-enhancer/releases/tag/v7.0) + 2. Enter the path to the unpacked package + 3. Press verify + ### Notes + - Package info is stored for future use on sucessful verification + - DLSS controller process is started on first use + - Use status to check the current state of the DLSS controller + - Use reset to restore the DLSS controller to its default state + - Use shutdown to stop the DLSS controller process + """) + return gr.update(value=note, visible=True) + + def verify(self, pkg_path): + log.info(f'DLSS verify: path="{pkg_path}"') + if not os.path.exists(pkg_path) or not os.path.isdir(pkg_path): + log.error(f'DLSS: path="{pkg_path}" not found') + return { 'error': 'package path not found' } + if not c.controller.get_python(pkg_path): + return { 'error': 'python not found in package path' } + shared.opts.dlss_pkg_path = pkg_path + response = c.controller.call(pkg_path, 'verify', { 'gpu_uuid': 'auto', 'options': { 'level': 'deep' } }) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return { 'error': error.get('message', 'unknown error') } + report = (response.get('result') or {}).get('report', {}) + if debug: + log.trace(f'DLSS raw: {report}') + checks = { 'passed': 0, 'failed': 0 } + for check in report.get('checks', []): + if check.get('passed', False): + checks['passed'] += 1 + else: + checks['failed'] += 1 + log.error(f'DLSS : {check}') + log.debug(f'DLSS: gpu={report.get("gpu", "unknown")} checks={checks}') + return report + + def status(self, pkg_path): + log.info(f'DLSS status: path="{pkg_path}"') + response = c.controller.call(pkg_path, 'status', {}) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return { 'error': error.get('message', 'unknown error') } + return response.get('result', {}) + + def reset(self, pkg_path): + log.info(f'DLSS reset: path="{pkg_path}"') + response = c.controller.call(pkg_path, 'reset', {}) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return { 'error': error.get('message', 'unknown error') } + return response.get('result', {}) + + def shutdown(self, pkg_path): + log.info(f'DLSS shutdown: path="{pkg_path}"') + if not c.controller.is_alive(): + return { 'shutdown': True, 'note': 'controller was not running' } + c.controller.stop() + return { 'shutdown': True } + + def supersample(self, pkg_path, images, ss_vsr_quality, ss_size_mode, ss_scale_factor, ss_width, ss_height): + try: + options = { + 'vsr_quality': int(ss_vsr_quality[0]), + 'size_mode': ss_size_mode, + 'scale_factor': float(ss_scale_factor), + 'width': int(ss_width), + 'height': int(ss_height), + 'aspect_lock': False, + } + response = c.controller.call(pkg_path, 'upscale', { 'images': c.images_to_nchw(images), 'options': options }) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return None + return c.nchw_to_images(response.get('result')) + except Exception as e: + log.error(f'DLSS: {e}') + errors.display(e, 'DLSS') + return None + + def neuralrender(self, pkg_path, images, nr_style, nr_intensity, nr_local_tone, nr_local_structure, nr_skin_structure, nr_upscaling_factor, nr_preset, nr_automatic_mask, nr_model_preset): + try: + options = { + 'nr_style': nr_style, + 'nr_intensity': float(nr_intensity), + 'local_tone_strength': float(nr_local_tone), + 'local_structure_strength': float(nr_local_structure), + 'skin_structure_strength': float(nr_skin_structure), + 'upscaling_factor': float(nr_upscaling_factor), + 'warmup_frames': 0, + 'nr_preset': nr_preset, + 'automatic_mask': bool(nr_automatic_mask), + 'dlss_model_preset': nr_model_preset, + } + response = c.controller.call(pkg_path, 'render', { 'images': c.images_to_nchw(images), 'options': options }) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return None + return c.nchw_to_images(response.get('result')) + except Exception as e: + log.error(f'DLSS: {e}') + errors.display(e, 'DLSS') + return None + + def framegen(self, pkg_path, images, fg_source_fps, fg_target_fps, fg_engine): + try: + if len(images) < 2: + log.warning('DLSS: FrameGen requires at least two frames, skipping') + return None + options = { 'ai_gpu_uuid': 'auto', 'engine': fg_engine } + response = c.controller.call( + pkg_path, 'framegen', + { 'frames': c.images_to_nchw(images), 'source_fps': fg_source_fps, 'target_fps': fg_target_fps, 'options': options }, + timeout=120.0, + ) + if response.get('status') != 'ok': + error = response.get('error') or {} + log.error(f'DLSS: {error.get("message")}') + return None + return c.nchw_to_images(response.get('result')) + except Exception as e: + log.error(f'DLSS: {e}') + errors.display(e, 'DLSS') + return None + + def dlss(self, p, pp, + nr_enabled, nr_append, nr_style, nr_intensity, nr_local_tone, nr_local_structure, nr_skin_structure, nr_upscaling_factor, nr_preset,nr_automatic_mask, nr_model_preset, + ss_enabled, ss_append, ss_vsr_quality, ss_size_mode, ss_scale_factor, ss_width, ss_height, + fg_enabled, fg_source_fps, fg_target_fps, fg_engine, + *args, **kwargs + ): + if not (ss_enabled or nr_enabled or fg_enabled): + return None + pkg_path = shared.opts.dlss_pkg_path + if not pkg_path or not c.controller.get_python(pkg_path): + log.error('DLSS: package path not configured') + return None + if debug: + log.trace(f'DLSS: path="{pkg_path}" args={args} kwargs={kwargs}') + + if hasattr(pp, 'images') and pp.images is not None and len(pp.images) > 0: + inputs = pp.images + elif hasattr(pp, 'image') and pp.image is not None: + inputs = [pp.image] + else: + return None + + # cast to appropriate types + nr_intensity = float(nr_intensity) + nr_local_tone = float(nr_local_tone) + nr_local_structure = float(nr_local_structure) + nr_skin_structure = float(nr_skin_structure) + nr_upscaling_factor = float(nr_upscaling_factor) + ss_width = int(ss_width) + ss_height = int(ss_height) + ss_scale_factor = float(ss_scale_factor) + fg_source_fps = float(fg_source_fps) + fg_target_fps = float(fg_target_fps) + + images = [] + originals = [] + current_images = inputs + t = timer.Timer() + + if ss_enabled: + t0 = time.time() + p.extra_generation_params["DLSSSuperSample"] = True + log.debug(f'DLSS: method=SuperSample quality="{ss_vsr_quality}" mode="{ss_size_mode}" scale={ss_scale_factor} width={ss_width} height={ss_height}') + if ss_append: + originals.extend(current_images) + output = self.supersample(pkg_path, current_images, ss_vsr_quality, ss_size_mode, ss_scale_factor, ss_width, ss_height) + if debug: + log.trace(f'DLSS: method=SuperSample images={len(output) if output else 0} time={time.time() - t0:.3f}') + if output: + images.extend(output) + current_images = output + t.ts('supersample', t0) + + if nr_enabled: + t0 = time.time() + p.extra_generation_params["DLSSNeuralRender"] = True + log.debug(f'DLSS: method=NeuralRender style={nr_style} intensity={nr_intensity} tone={nr_local_tone} structure={nr_local_structure} skin={nr_skin_structure} scale={nr_upscaling_factor} preset={nr_preset} mask={nr_automatic_mask} model={nr_model_preset}') + if nr_append: + originals.extend(current_images) + output = self.neuralrender(pkg_path, current_images, nr_style, nr_intensity, nr_local_tone, nr_local_structure, nr_skin_structure, nr_upscaling_factor, nr_preset, nr_automatic_mask, nr_model_preset) + if debug: + log.trace(f'DLSS: method=NeuralRender images={len(output) if output else 0} time={time.time() - t0:.3f}') + if output: + images.extend(output) + current_images = output + t.ts('neuralrender', t0) + + if fg_enabled: + t0 = time.time() + p.extra_generation_params["DLSSFrameGen"] = True + log.debug(f'DLSS: method=FrameGen source={fg_source_fps} target={fg_target_fps} engine={fg_engine}') + output = self.framegen(pkg_path, current_images, fg_source_fps, fg_target_fps, fg_engine) + if debug: + log.trace(f'DLSS: method=FrameGen images={len(output) if output else 0} time={time.time() - t0:.3f}') + if output: + images.extend(output) + current_images = output + t.ts('framegen', t0) + + log.debug(f'DLSS: images={len(images)} {t.summary(min_time=0)}') + pp.images = images + pp.originals = originals + return pp + + def postprocess_image(self, p: processing.StableDiffusionProcessing, pp: scripts_manager.PostprocessImageArgs, *args, **kwargs): + # postprocess_image is intended to modify single image in-place so not suited for dlss + pass + + def postprocess(self, p: processing.StableDiffusionProcessing, pp: processing.Processed, *args, **kwargs): # pylint: disable=arguments-differ,unused-argument + _pp = self.dlss(p, pp, *args, **kwargs) + # postprocess triggers after initial images have already been saved + if _pp is not None and hasattr(_pp, 'images') and _pp.images is not None: + pp = _pp + orig_infos = pp.infotexts if hasattr(pp, 'infotexts') else [] + out_images, out_infos = processing.process_samples(p, pp.images) + pp.images = out_images + pp.infotexts = out_infos + if hasattr(pp, 'originals') and pp.originals is not None and len(pp.originals) > 0: + pp.infotexts = orig_infos + pp.infotexts + pp.images = pp.originals + pp.images + return pp + + +""" +add install notes +add postprocessing +add framegen +add video +"""