Skip to content
Justin SilverTECHNOLOGY, TRAVEL & EVERYTHING IN BETWEEN
Libraries & Tools
Local AI image editingRustCPUGPU

OSOR-rs

OSOR-rs is a native Rust inference implementation of the upstream OSOR SDXL-Inpainting Phase II model. OSOR is short for One-Step Object Removal. The library accepts an RGB image, a removal mask, and explicit noise tensors, then returns the completed RGB image and the model's predicted edit alpha. Model weights load from local safetensors bundles. The Rust runtime has no Python sidecar or automatic model download.

The Cargo package and Rust import are named osor. The repository remains doublesharp/osor-rs. Candle provides CPU, Metal, and CUDA execution; an optional Burn backend provides WGPU execution.

Install the crate

Use the repository's pinned Rust toolchain. Add the Git dependency and commit Cargo.lock to retain the resolved revision. Enable the optional metal feature for a macOS GPU build; cuda and wgpu enable the other GPU implementations.

[dependencies.osor]
git = "https://github.com/doublesharp/osor-rs"
branch = "main"
# features = ["metal"]

Load a verified local bundle

A prepared bundle contains manifest.json, unet.safetensors, vae.safetensors, and prompt.safetensors. It includes the Phase II LoRA adapters merged into the UNet, the trained five-channel output head, the VAE, and cached fixed prompt embeddings. The offline Rust converter validates critical tensor names and shapes and records source revisions and SHA-256 hashes.

The builder selects a backend and can pin a trusted manifest digest retained during installation or conversion. Loading verifies that manifest before opening weight tensors, then verifies the tensor bytes against it. Retain the loaded model across jobs to reuse the weights.

use osor::{Backend, Osor, Result};

fn load_model(bundle: &str, manifest_sha256: &str) -> Result<Osor> {
    Osor::builder(bundle)
        .backend(Backend::Cpu)
        .trusted_manifest_sha256(manifest_sha256)
        .load()
}

Supply image, mask, and noise tensors

Inputs use NCHW layout on model.device(). The image is F32 RGB in the range zero to one. The mask is F32 in the same range, with white marking removal. Supply independent standard-normal tensors for VAE posterior sampling and diffusion noise, using the model dtype. WGPU uses CPU staging tensors for input and output even though its neural graph runs on the GPU; compute_device_description() identifies its compute adapter.

image:           [B, 3, H, W]
mask:            [B, 1, H, W]
posterior_noise: [B, 4, H/8, W/8]
diffusion_noise: [B, 4, H/8, W/8]

Height and width must be multiples of eight between 32 and 1024. CPU, Metal, and CUDA admit batches up to four. WGPU admits batch one with at most 262,144 pixels. Read model.inference_limits() for the effective limits of the selected backend.

model.seeded_noise(seed, batch, height, width) generates both noise tensors using a versioned, platform-independent algorithm. Retain explicit noise tensors when comparing different frameworks; a matching seed alone does not establish matching input noise.

use osor::{InferenceInput, InferenceResult, Osor, Result, Tensor};

fn remove_object(
    model: &Osor,
    image: &Tensor,
    mask: &Tensor,
    posterior_noise: &Tensor,
    diffusion_noise: &Tensor,
) -> Result<InferenceResult> {
    model.infer(InferenceInput {
        image,
        mask,
        posterior_noise,
        diffusion_noise,
    })
}

Use RGB and the predicted edit region

InferenceResult contains image with shape [B, 3, H, W], alpha with shape [B, 1, H, W], and latent_alpha with shape [B, 1, H/8, W/8]. The inference recipe blends denoised and original latents using the predicted alpha before decoding. The returned RGB already contains that blend. The alpha outputs describe edit support; they are not instructions to multiply the decoded image by alpha a second time.

That region can extend beyond the painted mask to include associated shadows and reflections. Compositing the result strictly through the original mask can restore effects that the model removed. The host application must choose how the predicted region interacts with selections and protected pixels.

Resize, pad, and restore dimensions

preprocess::resize_reference_inputs applies the reference dimension policy, bicubic RGB resizing, and nearest-neighbor mask resizing. After inference, reference_output_rgb8 restores the original dimensions from quantized output and optionally composites with predicted alpha. Tensor callers can use pad_to_model_size and Padding::crop_result to pad and crop instead of resampling.

Track jobs and bound resource use

infer runs synchronously. infer_with_observer reports 18 typed stages, infer_with_cancellation checks a shared cancellation token, and infer_profiled records synchronized stage timings. Cancellation is cooperative, with checkpoints between model operations, attention slices, and convolution bands.

ArtifactLimits bounds decoded weight admission. InferenceLimits bounds job geometry. These are separate budgets: a model can fit its weight budget while an inference job still needs substantial activation memory. The host application owns memory admission, background scheduling, model installation, and document history.

Code, weights, and integration

Original project code is MIT licensed, with preserved MIT and MIT-CMU notices for adapted components. Model weights have separate terms. This implementation uses the SDXL variant, whose checkpoints use CreativeML Open RAIL++-M. The repository documents bundle provenance, backend setup, numerical comparisons, and model licensing.

An in-progress PhotoCraft integration uses the library through its local model service. The service owns installation and reuse; the editor owns colour conversion, crop selection, preview, and apply or undo. The journal post explains why I started this work for the Remove and Smart Fill tools, including the different models those tools currently use.

The PhotoCraft story

Read why I built this library while working on local AI support for Remove and Smart Fill.

Read the journal post