Inference API

Conditional-inference utilities for discrete latent models: inpainting, super-resolution, masking helpers, large-volume tiling, and deployment-oriented optimizations (KV-cache, quantization, pruning, ONNX export).

Masking and Spatial Utilities

medlatents.inference.create_spatial_mask(shape, mask_type, mask_ratio=0.5, **mask_kwargs)

Create spatial masks for inpainting/super-resolution. True = keep, False = mask.

Parameters:

• shape (tuple[int, ...])

• mask_type (Literal['random', 'block', 'slice', 'checkerboard'])

• mask_ratio (float, default: 0.5)

Return type:

Tensor

medlatents.inference.spatial_to_sequence_mask(spatial_mask, rasterization_method='hilbert', **raster_kwargs)

Convert spatial mask to sequence mask using rasterization.

Parameters:

• spatial_mask (Tensor)

• rasterization_method (Literal['raster', 'hilbert', 'zorder'], default: 'hilbert')

Return type:

Tensor

medlatents.inference.apply_token_mask(tokens, mask, mask_value)

Apply mask to token sequence. True = keep, False = mask.

Accepts mask as either a 1-D (seq,) tensor (broadcast across batch) or a 2-D (batch, seq) tensor (per-sample mask). Both are common in practice — a single shared mask is natural for centre-mask inpainting of a fixed region, and per-sample masks are needed for anomaly-detection pipelines where each image has the anomaly at a different location.

Parameters:

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• mask_value (int)

Return type:

Int[Tensor, 'batch seq']

medlatents.inference.confidence_mask(logits, threshold)

Create mask based on prediction confidence. True = high confidence, False = low.

Parameters:

• logits (Float[Tensor, 'batch seq vocab'])

• threshold (float)

Return type:

Tensor

medlatents.inference.interpolate_spatial_upsampling(volume, target_shape, mode='trilinear')

Upsample volume using interpolation.

Parameters:

• volume (Tensor) – 4D (batch, channels, H, W) or 5D (batch, channels, D, H, W) tensor

• target_shape (tuple[int, ...]) – Target spatial dimensions

• mode (str, default: 'trilinear') – Interpolation mode (‘bilinear’ for 2D, ‘trilinear’ for 3D)

Return type:

Tensor

medlatents.inference.reslice_volume(volume, axis, target_slices, mode='trilinear')

Reslice volume along specific axis to target resolution.

Parameters:

• volume (Tensor)

• axis (int)

• target_slices (int)

• mode (str, default: 'trilinear')

Return type:

Tensor

Inpainting

End-to-end inpainting plus per-architecture entry points.

medlatents.inference.inpaint_volume(model, tokenizer, volume, mask, model_type, rasterization_method='hilbert', **inpaint_kwargs)

End-to-end volume inpainting: volume → tokens → inpaint → volume.

Parameters:

• model (Module)

• volume (Tensor)

• mask (Tensor)

• model_type (Literal['autoreg', 'maskgit', 'flow', 'diffusion', 'bayesian_flow'])

• rasterization_method (str, default: 'hilbert')

Return type:

Tensor

medlatents.inference.inpaint_autoregressive(model, tokens, mask, mask_token, num_iterations=5, temperature=1.0, top_p=0.95, confidence_threshold=0.1)

Inpaint using autoregressive model with iterative refinement.

Parameters:

• model (Module)

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• mask_token (int)

• num_iterations (int, default: 5)

• temperature (float, default: 1.0)

• top_p (float | None, default: 0.95)

• confidence_threshold (float, default: 0.1)

Return type:

Int[Tensor, 'batch seq']

medlatents.inference.inpaint_maskgit(model, tokens, mask, mask_token, num_steps=12, temperature=1.0, schedule='cosine')

Inpaint using MaskGIT with confidence-based scheduling.

mask is True at positions to keep (visible) and False at positions to inpaint. Accepts either a 1-D (seq,) mask (shared across the batch, natural for centre-mask demos) or a 2-D (batch, seq) mask (per-sample, natural for anomaly-detection pipelines where each image’s anomaly is at a different location).

Parameters:

• model (Module)

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• mask_token (int)

• num_steps (int, default: 12)

• temperature (float, default: 1.0)

• schedule (Literal['cosine', 'linear', 'sqrt'], default: 'cosine')

Return type:

Int[Tensor, 'batch seq']

medlatents.inference.inpaint_flow_matching(model, tokens, mask, mask_token, num_steps=100, temperature=1.0, path_type='polynomial', source_dist_type='uniform', vocab_size=1024)

Inpaint using flow matching with conditional generation.

Parameters:

• model (Module)

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• mask_token (int)

• num_steps (int, default: 100)

• temperature (float, default: 1.0)

• path_type (str, default: 'polynomial')

• source_dist_type (str, default: 'uniform')

• vocab_size (int, default: 1024)

Return type:

Int[Tensor, 'batch seq']

medlatents.inference.inpaint_diffusion_repaint(model, diffusion, tokens, mask, num_steps=50, jump_length=10, jump_n_sample=10)

Inpaint using RePaint algorithm for discrete diffusion (D3PM).

Parameters:

• model (Module)

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• num_steps (int, default: 50)

• jump_length (int, default: 10)

• jump_n_sample (int, default: 10)

Return type:

Int[Tensor, 'batch seq']

medlatents.inference.inpaint_bayesian_flow(model, tokens, mask, num_steps=100, temperature=1.0)

Inpaint using Bayesian Flow Network with conditional refinement.

Parameters:

• model (Module)

• tokens (Int[Tensor, 'batch seq'])

• mask (Tensor)

• num_steps (int, default: 100)

• temperature (float, default: 1.0)

Return type:

Int[Tensor, 'batch seq']

Super-Resolution

medlatents.inference.super_resolve_slices(model, tokenizer, volume, target_slices, axis=0, model_type='maskgit', interpolation_mode='trilinear', rasterization_method='hilbert', **inpaint_kwargs)

Super-resolve volume slices: upsample + inpaint interpolated slices.

Parameters:

• model (Module)

• volume (Tensor)

• target_slices (int)

• axis (int, default: 0)

• model_type (Literal['autoreg', 'maskgit', 'flow', 'diffusion', 'bayesian_flow'], default: 'maskgit')

• interpolation_mode (str, default: 'trilinear')

• rasterization_method (str, default: 'hilbert')

Return type:

Tensor

medlatents.inference.anisotropic_super_resolution(model, tokenizer, volume, target_shape, model_type='maskgit', interpolation_mode='trilinear', rasterization_method='hilbert', **inpaint_kwargs)

Super-resolve volume with anisotropic resolution.

Parameters:

• model (Module)

• volume (Tensor)

• target_shape (tuple[int, int, int])

• model_type (Literal['autoreg', 'maskgit', 'flow', 'diffusion', 'bayesian_flow'], default: 'maskgit')

• interpolation_mode (str, default: 'trilinear')

• rasterization_method (str, default: 'hilbert')

Return type:

Tensor

medlatents.inference.progressive_super_resolution(model, tokenizer, volume, target_slices, axis=0, num_stages=2, model_type='maskgit', **inpaint_kwargs)

Progressive super-resolution: upsample in multiple stages.

Parameters:

• model (Module)

• volume (Tensor)

• target_slices (int)

• axis (int, default: 0)

• num_stages (int, default: 2)

• model_type (Literal['autoreg', 'maskgit', 'flow', 'diffusion', 'bayesian_flow'], default: 'maskgit')

Return type:

Tensor

medlatents.inference.compare_with_interpolation(model, tokenizer, volume, target_slices, axis=0, model_type='maskgit', **inpaint_kwargs)

Compare model-based super-resolution with interpolation baseline.

Parameters:

• model (Module)

• volume (Tensor)

• target_slices (int)

• axis (int, default: 0)

• model_type (str, default: 'maskgit')

Return type:

tuple[Tensor, Tensor]

Large-Volume Processing

medlatents.inference.encode_large_volume(tokenizer, volume)

Encode large volume. The tokenizer handles sliding window inference automatically.

Parameters:

volume (Tensor)

Return type:

Tensor

medlatents.inference.decode_large_volume(tokenizer, tokens)

Decode large token volume. The tokenizer handles sliding window inference automatically.

Parameters:

tokens (Tensor)

Return type:

Tensor

medlatents.inference.reconstruct_large_volume(tokenizer, volume, roi_size=(64, 64, 64), overlap=0.5)

Reconstruct large volume using sliding window inference.

The tokenizer’s reconstruct() method handles sliding window inference internally for large volumes.

Parameters:

• volume (Tensor)

• roi_size (tuple[int, ...], default: (64, 64, 64))

• overlap (float, default: 0.5)

Return type:

Tensor

KV-Cache

class medlatents.inference.PagedKVCache(num_layers, num_heads, head_dim, block_size=16, num_blocks=1024, dtype=torch.float16, device='cuda')

Bases: object

Paged KV-cache for efficient memory management with variable-length sequences.

Inspired by vLLM’s PagedAttention. Allocates memory in fixed-size blocks to reduce fragmentation and enable efficient batched generation.

This is a simplified implementation - full paged attention requires custom CUDA kernels for optimal performance.

Parameters:

• num_layers (int)

• num_heads (int)

• head_dim (int)

• block_size (int, default: 16)

• num_blocks (int, default: 1024)

• dtype (dtype, default: torch.float16)

• device (device | str, default: 'cuda')

__init__(num_layers, num_heads, head_dim, block_size=16, num_blocks=1024, dtype=torch.float16, device='cuda')

Initialize paged KV-cache.

Parameters:

• num_layers (int) – Number of transformer layers

• num_heads (int) – Number of attention heads

• head_dim (int) – Dimension per head

• block_size (int, default: 16) – Tokens per block (default 16)

• num_blocks (int, default: 1024) – Total blocks in the pool

• dtype (dtype, default: torch.float16) – Storage dtype

• device (device | str, default: 'cuda') – Device for storage

allocate_sequence(seq_id, initial_length=0)

Allocate blocks for a new sequence.

Parameters:

• seq_id (int) – Unique sequence identifier

• initial_length (int, default: 0) – Initial sequence length (allocates blocks)

Return type:

None

free_sequence(seq_id)

Free all blocks for a sequence.

Parameters:

seq_id (int) – Sequence identifier

Return type:

None

append(seq_id, layer_idx, key, value)

Append new K/V to a sequence’s cache.

Parameters:

• seq_id (int) – Sequence identifier

• layer_idx (int) – Layer index

• key (Tensor) – New key tensor [1, num_heads, 1, head_dim] (single token)

• value (Tensor) – New value tensor [1, num_heads, 1, head_dim]

Return type:

None

get(seq_id, layer_idx)

Get K/V cache for a sequence and layer.

Parameters:

• seq_id (int) – Sequence identifier

• layer_idx (int) – Layer index

Return type:

tuple[Tensor, Tensor] | None

Returns:

(key, value) tensors or None if empty

memory_stats()

Get memory usage statistics.

Return type:

dict[str, float]

Returns:

Dict with memory info

class medlatents.inference.QuantizedKVCache(num_layers, config=None, max_length=None)

Bases: object

Memory-efficient KV-cache with INT8 quantization.

Reduces memory usage by ~4x compared to FP32 or ~2x compared to FP16. Uses per-token symmetric quantization for a good accuracy/speed tradeoff.

Usage:

cache = QuantizedKVCache(num_layers=12, config=QuantizationConfig())

# During generation
for token in tokens:
    for layer_idx, block in enumerate(model.blocks):
        past_kv = cache.get(layer_idx)
        output, new_kv = block.forward_with_cache(x, freqs, past_kv)
        cache.update(layer_idx, new_kv)

Parameters:

• num_layers (int)

• config (QuantizationConfig | None, default: None)

• max_length (int | None, default: None)

__init__(num_layers, config=None, max_length=None)

Initialize quantized KV-cache.

Parameters:

• num_layers (int) – Number of transformer layers

• config (QuantizationConfig | None, default: None) – Quantization configuration

• max_length (int | None, default: None) – Optional maximum sequence length (for pre-allocation)

get(layer_idx)

Get dequantized K/V for a layer.

Parameters:

layer_idx (int) – Layer index

Return type:

tuple[Tensor, Tensor] | None

Returns:

(key, value) tuple or None if cache is empty

update(layer_idx, new_kv)

Update cache with new K/V, quantizing the new values.

Parameters:

• layer_idx (int) – Layer index

• new_kv (tuple[Tensor, Tensor]) – (key, value) tuple from forward_with_cache

Return type:

None

clear()

Clear the cache.

Return type:

None

property length: int

Current cached sequence length.

memory_usage()

Compute current memory usage in MB.

Return type:

dict[str, float]

Returns:

Dict with memory breakdown

Quantization and Pruning

class medlatents.inference.ModelQuantizer(config=None)

Bases: object

Quantize PyTorch models for efficient inference.

Supports: - Dynamic quantization (quick, no calibration needed) - Static quantization (requires calibration data) - SmoothQuant for transformer-friendly quantization

Example

quantizer = ModelQuantizer(config=QuantizationConfig(method=”dynamic”)) quantized_model = quantizer.quantize(model)

Parameters:

config (QuantizationConfig | None, default: None)

__init__(config=None)

Initialize quantizer.

Parameters:

config (QuantizationConfig | None, default: None) – Quantization configuration

quantize(model, calibration_data=None)

Quantize a model.

Parameters:

• model (Module) – Model to quantize

• calibration_data (Iterator[Tensor] | None, default: None) – Iterator of calibration inputs (for static quantization)

Return type:

Module

Returns:

Quantized model

class medlatents.inference.QuantizationConfig(method=QuantizationMethod.DYNAMIC, dtype='int8', per_channel=True, symmetric=True, calibration_batches=100, modules_to_quantize=None, modules_to_skip=<factory>, smoothing_alpha=0.5)

Bases: object

Configuration for model quantization.

Variables:

• method – Quantization method to use

• dtype – Target dtype for quantized weights (int8/int4 for quantization, fp16/bf16 for casting)

• per_channel – Use per-channel quantization (more accurate)

• symmetric – Use symmetric quantization

• calibration_batches – Number of batches for static quantization calibration

• modules_to_quantize – List of module types to quantize (None = all supported)

• modules_to_skip – List of module names to skip

• smoothing_alpha – SmoothQuant smoothing factor (0.5 is typical)

Parameters:

• method (QuantizationMethod, default: <QuantizationMethod.DYNAMIC: 'dynamic'>)

• dtype (Literal['int8', 'int4', 'fp16', 'bf16'], default: 'int8')

• per_channel (bool, default: True)

• symmetric (bool, default: True)

• calibration_batches (int, default: 100)

• modules_to_quantize (list[str] | None, default: None)

• modules_to_skip (list[str], default: <factory>)

• smoothing_alpha (float, default: 0.5)

method: medlatents.inference.quantization.QuantizationMethod = 'dynamic'

dtype: Literal['int8', 'int4', 'fp16', 'bf16'] = 'int8'

per_channel: bool = True

symmetric: bool = True

calibration_batches: int = 100

modules_to_quantize: list[str] | None = None

modules_to_skip: list[str]

smoothing_alpha: float = 0.5

get_quant_dtype()

Get quantization dtype, defaulting to int8 for non-quantized types.

Return type:

Literal['int8', 'int4']

__init__(method=QuantizationMethod.DYNAMIC, dtype='int8', per_channel=True, symmetric=True, calibration_batches=100, modules_to_quantize=None, modules_to_skip=<factory>, smoothing_alpha=0.5)

Parameters:

• method (QuantizationMethod, default: <QuantizationMethod.DYNAMIC: 'dynamic'>)

• dtype (Literal['int8', 'int4', 'fp16', 'bf16'], default: 'int8')

• per_channel (bool, default: True)

• symmetric (bool, default: True)

• calibration_batches (int, default: 100)

• modules_to_quantize (list[str] | None, default: None)

• modules_to_skip (list[str], default: <factory>)

• smoothing_alpha (float, default: 0.5)

class medlatents.inference.ModelPruner(config=None)

Bases: object

Prune PyTorch models to reduce size and computation.

Supports: - Magnitude-based pruning (unstructured and structured) - WANDA pruning (weights + activations) - Iterative pruning with gradual sparsity increase

Example

pruner = ModelPruner(config=PruningConfig(sparsity=0.5)) pruned_model = pruner.prune(model)

Parameters:

config (PruningConfig | None, default: None)

__init__(config=None)

Initialize pruner.

Parameters:

config (PruningConfig | None, default: None) – Pruning configuration

prune(model, calibration_data=None)

Prune a model.

Parameters:

• model (Module) – Model to prune

• calibration_data (Iterator[Tensor] | None, default: None) – Calibration data for activation-based pruning

Return type:

Module

Returns:

Pruned model

class medlatents.inference.PruningConfig(method=PruningMethod.MAGNITUDE, sparsity=0.5, structured_dim=0, granularity='element', iterative_steps=1, importance_scores=None)

Bases: object

Configuration for model pruning.

Variables:

• method – Pruning method to use

• sparsity – Target sparsity ratio (0.0 to 1.0)

• structured_dim – Dimension for structured pruning (0=output, 1=input)

• granularity – Pruning granularity (‘element’, ‘row’, ‘column’)

• iterative_steps – Number of iterative pruning steps

• importance_scores – Pre-computed importance scores (for custom pruning)

Parameters:

• method (PruningMethod, default: <PruningMethod.MAGNITUDE: 'magnitude'>)

• sparsity (float, default: 0.5)

• structured_dim (int, default: 0)

• granularity (Literal['element', 'row', 'column'], default: 'element')

• iterative_steps (int, default: 1)

• importance_scores (dict[str, Tensor] | None, default: None)

method: medlatents.inference.quantization.PruningMethod = 'magnitude'

sparsity: float = 0.5

structured_dim: int = 0

granularity: Literal['element', 'row', 'column'] = 'element'

iterative_steps: int = 1

importance_scores: dict[str, torch.Tensor] | None = None

__init__(method=PruningMethod.MAGNITUDE, sparsity=0.5, structured_dim=0, granularity='element', iterative_steps=1, importance_scores=None)

Parameters:

• method (PruningMethod, default: <PruningMethod.MAGNITUDE: 'magnitude'>)

• sparsity (float, default: 0.5)

• structured_dim (int, default: 0)

• granularity (Literal['element', 'row', 'column'], default: 'element')

• iterative_steps (int, default: 1)

• importance_scores (dict[str, Tensor] | None, default: None)

medlatents.inference.prune_attention_heads(model, heads_to_prune, num_heads, head_dim)

Prune specific attention heads from a model.

Parameters:

• model (Module) – Model containing attention layers

• heads_to_prune (dict[int, list[int]]) – Dict mapping layer index to list of head indices to prune

• num_heads (int) – Total number of heads per layer

• head_dim (int) – Dimension per head

Return type:

Module

Returns:

Model with pruned heads (zeroed out)

medlatents.inference.export_to_onnx(model, sample_input, output_path, opset_version=17, dynamic_axes=None)

Export model to ONNX format for deployment.

Parameters:

• model (Module) – Model to export

• sample_input (Tensor | dict[str, Tensor]) – Sample input for tracing

• output_path (str) – Output file path (.onnx)

• opset_version (int, default: 17) – ONNX opset version

• dynamic_axes (dict[str, dict[int, str]] | None, default: None) – Dynamic axis specification

Return type:

None

medlatents.inference.estimate_model_size(model)

Estimate model size in different formats.

Parameters:

model (Module) – Model to analyze

Return type:

dict[str, float]

Returns:

Dict with size estimates in MB