depth_fm

depth_fm

Mars DepthFM — flow-matching DTM generation from monocular Mars imagery.

Public API:

from depth_fm import MarsDepthFM, build_model
from depth_fm import DepthFMLightningModule

Subpackages:

• models/ — MarsDepthFM wrapper + CompVis UNet backbone + experimental variants.
• training/ — Lightning module and train_lightning.py CLI entry point.
• data/ — Map-style adapter, LitData streaming datamodule, normalization (scalers).
• objectives/ — Loss functions and evaluation metrics.
• flow/ — Flow-matching noise schedule.
• viz/ — Publication figures (train_viz) and training-side diagnostics (debug_viz).

MarsDepthFM

MarsDepthFM(backbone: Module, vae: AutoencoderKL, noising_step: int, empty_text_embed: Tensor, freeze_encoder: bool = False)

Bases: Module

Wrapper for DepthFM flow matching depth estimation.

Exposes a clean interface for the training loop

encode_to_latent — pixel → latent (frozen VAE, no grad) decode_from_latent — latent → pixel (frozen VAE, grad allowed for losses) predict_velocity — core UNet forward

Source code in src/depth_fm/models/mars_depthfm.py

def __init__(
        self,
        backbone: nn.Module,
        vae: AutoencoderKL,
        noising_step: int,
        empty_text_embed: torch.Tensor,
        freeze_encoder: bool = False,
):
    super().__init__()
    self.backbone = backbone
    self.vae = vae
    self.noising_step = noising_step
    # Register as buffer so it moves with .to(device) and gets saved
    self.register_buffer("empty_text_embed", empty_text_embed)

    self.scale_factor = _SCALE_FACTOR
    self.latent_channels = 4  # SD 1.5 VAE always outputs 4-channel latents

    if freeze_encoder:
        self._freeze_encoder()

encode_to_latent

encode_to_latent(pixel_images: Tensor) -> torch.Tensor

Encode pixel-space images (B, 3, H, W) ∈ [-1, 1] to latent space. Returns: (B, 4, h, w) scaled latents.

Source code in src/depth_fm/models/mars_depthfm.py

@torch.no_grad()
def encode_to_latent(self, pixel_images: torch.Tensor) -> torch.Tensor:
    """
    Encode pixel-space images (B, 3, H, W) ∈ [-1, 1] to latent space.
    Returns: (B, 4, h, w) scaled latents.
    """
    posterior = self.vae.encode(pixel_images).latent_dist
    latent = posterior.mode()
    return latent * self.scale_factor

decode_from_latent

decode_from_latent(latent: Tensor) -> torch.Tensor

Decode latent (B, 4, h, w) to pixel space (B, 3, H, W) ∈ [-1, 1].

Called WITHOUT torch.no_grad() when used for pixel-space training losses so gradients flow back through the frozen VAE decoder to v_pred.

Source code in src/depth_fm/models/mars_depthfm.py

def decode_from_latent(self, latent: torch.Tensor) -> torch.Tensor:
    """
    Decode latent (B, 4, h, w) to pixel space (B, 3, H, W) ∈ [-1, 1].

    Called WITHOUT torch.no_grad() when used for pixel-space training losses
    so gradients flow back through the frozen VAE decoder to v_pred.
    """
    latent_unscaled = latent / self.scale_factor
    return self.vae.decode(latent_unscaled).sample

forward

forward(ims: Tensor, num_steps: int = 2, ensemble_size: int = 4) -> torch.Tensor

Full inference forward — mirrors DepthFM.forward() line-for-line.

Parameters:

Name	Type	Description	Default
ims	Tensor	(1, 3, H, W) in [-1, 1]	required
num_steps	int	Euler ODE steps	2
ensemble_size	int	repeat image in batch; average depth at end	4

Returns:

Name	Type	Description
depth	Tensor	(1, 1, H, W) in [0, 1]

Source code in src/depth_fm/models/mars_depthfm.py

def forward(
        self,
        ims: torch.Tensor,
        num_steps: int = 2,
        ensemble_size: int = 4,
) -> torch.Tensor:
    """
    Full inference forward — mirrors DepthFM.forward() line-for-line.

    Args:
        ims: (1, 3, H, W) in [-1, 1]
        num_steps: Euler ODE steps
        ensemble_size: repeat image in batch; average depth at end

    Returns:
        depth: (1, 1, H, W) in [0, 1]
    """
    from depth_fm.flow.noise import q_sample, per_sample_min_max_normalization

    if ensemble_size > 1:
        ims = ims.repeat(ensemble_size, 1, 1, 1)

    bs = ims.shape[0]
    device = ims.device

    # Encode image latent (mode, no sampling) — mirrors DepthFM.encode(sample_posterior=False)
    ims_z = self.encode_to_latent(ims)

    # Null text conditioning (cross-attention) — matches dfm.py repeat semantics
    conditioning = self.empty_text_embed.to(device=device, dtype=ims_z.dtype).expand(bs, -1, -1)

    # Clean image latent for channel-concat conditioning
    context = ims_z

    # Noise-augment the ODE starting point — matches dfm.py exactly:
    #   if self.noising_step > 0: x_source = q_sample(x_source, self.noising_step)
    x_source = ims_z.clone()
    if self.noising_step > 0:
        x_source = q_sample(x_source, self.noising_step)

    # Euler ODE: x_source → depth_z  (equivalent to torchdiffeq Euler in dfm.py)
    z_t = x_source
    dt = 1.0 / num_steps
    for step in range(num_steps):
        t_val = step * dt
        t = torch.full((bs,), t_val, device=device, dtype=z_t.dtype)
        # UNet forward: channel-concat with context (clean image), cross-attn with conditioning
        v = self.backbone(x=z_t, t=t, context=context, context_ca=conditioning)
        z_t = z_t + dt * v

    depth_z = z_t

    # Decode + channel mean — matches dfm.py: depth.mean(dim=1, keepdim=True)
    depth = self.decode_from_latent(depth_z)  # (E, 3, H, W)
    depth = depth.mean(dim=1, keepdim=True)  # (E, 1, H, W)

    if ensemble_size > 1:
        depth = depth.mean(dim=0, keepdim=True)  # (1, 1, H, W)

    # exp() then per-sample min-max normalize → [0, 1]  — matches dfm.py
    depth = per_sample_min_max_normalization(depth.exp())

    return depth

predict_depth

predict_depth(ims: Tensor, num_steps: int = 2, ensemble_size: int = 4) -> torch.Tensor

Public inference API — mirrors DepthFM.predict_depth() exactly.

Returns:

Name	Type	Description
depth	Tensor	(1, 1, H, W) in [0, 1]

Source code in src/depth_fm/models/mars_depthfm.py

@torch.no_grad()
def predict_depth(
        self,
        ims: torch.Tensor,
        num_steps: int = 2,
        ensemble_size: int = 4,
) -> torch.Tensor:
    """
    Public inference API — mirrors DepthFM.predict_depth() exactly.

    Returns:
        depth: (1, 1, H, W) in [0, 1]
    """
    return self.forward(ims, num_steps, ensemble_size)

predict_velocity

predict_velocity(z_t: Tensor, t: Tensor, z_img_cond: Tensor) -> torch.Tensor

Predict velocity field v_θ(z_t, t; z_img).

The UNet concatenates z_img_cond onto z_t INTERNALLY before input_blocks, so we pass them as separate arguments (x and context).

Parameters:

Name	Type	Description	Default
z_t	Tensor	(B, 4, h, w) — noisy interpolant on the flow path	required
t	Tensor	(B,) in [0, 1] — continuous timestep	required
z_img_cond	Tensor	(B, 4, h, w) — clean image latent (conditioning)	required

Returns:

Name	Type	Description
v_pred	Tensor	(B, 4, h, w) — predicted velocity

Source code in src/depth_fm/models/mars_depthfm.py

def predict_velocity(
        self,
        z_t: torch.Tensor,
        t: torch.Tensor,
        z_img_cond: torch.Tensor,
) -> torch.Tensor:
    """
    Predict velocity field v_θ(z_t, t; z_img).

    The UNet concatenates z_img_cond onto z_t INTERNALLY before input_blocks,
    so we pass them as separate arguments (x and context).

    Args:
        z_t:        (B, 4, h, w) — noisy interpolant on the flow path
        t:          (B,) in [0, 1] — continuous timestep
        z_img_cond: (B, 4, h, w) — clean image latent (conditioning)

    Returns:
        v_pred: (B, 4, h, w) — predicted velocity
    """

    B = z_t.size(0)

    # Expand stored null text embedding to batch size and cast to current dtype
    null_ca = self.empty_text_embed.to(dtype=z_t.dtype).expand(B, -1, -1)

    # UNet forward: x=z_t, context=z_img_cond (channel-concat inside),
    # context_ca=null_ca (cross-attention in transformer blocks)
    # t is passed as-is in [0,1]; timestep_embedding handles sinusoidal encoding
    # (B, 4, h, w) — no slicing needed, out_channels=4
    return self.backbone(
        x=z_t,
        t=t,
        context=z_img_cond,
        context_ca=null_ca,
    )

DepthFMLightningModule

DepthFMLightningModule(config)

Bases: LightningModule

Lightning module for DepthFM flow matching training.

Handles: - Conditional flow matching velocity loss - Surface normals auxiliary loss with warmup - Logit-normal timestep sampling - Noise augmentation on source distribution - Multi-step Euler inference for validation - Full metric computation and logging

Source code in src/depth_fm/training/lightning_module.py

def __init__(self, config):
    super().__init__()
    self.save_hyperparameters(config)
    self.config = config

    # Build model
    torch.backends.cudnn.benchmark = True
    self.model = build_model(config)

    if config.training.get("optimizer", "adamw") != "shampoo":
        self.model = self.model.to(memory_format=torch.channels_last)
    else:
        logger.warning("Disabled channels_last memory format for shampoo training")

    # Loss
    lc = config.training.losses
    self.loss_fn = CombinedLoss(
        velocity_weight=lc.velocity_weight,
        normals_weight=lc.get("normals_weight", 0.1),
        normals_start_step=lc.get("normals_start_step", 2000),
        use_confidence_weighting=lc.get("use_confidence_weighting", False),
        freq_weight=lc.get("freq_weight", 0.0),
        freq_start_step=lc.get("freq_start_step", 0),
        freq_alpha=lc.get("freq_alpha", 1.0),
        grad_weight=lc.get("grad_weight", 0.0),
        grad_start_step=lc.get("grad_start_step", 0),
        grad_scales=tuple(lc.get("grad_scales", [1, 2, 4])),
        photo_weight=lc.get("photo_weight", 0.0),
        photo_start_step=lc.get("photo_start_step", 2000),
        huber_weight=lc.get("huber_weight", 0.0),
        huber_start_step=lc.get("huber_start_step", 0.0),
        laplacian_weight=lc.get("laplacian_weight", 0.0),
        laplacian_start_step=lc.get("laplacian_start_step", 0.0),
        ordinal_weight=lc.get("ordinal_weight", 0.0),
        ordinal_start_step=lc.get("ordinal_start_step", 0.0),

    )

    # Elevation normaliser — used as fallback when batch has no residual_scale.
    # clip must match the adapter so clip=False data (values outside [-1,1]) is
    # denormalized correctly (clip has no effect on denormalize_prediction but
    # keeps the object consistent with the data that was produced).
    self.evel_normalizer = GlobalLogNormalizer(
        config.data.get("elev_ref_scale", DEFAULT_ELEV_REF_SCALE),
        clip=config.data.get("clip", False),
    )

    # Flow matching config
    self.fm = config.training.flow_matching

    # Metric aggregators (reset each epoch)
    self._val_aggregator = MetricsAggregator()
    self._test_aggregator = MetricsAggregator()

    # Optional tag suffix for test-output files (used by the ablation
    # orchestrator when running test twice — once per checkpoint type).
    # Empty string → test_per_patch.npz / test_results.json (default).
    self._test_tag: str = ""

    # Track validation history for convergence plotting
    self.val_history: dict[str, list[float]] = {
        "val/rmse": [], "val/abs_rel": [], "val/delta_1": [],
        "val/loss": [], "train/loss": [],
    }

    # Cache for grad norm logging (avoid recomputing every step)
    self._grad_norm_log_interval = 50

    # ── Visualisation system (initialised lazily in on_train_start) ──
    # These are set up once training begins and the Trainer context is
    # available (world_size, default_root_dir, logger, etc.).
    self._vis_worker: mp.Process | None = None
    self._vis_queue: mp.SimpleQueue | None = None
    self._vis_staging_dir: str | None = None
    self._vector_fig_dir: str | None = None

estimate_uncertainty

estimate_uncertainty(z_img: Tensor, num_samples: int = 10, num_steps: int = 4) -> tuple[torch.Tensor, torch.Tensor]

Computes the epistemic uncertainty via stochastic ODE sampling.

Parameters:

Name	Type	Description	Default
z_img	Tensor	(B, C, H, W) Encoded input orthoimage.	required
num_samples	int	Number of stochastic noise initializations (N).	10
num_steps	int	Euler integration steps per sample.	4

Returns:

Name	Type	Description
mean_dtm	Tensor	(B, 1, H, W) The expected topographic surface.
var_dtm	Tensor	(B, 1, H, W) The pixel-wise epistemic variance.

Source code in src/depth_fm/training/lightning_module.py

@torch.no_grad()
def estimate_uncertainty(
        self,
        z_img: torch.Tensor,
        num_samples: int = 10,
        num_steps: int = 4
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Computes the epistemic uncertainty via stochastic ODE sampling.

    Args:
        z_img: (B, C, H, W) Encoded input orthoimage.
        num_samples: Number of stochastic noise initializations (N).
        num_steps: Euler integration steps per sample.

    Returns:
        mean_dtm: (B, 1, H, W) The expected topographic surface.
        var_dtm: (B, 1, H, W) The pixel-wise epistemic variance.
    """
    self._trace(f"Starting epistemic uncertainty estimation with N={num_samples}")

    dtm_hypotheses = []

    for i in range(num_samples):
        # _predict_depth automatically samples a new x_source ~ N(0, I)
        # via the _get_x_source() method internally.
        z_pred = self._predict_depth(z_img, num_steps=num_steps)

        pred_pix = self.evel_normalizer.denormalize_prediction(
            self._decode(z_pred)[:, 0])
        dtm_hypotheses.append(pred_pix)

    # Stack into shape: (N, B, 1, H, W)
    dtm_tensor = torch.stack(dtm_hypotheses, dim=0)

    # Calculate moments
    mean_dtm = torch.mean(dtm_tensor, dim=0)
    var_dtm = torch.var(dtm_tensor, dim=0, unbiased=True)

    return mean_dtm, var_dtm

on_train_end

on_train_end() -> None

Final wandb upload and vis worker shutdown.

Source code in src/depth_fm/training/lightning_module.py

def on_train_end(self) -> None:
    """Final wandb upload and vis worker shutdown."""
    # Upload any remaining staged figures
    self._deferred_wandb_upload()
    self._upload_vector_figures_artifact(prefix="val")

    # Shut down the worker process
    self._shutdown_vis_worker(timeout=60.0)
    self._trace("Training complete, vis worker shut down")

build_model

build_model(config) -> MarsDepthFM

Factory function: build the MarsDepthFM model from config.

Only supports backend="sd21" (the actual DepthFM architecture).

Source code in src/depth_fm/models/mars_depthfm.py

def build_model(config) -> MarsDepthFM:
    """
    Factory function: build the MarsDepthFM model from config.

    Only supports backend="sd21" (the actual DepthFM architecture).
    """
    backend = config.model.backend
    if backend != "sd21":
        raise ValueError(
            f"Unsupported backend: '{backend}'. Only 'sd21' is supported "
            "(the original DepthFM CompVis LDM UNet checkpoint)."
        )

    backbone, vae, noising_step, empty_text_embed = load_sd21_backend(
        depthfm_checkpoint=config.model.depthfm_checkpoint,
        vae_id=config.model.vae_id,
        use_checkpoint=config.model.use_checkpoint
    )

    model_type = config.model.get("model_type", "debug")

    if model_type == "debug":
        backbone = DebugUNet()

    elif model_type == "modulated":
        backbone = ModulatedMicroFlowNet()

    model = MarsDepthFM(
        backbone=backbone,
        vae=vae,
        noising_step=noising_step,
        empty_text_embed=empty_text_embed,
        freeze_encoder=config.model.get("freeze_encoder", False),
    )

    # Gradient checkpointing — UNetModel supports use_checkpoint per-block,
    # but it's set at construction time via ldm_hparams. If requested post-hoc,
    # patch the flag on all ResBlock/AttentionBlock instances.
    if config.model.get("gradient_checkpointing", False):
        from depth_fm.models.unet.openaimodel import ResBlock, AttentionBlock
        count = 0
        for module in model.backbone.modules():
            if isinstance(module, (ResBlock, AttentionBlock)):
                module.use_checkpoint = True
                count += 1
        logger.info("Gradient checkpointing enabled on %d blocks", count)

    # torch.compile for throughput — use "reduce-overhead" to avoid the extra
    # VRAM that "max-autotune" reserves for CUDA-graph autotuning. This mode
    # still captures CUDA graphs but skips expensive kernel search.
    # disable=True skips compile entirely (useful for debugging).
    if config.model.get("compile_backbone", False):
        compile_mode = config.model.get("compile_mode", "max-autotune")
        logger.info("Compiling backbone with mode=%s", compile_mode)
        model.backbone = torch.compile(
            model.backbone,
            mode=compile_mode,
            fullgraph=False,  # allow graph breaks rather than erroring
        )

    return model

load_sd21_backend

load_sd21_backend(depthfm_checkpoint: str, vae_id: str, device: str = 'cpu', use_checkpoint=True) -> Tuple[nn.Module, AutoencoderKL, int, torch.Tensor]

Load the actual DepthFM model from the official checkpoint.

Checkpoint format

noising_step — int, e.g. 200 ldm_hparams — dict of UNetModel constructor kwargs empty_text_embedding — (1, seq_len, context_dim) null text conditioning state_dict — UNetModel weights (keys match CompVis naming)

Returns:

Name	Type	Description
unet	Module	trainable UNetModel
vae	AutoencoderKL	frozen SD 1.5 AutoencoderKL
noising_step	int	int
empty_text_embed	Tensor	(1, seq_len, context_dim) tensor

Source code in src/depth_fm/models/mars_depthfm.py

def load_sd21_backend(depthfm_checkpoint: str, vae_id: str, device: str = "cpu", use_checkpoint=True) -> Tuple[
    nn.Module, AutoencoderKL, int, torch.Tensor]:
    """
    Load the actual DepthFM model from the official checkpoint.

    Checkpoint format:
        noising_step          — int, e.g. 200
        ldm_hparams           — dict of UNetModel constructor kwargs
        empty_text_embedding  — (1, seq_len, context_dim) null text conditioning
        state_dict            — UNetModel weights (keys match CompVis naming)

    Returns:
        unet: trainable UNetModel
        vae: frozen SD 1.5 AutoencoderKL
        noising_step: int
        empty_text_embed: (1, seq_len, context_dim) tensor
    """
    from depth_fm.models.unet import UNetModel

    ckpt = torch.load(depthfm_checkpoint, map_location=device, weights_only=False)

    noising_step = ckpt["noising_step"]
    empty_text_embed = torch.from_numpy(ckpt["empty_text_embedding"]).to(device=device)  # (1, seq_len, 1024)
    ldm_hparams = dict(ckpt["ldm_hparams"])  # copy to avoid mutation
    ldm_hparams["use_checkpoint"] = use_checkpoint

    unet = UNetModel(**ldm_hparams)
    missing, unexpected = unet.load_state_dict(ckpt["state_dict"], strict=True)
    if missing:
        logger.warning("Missing keys loading DepthFM UNet: %d", len(missing))
        for k in missing[:5]:
            logger.warning("  %s", k)
    if unexpected:
        logger.warning("Unexpected keys loading DepthFM UNet: %d", len(unexpected))
        for k in unexpected[:5]:
            logger.warning("  %s", k)

    # SD 1.5 VAE — frozen encoder/decoder
    vae: AutoencoderKL = AutoencoderKL.from_pretrained(vae_id, subfolder="vae").to(device=device)
    vae.eval()
    vae.requires_grad_(False)

    n_params = sum(p.numel() for p in unet.parameters()) / 1e6
    logger.info(
        "Loaded DepthFM: UNet=%.1fM params, noising_step=%d, "
        "empty_text_embed shape=%s",
        n_params, noising_step, tuple(empty_text_embed.shape),
    )

    return unet, vae, noising_step, empty_text_embed