Inference Engine Model Slots and Tiling

Inference Engine, Model Slots and Tiling

Relevant source files

The following files were used as context for generating this wiki page:

Purpose and Scope

This page covers the implementation details of the neural super-resolution inference engine in clarity/processing/engine.py. The inference subsystem provides GPU-accelerated upscaling using ESRGAN-family models loaded via the spandrel library, with support for FP16 precision, custom model slots, automatic tiled inference, batched tile execution, CUDA Out-Of-Memory (OOM) recovery, and a Lanczos CPU fallback path when PyTorch or model weights are absent clarity/processing/engine.py:1-16.

Model Registry and Slots

The engine maps logical texture processing roles to specific weight files using a registry system. Two core registries are defined: DEFAULT_REGISTRY (mirroring classic ChaiNNer pipelines) and RECOMMENDED_REGISTRY (optimized 2024–2025 picks such as RealPLKSR-DySample and SPAN models) clarity/processing/engine.py:25-47.

Registry Definitions

Slot Name	Default Model File	Recommended Model File	Role Description
`bc1clean`	`1x_BC1-smooth2.pth`	`None`	Removes BC1 compression artifacts (unset for PBRify models) clarity/processing/engine.py:25-38
`normal`	`4x-Normal-RG0-BC7.pth`	`4x-Normal-RG0-BC7.pth`	Tangent space normals with the blue/z channel zeroed clarity/processing/engine.py:27-39
`normal_bc1`	`4x-Normal-RG0-BC1.pth`	`4x-Normal-RG0-BC1.pth`	BC1-encoded normal sources clarity/processing/engine.py:28-40
`color`	`4x_scalenx_90k.pth`	`4x-PBRify_RPLKSRd_V3.pth`	General diffuse, specular, and base color textures clarity/processing/engine.py:29-41
`mask`	`None`	`4x-PBRify_UpscalerSPANV4.pth`	Scalar channels (falls back to `color` if unset) clarity/processing/engine.py:30-42
`face`	`4xFaceUpDAT.pth`	`4xFaceUpDAT.pth`	Character face base color textures clarity/processing/engine.py:31-43
`skin`	`x1_ITF_SkinDiffDDS_v1.pth`	`x1_ITF_SkinDiffDDS_v1.pth`	Skin de-artifacting filter clarity/processing/engine.py:32-44
`hair`	`4x_UltraFArt_v3.pth`	`4x_UltraFArt_v3.pth`	Hair textures (typically covered by specialized mods) clarity/processing/engine.py:33-45
`ui`	`4x_foolhardy_Remacri.pth`	`4x-UltraSharpV2.safetensors`	UI sheets and game icons clarity/processing/engine.py:34-46

Configuration resolution checks for a local registry.json file beside the model weights directory, falling back to the tracked package registry configuration at paths.REGISTRY clarity/processing/engine.py:65-71. A slot mapped to None is deliberately disabled, bypassing fallback routines clarity/processing/engine.py:100-103. Slot aliases are resolved dynamically (e.g., normal_bc1 falls back to normal, and mask falls back to color when unconfigured) clarity/processing/engine.py:104-111.

graph TD
    A["Engine.__init__"] --> B["Read local registry.json"]
    B --> C["Read paths.REGISTRY"]
    C --> D["Merge into self.registry"]
    D --> E{"Slot requested"}
    E -->|normal_bc1 missing| F["Resolve to normal"]
    E -->|mask missing| G["Resolve to color"]
    E -->|Enabled| H["ModelLoader.load_from_file"]

Sources: clarity/processing/engine.py:25-111, clarity/paths.py [Implicit]

Engine Initialization and Spandrel Loading

The Engine class manages device allocation, precision settings, and model caching clarity/processing/engine.py:50-89.

class Engine:
    def __init__(
        self,
        models_dir,
        device=None,
        tile=512,
        pad=16,
        fp16=True,
        allow_fallback=True,
        tile_batch=4,
    ):
        ...

Sources: clarity/processing/engine.py:50-62

Key Initialization Attributes

models_dir: Directory containing .pth or .safetensors weight files clarity/processing/engine.py:50-63.
device: Automatically selects "cuda" if torch.cuda.is_available() evaluates to true, otherwise defaults to "cpu" clarity/processing/engine.py:82-88.
fp16: Enables half-precision (torch.half()) inference on CUDA devices for models supporting half-precision (m.supports_half) clarity/processing/engine.py:124-126.
tile_batch: Maximum number of tiles processed concurrently in a single forward pass clarity/processing/engine.py:73-76.

Models are lazily loaded using spandrel.ModelLoader, moved to the target device, set to evaluation mode (.eval()), and cached in self._models clarity/processing/engine.py:117-128.

Inference and Lanczos CPU Fallback

When an image is passed to Engine.run(slot, img, scale) or Engine.run_batch(slot, imgs, scale), the engine checks model availability clarity/processing/engine.py:130-147.

If the requested model or PyTorch is missing:

If allow_fallback is False, a RuntimeError is raised clarity/processing/engine.py:138-139.
If allow_fallback is True, the slot name is added to self.missing, and execution falls back to a CPU-based Lanczos scaling implementation (lanczos(img, scale)) clarity/processing/engine.py:140-141.

graph TD
    A["Engine.run(slot, img, scale)"] --> B{"Model available & Torch loaded?"}
    B -->|Yes| C["_run_model(slot, img, scale)"]
    B -->|No| D{"allow_fallback?"}
    D -->|False| E["Raise RuntimeError"]
    D -->|True| F["Record missing slot"]
    F --> G["lanczos(img, scale)"]

Sources: clarity/processing/engine.py:130-142

Tiling and Batched Inference

To process large textures without exhausting VRAM, Engine splits incoming images into overlapping patches (tiles) using configurable dimensions (tile) and boundary padding (pad) clarity/processing/engine.py:53-62.

Tiling Execution Flow (`_tiled`)

Dimensions & Padding: The image tensor is padded to handle boundary artifacts. If the source dimensions are smaller than tile + 2 * pad, tiling is bypassed entirely tests/test_engine_tiling.py:138-144.
Tile Extraction: Slices are extracted with an overlap equal to pad to prevent boundary seam artifacts tests/test_engine_tiling.py:51-65.
Batching: Extracted tiles from one or multiple images are grouped into batches up to tile_batch size and pushed through a single model forward pass clarity/processing/engine.py:73-76, tests/test_engine_tiling.py:81-86.
Reassembly: Output tiles are cropped to remove padding margins and stitched back into the full-resolution target tensor tests/test_engine_tiling.py:51-65.

CUDA OOM Recovery

During execution, GPU memory exhaustion can trigger a RuntimeError containing CUDA OOM messages tests/test_engine_tiling.py:37-48. The engine handles this robustly without crashing the entire batch processing pipeline:

Catch OOM: The tiling loop catches CUDA out-of-memory exceptions during model forward passes tests/test_engine_tiling.py:115-127.
Cache Clearing: torch.cuda.empty_cache() is invoked to reclaim unreferenced memory tests/test_engine_tiling.py:119.
Dynamic Batch Reduction: self.tile_batch is halved in-place (bounded at a minimum of 1) based on the batch size that triggered the failure clarity/processing/engine.py:73-76, tests/test_engine_tiling.py:122-126.
Retry: The failed tile group is re-processed with the smaller batch size. Once lowered, self.tile_batch remains reduced for the remainder of the session clarity/processing/engine.py:73-76.

sequenceDiagram
    participant Engine as Engine._tiled
    participant Model as Spandrel Model
    participant CUDA as torch.cuda

    Engine->>Model: Forward pass (tile_batch = N)
    Model--xEngine: RuntimeError: CUDA out of memory
    Engine->>CUDA: empty_cache()
    Note over Engine: Halve tile_batch: N = max(1, N // 2)
    Engine->>Model: Retry forward pass (tile_batch = N // 2)
    Model-->>Engine: Success (Tensor output)

Sources: clarity/processing/engine.py:73-76, tests/test_engine_tiling.py:115-127