schrodinger.application.matsci.flywheel.ldmol.generator module

High-level LDMol generator for text-to-molecule generation.

Integrates MolT5 text encoder, VAE autoencoder, DiT, and DDIM sampling.

class schrodinger.application.matsci.flywheel.ldmol.generator.LDMolGenerator(device='cpu', logger=None)

Bases: BaseModel

Text-to-molecule generator using LDMol.

Generates molecules from text descriptions via latent diffusion. All models are loaded lazily on first use.

__init__(device='cpu', logger=None)
Parameters:

  • device (str) – Compute device (‘cpu’, ‘cuda’, ‘mps’)

  • logger – Logger instance for debug output

ensureModelsLoaded()

Load all models on first use (lazy loading).

Loads MolT5 text encoder, VAE autoencoder, and DiT in eval mode.

loadDiT()

Download (if needed), create, and load the DiT checkpoint.

Weights are loaded with strict=False to allow EMA checkpoints that may contain extra keys not in the model definition.

encodePrompts(text_prompts, num_samples)

Encode text prompts to MolT5 embeddings and tile for batch.

Parameters:

  • text_prompts (list(str)) – Text descriptions

  • num_samples (int) – Number of molecules per prompt

Return type:

tuple(torch.Tensor, torch.Tensor, int)

Returns:

text_embs (batch, seq, dim), attn_mask (batch, seq), batch_size

runDiffusionSampling(text_embs, attn_mask, batch_size, diffusion, cfg_scale)

Run DDIM reverse diffusion with optional classifier-free guidance.

Parameters:

  • text_embs (torch.Tensor) – Conditioned text embeddings (batch, seq, dim)

  • attn_mask (torch.Tensor) – Text attention mask (batch, seq)

  • batch_size (int) – Number of molecules to generate

  • diffusion – SpacedDiffusion instance

  • cfg_scale (float) – Guidance scale (>= 1.0)

Return type:

torch.Tensor

Returns:

Denoised latents (batch, seq_len, channels)

generate(text_prompts, num_samples=1, num_sampling_steps=100, cfg_scale=7.5, max_length=150, return_latents=False)

Generate molecules from text descriptions.

Pipeline: text → MolT5 embeddings → DiT DDIM denoising → VAE decode → SMILES.

Parameters:

  • text_prompts (str or list(str)) – Text description(s) to condition generation

  • num_samples (int) – Number of molecules to generate per prompt

  • num_sampling_steps (int) – DDIM reverse diffusion steps

  • cfg_scale (float) – Classifier-free guidance scale (>= 1.0)

  • max_length (int) – Maximum autoregressive SMILES decode length

  • return_latents (bool) – If True, return (smiles, latents) tuple

Return type:

list(str) or tuple(list(str), torch.Tensor)

Returns:

Generated SMILES strings, or (smiles, latents) if return_latents

decodeLatents(latents, max_length=150)

Decode latent tensors to SMILES strings.

Parameters:

  • latents (torch.Tensor) – Latents (batch, seq_len, channels)

  • max_length (int) – Maximum autoregressive decode length

Return type:

list(str)

Returns:

Decoded SMILES strings