Tutorial 3: Inpainting

Inpainting fills missing or corrupted regions of a medical image while keeping the known region fixed. MedLatents supports inpainting for every discrete architecture and provides both an end-to-end helper and per-architecture entry points. The relevant API lives in Inference API.

The workflow

Encode the known image to latent tokens.
Build a mask marking which tokens are unknown.
Resample only the masked tokens, conditioned on the known ones.
Decode back to image space.

The inpaint_volume() helper bundles all four steps; the other medlatents.inference functions expose them individually for token-level control.

Creating a mask

Use create_spatial_mask() to mark the region to fill. Masks are defined in image space and converted to token space internally.

from medlatents.inference import create_spatial_mask

# Block region in a single 2D slice
mask = create_spatial_mask(
    shape=(1, 256, 256),
    mask_type="block",
    start=(64, 64),
    end=(192, 192),
)

End-to-end inpainting

With a loaded generator, inpaint a masked region in one call using inpaint_volume(), which encodes, resamples the masked tokens, and decodes for you:

from medlatents.generation import DiscreteLatentGenerator
from medlatents.inference import inpaint_volume

generator = DiscreteLatentGenerator.from_checkpoints(
    model_type="maskgit",
    model_path="checkpoints/model.pt",
    tokenizer_path="tokenizer.pt",
    device="cuda",
)

inpainted = inpaint_volume(
    model=generator.model,
    tokenizer=generator.tokenizer,
    volume=corrupted_image,
    mask=mask,
    model_type="maskgit",
    num_steps=12,
)

Per-architecture entry points

When you are working directly with a model and tokens, choose the function that matches your architecture:

inpaint_maskgit() – confidence-guided parallel refilling for MaskGIT.
inpaint_autoregressive() – causal refilling for autoregressive models.
inpaint_flow_matching() – flow-based inpainting.
inpaint_diffusion_repaint() – RePaint-style resampling for D3PM.
inpaint_bayesian_flow() – inpainting for Bayesian flow models.

For 3D data, inpaint_volume() orchestrates masking and decoding over a whole volume.

Super-resolution

The same conditional-inference machinery upscales low-resolution scans. super_resolve_slices() and anisotropic_super_resolution() are the main entry points, with compare_with_interpolation() for a baseline comparison.

from medlatents.inference import anisotropic_super_resolution

# target_shape is the absolute (D, H, W) to upscale to; here we 4x the
# through-plane axis of a (1, 1, 64, 256, 256) volume and keep the rest.
high_res = anisotropic_super_resolution(
    model=generator.model,
    tokenizer=generator.tokenizer,
    volume=low_res_volume,
    target_shape=(256, 256, 256),
)

Working with large volumes

For volumes that do not fit in memory, tile encode/decode with encode_large_volume(), decode_large_volume(), and reconstruct_large_volume().

Preserving metadata

Inpainting and super-resolution operate on latents, but downstream reconstruction fidelity depends on orientation and spacing. Carry the affine / spacing metadata attached at load time through your pipeline – see Token Interface for the metadata contract.

Next steps

Tutorial 4: Evaluating Models – quantify reconstruction and generation quality.