Module

The deepspatial.module serves as the training and inference orchestrator for the framework. Built upon PyTorch Lightning, it elegantly encapsulates the Flow Matching training objective, multi-modal loss computation, Exponential Moving Average (EMA) weight updates, and continuous integration solvers.

Lightning Module

The core engine responsible for managing the optimization lifecycle and the generation phase.

DeepSpatialModule

DeepSpatial Module for Training & Inference.

Key Methods

Fundamental operations managed by the module. The sample method is particularly critical as it executes the ODE/SDE integration process for 3D volume reconstruction.

DeepSpatialModule.training_step	Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.
DeepSpatialModule.sample	Integrates the learned flow field to reconstruct intermediate biological states.