| --- |
| license: cc-by-nc-4.0 |
| tags: |
| - diffusion |
| - ddpm |
| - 3d |
| - volumetric |
| - monai |
| - generative |
| - synthetic-data |
| - atmospheric-science |
| --- |
| |
| # 3D Convective-Cloud Denoising Diffusion Probabilistic Model |
|
|
| An unconditional **3D Denoising Diffusion Probabilistic Model (DDPM)** that |
| generates synthetic **64×64×64 convective-cloud cells** (vertical-velocity |
| volumes). Roughly 190M parameters. |
|
|
| **What it is for:** producing synthetic 3D cloud cells to augment scarce |
| atmospheric-simulation data — for example as additional training data for 3D |
| cloud segmentation. |
|
|
| **What it is not:** this is not a weather-forecasting model and not a physical |
| simulation. It samples volumes resembling its training distribution; it does not |
| solve atmospheric dynamics. |
|
|
| **Project:** [3D Convective-Cloud Generation collection](https://huggingface.co/collections/HoqueMahmudul/3d-convective-cloud-generation-ddpm-190m-params-6a6ba230c89356ff617d481a) |
|
|
| > **Status:** research artifact from work with an associated manuscript **in |
| > preparation**. Methodology, training details, and evaluation are deliberately |
| > not described here and will appear in the publication. |
|
|
| ## Requirements |
|
|
| This is a **MONAI** model, not a `transformers` model. You need |
| **`monai-generative`** (which provides the `generative` package) to build the |
| network before loading the weights: |
|
|
| ```bash |
| pip install monai-generative==0.2.3 monai==1.5.2 torch safetensors huggingface_hub |
| ``` |
|
|
| A `requirements.txt` with the verified-working versions is included in this repo: |
|
|
| ```bash |
| pip install -r <(curl -sL https://huggingface.co/HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d/raw/main/requirements.txt) |
| ``` |
|
|
| > [!IMPORTANT] |
| > **Use `monai-generative`, not MONAI core.** MONAI core (>= 1.4) ships its own |
| > `DiffusionModelUNet`, `DDPMScheduler` and `DiffusionInferer`, and they look |
| > like drop-in replacements. **They are not weight-compatible with this |
| > checkpoint.** |
| > |
| > | | `monai-generative` 0.2.3 | MONAI core 1.5 | |
| > |---|---|---| |
| > | Attention keys | `...attentions.0.to_q.weight` | `...attentions.0.`**`attn`**`.to_q.weight` | |
| > | Constructor arg | `num_channels=[...]` | `channels=[...]` | |
| > | Strict load here | works | **52 missing / 52 unexpected keys** | |
| > |
| > Both classes have the *identical* parameter count (189,815,297), so if you hit |
| > the key error and "fix" it with `strict=False`, you will get a model with **52 |
| > randomly-initialised attention tensors**. It will run and produce |
| > plausible-looking output, and it will be silently wrong. |
| > |
| > **Never load these weights with `strict=False`.** If a strict load fails, the |
| > environment is wrong — fix the environment. |
| |
| ## Files |
| |
| | File | What it is | |
| |---|---| |
| | `best_model.safetensors` | The model weights (a plain PyTorch `state_dict`) | |
| | `requirements.txt` | Verified-working dependency versions | |
| |
| ## How to load and generate a sample |
| |
| The configuration below is **required** — the weights are a bare `state_dict`, so |
| the network must be constructed exactly this way before loading. |
|
|
| ```python |
| import torch, numpy as np |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| from generative.networks.nets import DiffusionModelUNet |
| from generative.networks.schedulers import DDPMScheduler |
| from generative.inferers import DiffusionInferer |
| |
| REPO = "HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d" |
| device = "cuda" |
| |
| # 1. build the network with exactly this configuration |
| model = DiffusionModelUNet( |
| spatial_dims=3, in_channels=1, out_channels=1, |
| num_channels=[256, 256, 512], |
| attention_levels=[False, False, True], |
| num_head_channels=[0, 0, 512], |
| num_res_blocks=2, |
| ) |
| |
| # 2. load the weights -- strict=True is deliberate; see the warning above |
| weights = hf_hub_download(REPO, "best_model.safetensors") |
| model.load_state_dict(load_file(weights), strict=True) |
| model.to(device).eval() |
| |
| # 3. the sampler must use these settings |
| scheduler = DDPMScheduler( |
| num_train_timesteps=1000, schedule="scaled_linear_beta", |
| beta_start=0.0005, beta_end=0.0195, |
| ) |
| scheduler.set_timesteps(1000) |
| inferer = DiffusionInferer(scheduler) |
| |
| # 4. sample one volume (1000 reverse steps -- use a GPU) |
| noise = torch.randn((1, 1, 64, 64, 64), device=device) |
| with torch.no_grad(), torch.amp.autocast("cuda"): |
| x = inferer.sample(input_noise=noise, diffusion_model=model, scheduler=scheduler) |
| |
| # 5. convert the [0,1] output back to raw units |
| # clip the LOWER bound at 1, not 0: the model's background output is very |
| # slightly negative, which would otherwise denormalise to 0 and leave you |
| # with a background value the format does not use. |
| blob = x[0, 0].cpu().numpy() |
| blob = np.clip(blob * (3100 - 1) + 1, 1, 3100).astype(np.float32) |
| |
| # 6. enforce the value gap the data format requires (see below) |
| gap = (blob > 1) & (blob < 300) |
| blob[gap] = np.where(blob[gap] < 150, 1, 300) |
| |
| blob = np.round(blob).astype(np.uint16) # 64^3 synthetic cloud cell |
| mask = blob >= 300 # binary cloud mask |
| ``` |
|
|
| ## Understanding the output |
|
|
| You need these conventions to interpret a generated volume: |
|
|
| | Property | Value | |
| |---|---| |
| | Shape / dtype | `64 × 64 × 64`, `uint16` | |
| | Background | `1` | |
| | Cloud voxels | `300` – `3022` (vertical velocity, scaled units) | |
| | Value gap | nothing lies strictly between `1` and `300` | |
| | Voxel size | 0.5 km isotropic (so 20 voxels = 10 km) | |
| | Normalisation | `(raw - 1) / (3100 - 1)`, clipped to `[0, 1]` | |
|
|
| The model operates in the normalised `[0, 1]` space, so step 5 above is required |
| to return to raw units, and step 6 restores the value gap that real volumes have. |
|
|
| ## Post-processing — extracting a single cloud cell |
|
|
| Steps 5 and 6 of the snippet give you a valid volume, but a raw sample often |
| contains several disconnected pieces (or none). To get a usable single cloud |
| cell you need to threshold, find connected components, keep the largest, and |
| blank the rest. Requires `scipy`: |
|
|
| ```python |
| import numpy as np |
| from scipy import ndimage |
| |
| def extract_cloud(blob, min_voxels=6, max_voxels=None): |
| """Keep only the largest 26-connected cloud component; blank everything else. |
| |
| Returns (cleaned, mask, size). |
| cleaned : uint16 volume containing at most one cloud component |
| mask : boolean cloud mask for `cleaned` |
| size : voxel count of the kept component; 0 means "discard this sample" |
| """ |
| mask = blob >= 300 |
| if not mask.any(): |
| return blob.copy(), mask, 0 # empty sample |
| |
| # 26-connectivity: faces, edges AND corners count as connected |
| labels, n = ndimage.label(mask, structure=np.ones((3, 3, 3))) |
| sizes = ndimage.sum(mask, labels, range(1, n + 1)) |
| keep = int(np.argmax(sizes)) + 1 |
| largest = int(sizes.max()) |
| |
| cleaned = blob.copy() |
| cleaned[labels != keep] = 1 # fragments -> background |
| |
| if largest < min_voxels: |
| return cleaned, cleaned >= 300, 0 # noise only -> discard |
| if max_voxels is not None and largest > max_voxels: |
| return cleaned, cleaned >= 300, largest # implausibly large -> review |
| return cleaned, cleaned >= 300, largest |
| |
| |
| cleaned, mask, size = extract_cloud(blob) |
| if size == 0: |
| print("empty or noise-only sample -- generate another") |
| else: |
| print(f"single cloud cell, {size} voxels " |
| f"({size * 0.5**3:.1f} km^3 at 0.5 km voxels)") |
| ``` |
|
|
| `min_voxels` rejects noise-only samples and `max_voxels` (optional) flags blobs |
| that fill an implausible share of the volume. Both are yours to choose — pick |
| them from the size range your application needs. Expect to discard a large |
| fraction of samples; see the note below. |
|
|
| ## Practical notes before you use this |
|
|
| - **Sampling is variable and low-yield.** Only a minority of raw samples are |
| usable single cloud cells; the rest come out empty, as noise fragments, or |
| occasionally as a blob filling the whole volume. **Plan on generating a surplus |
| and filtering** with `extract_cloud()` above, discarding samples where it |
| returns `size == 0`. This is the intended workflow, not a sign of a defect. |
| - **The size distribution is skewed.** Generated cells tend to be smaller and |
| weaker than those in the source data, so a raw synthetic set is not a |
| distribution-matched substitute for real data. |
| - **One source dataset.** Behaviour outside that simulation regime is unknown. |
| - **Cost.** Each sample takes 1000 reverse diffusion steps. A GPU is strongly |
| recommended. |
|
|
| ## Intended use and limitations |
|
|
| Intended for **research and educational use** — synthetic data augmentation and |
| generative-modelling experiments on volumetric atmospheric data. |
|
|
| No perceptual or physical-fidelity metrics are published for this model, so **no |
| fidelity or accuracy claims are made**. Do not use it for operational |
| forecasting, hazard assessment, or any decision-making application. |
|
|
| ## Data and credit |
|
|
| Trained on convective-cloud cells extracted from atmospheric-simulation output. |
| **The dataset is not included or redistributed here.** |
|
|
| The underlying simulations come from **Dr. Xiaowen Li's lab at Morgan State |
| University**; please credit that lab as the source of the simulation data if you |
| use this model. The diffusion model in this repository was developed by |
| Mahmudul Hoque. |
|
|
| ## License |
|
|
| **CC BY-NC 4.0** — [Creative Commons Attribution-NonCommercial 4.0 |
| International](https://creativecommons.org/licenses/by-nc/4.0/). |
|
|
| You may use, share, and adapt these weights for **non-commercial** purposes with |
| attribution. Commercial use is not permitted. Please also credit the source of |
| the simulation data as described above. |
|
|
| ## References |
|
|
| - Ho, Jain & Abbeel, *Denoising Diffusion Probabilistic Models*, NeurIPS 2020 |
| ([arXiv:2006.11239](https://arxiv.org/abs/2006.11239)) |
| - [MONAI GenerativeModels](https://github.com/Project-MONAI/GenerativeModels) |
|
|
| ## Contact |
|
|
| Mahmudul Hoque, Morgan State University — <mahoq1@morgan.edu> |
|
|
