mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
Vladimir Mandic
@@ -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"]
|
||||
Reference in New Issue
Block a user