| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| """Latent-geometry helpers for the FLIP latent diffusion tutorial. |
| |
| The diffusion model denoises inside the frozen autoencoder's latent space, so both the training and |
| the validation pass need to know that space's shape and scale. Those helpers live here because both |
| passes in ``trainer.py`` share them. |
| |
| This tutorial is a `standard` job, whose required upload set is |
| ``trainer.py``/``config.json``/``models.py``; this module is uploaded beside them. |
| """ |
|
|
| import hashlib |
| import logging |
| import math |
| from collections.abc import Mapping |
| from typing import Any |
|
|
| import torch |
| from monai.inferers import LatentDiffusionInferer |
| from monai.networks.schedulers import DDPMScheduler |
| from torch.amp import autocast |
|
|
| logger = logging.getLogger(__name__) |
|
|
| _AUTOENCODER_PREFIX = "autoencoder." |
|
|
|
|
| def autoencoder_fingerprint(state_dict: Mapping[str, Any]) -> str: |
| """SHA-256 over an autoencoder's weights: each tensor's name, dtype, shape and bytes, by name. |
| |
| Keys may carry the ``autoencoder.`` prefix (an extracted checkpoint) or not (the submodule's own |
| ``state_dict()``); both spellings hash alike. ``extract_autoencoder.py`` records it in |
| ``autoencoder_config.yaml`` and the trainer recomputes it on the weights it actually received. |
| """ |
| tensors = {key.removeprefix(_AUTOENCODER_PREFIX): torch.as_tensor(value) for key, value in state_dict.items()} |
| digest = hashlib.sha256() |
| for name in sorted(tensors): |
| tensor = tensors[name].detach().cpu().contiguous() |
| digest.update(f"{name}|{tensor.dtype}|{tuple(tensor.shape)}|".encode()) |
| digest.update(tensor.reshape(-1).view(torch.uint8).numpy().tobytes()) |
| return digest.hexdigest() |
|
|
|
|
| def derive_new_latent_shape(input_shape: list, num_downsamplings: int) -> list: |
| """Round each latent dimension up to a multiple of ``2**num_downsamplings``, so the UNet can halve it that often.""" |
| factor = 2**num_downsamplings |
| return [math.ceil(size / factor) * factor for size in input_shape] |
|
|
|
|
| def latent_shapes(model: torch.nn.Module, spatial_shape: list) -> tuple[list, list]: |
| """Infer the autoencoder latent shape and the diffusion-model-compatible (padded) latent shape. |
| |
| Args: |
| model (torch.nn.Module): The composite LDM network (``models.get_model()``). |
| spatial_shape (list): The input spatial shape from ``config.json``. |
| |
| Returns: |
| tuple[list, list]: ``(autoencoder_latent_shape, ldm_latent_shape)``. |
| """ |
| autoencoder_latent_shape = [i / (2 ** (len(model.autoencoder.decoder.channels) - 1)) for i in spatial_shape] |
| ldm_latent_shape = derive_new_latent_shape( |
| autoencoder_latent_shape, len(model.diffusion_model.block_out_channels) - 1 |
| ) |
| return autoencoder_latent_shape, ldm_latent_shape |
|
|
|
|
| def resolve_scale_factor( |
| model: torch.nn.Module, |
| sample_images: torch.Tensor, |
| device: torch.device, |
| configured: float | None = None, |
| ) -> float: |
| """Resolve the latent scale factor, preferring a pinned config value over a derived one. |
| |
| The diffusion model trains on latents normalised by this factor, so **every site must use the |
| same value or their updates are not averaging comparable models.** Deriving it from a local |
| batch (what happens when ``LATENT_SCALE_FACTOR`` is null) gives each site a slightly different value; with a frozen |
| autoencoder the factor is a fixed property of that autoencoder and the data, so it can simply be |
| pinned in ``config.json`` as ``LATENT_SCALE_FACTOR``. |
| |
| Args: |
| model (torch.nn.Module): The composite LDM network with the autoencoder weights loaded. |
| sample_images (torch.Tensor): A batch used to derive the factor when none is configured. |
| device (torch.device): Device the model lives on. |
| configured (float | None): ``LATENT_SCALE_FACTOR`` from ``config.json``, or None to derive. |
| |
| Returns: |
| float: The scale factor to hand to :class:`LatentDiffusionInferer`. |
| """ |
| if configured is not None: |
| logger.info(f"Latent scale factor: {configured} (pinned via LATENT_SCALE_FACTOR).") |
| return float(configured) |
|
|
| with torch.no_grad(): |
| with autocast(enabled=True, device_type=device.type): |
| sample_z = model.autoencoder.encode_stage_2_inputs(sample_images.to(device)) |
| scale_factor = 1 / torch.std(sample_z) |
| del sample_z |
|
|
| derived = float(scale_factor.item()) |
| logger.info( |
| f"Latent scale factor: {derived} (derived from a local batch — LATENT_SCALE_FACTOR is unset, " |
| "so other sites will derive their own and may differ; pin it in config.json for a " |
| "multi-site run)." |
| ) |
| return derived |
|
|
|
|
| def build_inferer( |
| model: torch.nn.Module, |
| scheduler: DDPMScheduler, |
| spatial_shape: list, |
| sample_images: torch.Tensor, |
| device: torch.device, |
| scale_factor: float | None = None, |
| ) -> tuple[LatentDiffusionInferer, list]: |
| """Build the ``LatentDiffusionInferer`` for the diffusion phase. |
| |
| Args: |
| model (torch.nn.Module): The composite LDM network with the frozen autoencoder loaded. |
| scheduler (DDPMScheduler): The DDPM noise scheduler. |
| spatial_shape (list): The input spatial shape from ``config.json``. |
| sample_images (torch.Tensor): A batch of images used to infer the latent scale factor when |
| none is pinned in the config. |
| device (torch.device): Device the model lives on. |
| scale_factor (float | None): ``LATENT_SCALE_FACTOR`` from ``config.json``, or None to derive |
| it from ``sample_images``. |
| |
| Returns: |
| tuple[LatentDiffusionInferer, list]: The inferer and the padded LDM latent shape (needed by |
| callers to shape the sampling noise). |
| """ |
| autoencoder_latent_shape, ldm_latent_shape = latent_shapes(model, spatial_shape) |
| resolved = resolve_scale_factor(model, sample_images, device, scale_factor) |
|
|
| inferer = LatentDiffusionInferer( |
| scheduler=scheduler, |
| ldm_latent_shape=ldm_latent_shape, |
| autoencoder_latent_shape=autoencoder_latent_shape, |
| scale_factor=resolved, |
| ) |
| return inferer, ldm_latent_shape |
|
|
