Generation API

High-level, checkpoint-driven generation for discrete and continuous latent models. These classes wrap a trained model together with its tokenizer so you can go from a checkpoint to decoded images in a few lines. For low-level sampling primitives (schedulers, guidance, early stopping), see Sampling.

Discrete Latent Generation

class medlatents.generation.DiscreteLatentGenerator(model_type, model, tokenizer, device)[source][source]

Bases: object

Unified generator for discrete latent generative models.

Supports: - autoreg: Autoregressive generation - maskgit: Iterative parallel generation (MaskGIT) - flow: Discrete flow matching generation - diffusion: D3PM discrete diffusion generation

Parameters:
__init__(model_type, model, tokenizer, device)[source][source]
Parameters:
classmethod from_checkpoints(model_type, model_path, tokenizer_path, device=None, weights_only=True)[source][source]

Load generator from model and tokenizer checkpoints.

Parameters:

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

  • model_path (str)

  • tokenizer_path (str)

  • device (device | None, default: None)

  • weights_only (bool, default: True)

generate(num_samples, seq_length, temperature=1.0, top_k=None, num_steps=None, seed=None, **kwargs)[source][source]

Generate samples based on model type.

Parameters:

  • num_samples (int) – Number of samples to generate

  • seq_length (int) – Length of sequences

  • temperature (float, default: 1.0) – Sampling temperature

  • top_k (int | None, default: None) – Top-k sampling parameter

  • num_steps (int | None, default: None) – Number of iterative steps (maskgit/flow/diffusion)

  • seed (int | None, default: None) – Random seed for reproducibility (None for random)

  • **kwargs – Additional model-specific parameters

Return type:

Tensor

Returns:

Generated token sequences [num_samples, seq_length]

generate_volumes(num_samples, seq_length, output_dir, temperature=1.0, top_k=None, num_steps=None, seed=None, **kwargs)[source][source]

Generate samples and save as NIfTI volumes.

Parameters:

  • num_samples (int)

  • seq_length (int)

  • output_dir (str)

  • temperature (float, default: 1.0)

  • top_k (int | None, default: None)

  • num_steps (int | None, default: None)

  • seed (int | None, default: None)

inpaint_volumes(input_pattern, output_dir, likelihood_threshold=0.005, temperature=1.0, top_k=None)[source][source]

Inpaint anomalies in volumes.

Parameters:

  • input_pattern (str)

  • output_dir (str)

  • likelihood_threshold (float, default: 0.005)

  • temperature (float, default: 1.0)

  • top_k (int | None, default: None)

Continuous Latent Generation

class medlatents.generation.ContinuousLatentGenerator(model_type, model, tokenizer, device, diffusion=None, flow=None)[source][source]

Bases: object

Generator for continuous latent diffusion and flow models.

Parameters:
__init__(model_type, model, tokenizer, device, diffusion=None, flow=None)[source][source]
Parameters:
generate(num_samples, latent_shape=None, guidance_scale=1.0, condition=None, null_condition=None, num_steps=None, temperature=1.0, initial_noise=None, model_kwargs=None, seed=None)[source][source]
Parameters:

  • num_samples (int)

  • latent_shape (tuple[int, ...] | None, default: None)

  • guidance_scale (float, default: 1.0)

  • num_steps (int | None, default: None)

  • temperature (float, default: 1.0)

  • initial_noise (Tensor | None, default: None)

  • model_kwargs (dict | None, default: None)

  • seed (int | None, default: None)

Return type:

Tensor

decode_latents(latents, **decode_kwargs)[source][source]

Decode latents to reconstruction using tokenizer.

Parameters:

latents (Tensor)

Return type:

Tensor

generate_volumes(num_samples, output_dir, latent_shape=None, guidance_scale=1.0, condition=None, null_condition=None, num_steps=None, temperature=1.0, initial_noise=None, model_kwargs=None, affine=None, seed=None)[source][source]
Parameters:
Return type:

None