From 64b281e4afb5bcf1ea9a2034ac6f2e8711d25311 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 9 Sep 2026 13:32:49 +0200 Subject: [PATCH] add dlss to postprocessing Signed-off-by: Vladimir Mandic --- extensions-builtin/sdnext-modernui | 2 +- scripts/dlss_ext.py | 633 ++++++++++++++++------------- 2 files changed, 343 insertions(+), 292 deletions(-) diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 2840cb952..b3e6c8f75 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 2840cb952665d416768917c780d1acd9175d09cf +Subproject commit b3e6c8f75cf65b19a3c3309c2fb6a69336471625 diff --git a/scripts/dlss_ext.py b/scripts/dlss_ext.py index 13c27fed5..94a6490da 100644 --- a/scripts/dlss_ext.py +++ b/scripts/dlss_ext.py @@ -3,7 +3,7 @@ import time import textwrap import gradio as gr from modules.logger import log -from modules import scripts_manager, shared, devices, processing, timer, errors +from modules import shared, devices, processing, timer, errors, scripts_manager, scripts_postprocessing from scripts.dlss import controller_cli as c @@ -11,6 +11,301 @@ 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'] +def create_ui(parent): + with gr.Accordion('nVidia DLSS', open=False, elem_id=f'{parent}_dlss_accordion'): + with gr.Row(): + 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 NeuralRender', open=False, 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=False, 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=False, 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') + + 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'}) + + 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') + + btn_install.click(install, inputs=[], outputs=[install_note]) + btn_verify.click(verify, inputs=[pkg_path], outputs=[ss_status]) + btn_status.click(status, inputs=[pkg_path], outputs=[ss_status]) + btn_reset.click(reset, inputs=[pkg_path], outputs=[ss_status]) + btn_shutdown.click(shutdown, inputs=[pkg_path], outputs=[ss_status]) + + 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(): + 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(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(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(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(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(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(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(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(p: processing.StableDiffusionProcessing | None, pp: processing.Processed | scripts_postprocessing.PostprocessedImage, + 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() + if p: + 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 = 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() + if p: + 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 = 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() + if p: + p.extra_generation_params["DLSSFrameGen"] = True + log.debug(f'DLSS: method=FrameGen source={fg_source_fps} target={fg_target_fps} engine={fg_engine}') + output = 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 + + class DLSSScript(scripts_manager.Script): def title(self): return 'nVidia DLSS' @@ -21,296 +316,14 @@ class DLSSScript(scripts_manager.Script): 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 + return create_ui(self.parent) 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) + _pp = 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 @@ -324,9 +337,47 @@ class DLSSScript(scripts_manager.Script): return pp -""" -add install notes -add postprocessing -add framegen -add video -""" +class DLSSPostprocessingScript(scripts_postprocessing.ScriptPostprocessing): + name = "nVidia DLSS" + order = 30000 + + def ui(self): + 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 = create_ui('postprocess') + return { + "nr_enabled": nr_enabled, + "nr_append": nr_append, + "nr_style": nr_style, + "nr_intensity": nr_intensity, + "nr_local_tone": nr_local_tone, + "nr_local_structure": nr_local_structure, + "nr_skin_structure": nr_skin_structure, + "nr_upscaling_factor": nr_upscaling_factor, + "nr_preset": nr_preset, + "nr_automatic_mask": nr_automatic_mask, + "nr_model_preset": nr_model_preset, + "ss_enabled": ss_enabled, + "ss_append": ss_append, + "ss_vsr_quality": ss_vsr_quality, + "ss_size_mode": ss_size_mode, + "ss_scale_factor": ss_scale_factor, + "ss_width": ss_width, + "ss_height": ss_height, + "fg_enabled": fg_enabled, + "fg_source_fps": fg_source_fps, + "fg_target_fps": fg_target_fps, + "fg_engine": fg_engine, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, *args, **kwargs): + nr_enabled = kwargs.get("nr_enabled", False) + ss_enabled = kwargs.get("ss_enabled", False) + fg_enabled = kwargs.get("fg_enabled", False) + if not (nr_enabled or ss_enabled or fg_enabled): + return + if pp.image is None: + return + result = dlss(None, pp, *args, **kwargs) + if result is None or not hasattr(result, "images") or len(result.images) == 0: + return + pp.image = result.images[0] + pp.info["DLSS"] = f'NR: {nr_enabled} SS: {ss_enabled} FG: {fg_enabled}'