# SAII-CLDM LDM Checkpoints

This repository hosts the raw CompVis latent-diffusion checkpoints for SAII-CLDM.
For the Diffusers-format release with bundled inference code, see
[mally-2000/saii-cldm-synthetic](https://huggingface.co/mally-2000/saii-cldm-synthetic).

## Files

| File | Description |
| --- | --- |
| `stage1_vqgan.ckpt` | Stage 1 VQGAN checkpoint used as the first-stage autoencoder. |
| `stage2_ldm.ckpt` | Stage 2 SAII-CLDM latent diffusion checkpoint. |
| `stage1_vqgan_marmousi.yaml` | Public Stage 1 VQGAN model/data config. |
| `stage2_saii_cldm_marmousi.yaml` | Public Stage 2 SAII-CLDM model/data config. |

There are two public YAML files because there are two model checkpoints. The Lightning trainer/logger/callback config used for training is part of the code repository, not this model-weight repository.

## Download

```bash
pip install huggingface_hub
huggingface-cli download \
  --repo-type model \
  --local-dir ./saii-cldm-ldm-checkpoints \
  mally-2000/saii-cldm-ldm-checkpoints
```

Or in Python:

```python
from huggingface_hub import snapshot_download

snapshot_download(
    "mally-2000/saii-cldm-ldm-checkpoints",
    repo_type="model",
    local_dir="./saii-cldm-ldm-checkpoints",
)
```

## Use With The Code Repository

Clone the SAII-CLDM code repository:

```bash
git clone https://github.com/Mally-cj/cldm-diffusers.git
cd cldm-diffusers
```

Set the required local paths in `.env`:

```bash
cp .env.example .env
# edit FIRST_STAGE_CKPT, MARMousi_NPZ, and OVERTHRUST_DATA_DIR as needed
```

If you downloaded this repository next to the code repository, copy the checkpoints
and public configs into the code repository:

```bash
mkdir -p models configs
cp ../saii-cldm-ldm-checkpoints/stage1_vqgan.ckpt models/
cp ../saii-cldm-ldm-checkpoints/stage2_ldm.ckpt models/
cp ../saii-cldm-ldm-checkpoints/stage1_vqgan_marmousi.yaml configs/
cp ../saii-cldm-ldm-checkpoints/stage2_saii_cldm_marmousi.yaml configs/
```

Set the required local paths in `.env`:

```bash
FIRST_STAGE_CKPT=./models/stage1_vqgan.ckpt
# set MARMousi_NPZ and OVERTHRUST_DATA_DIR to your local data paths
```

Run CLDM inference on the Overthrust benchmark:

```bash
CUDA_VISIBLE_DEVICES=0 python eval_overthrust.py CLDM \
  --ckpt ./models/stage2_ldm.ckpt \
  --config ./configs/stage2_saii_cldm_marmousi.yaml \
  --output runs/eval_cldm \
  --device cuda \
  --steps 30
```

If GPU 0 is occupied or you hit OOM, switch to another GPU:

```bash
CUDA_VISIBLE_DEVICES=1 python eval_overthrust.py CLDM \
  --ckpt ./models/stage2_ldm.ckpt \
  --config ./configs/stage2_saii_cldm_marmousi.yaml \
  --output runs/eval_cldm \
  --device cuda \
  --steps 30
```

To train Stage 1 VQGAN from scratch:

```bash
CUDA_VISIBLE_DEVICES=0 python train_latent_diffusion.py -t \
  --base ./configs/stage1_vqgan_marmousi.yaml
```

To train Stage 2 SAII-CLDM from scratch, use the code repository's
`configs/training_lightning.yaml` together with the Stage 2 model/data config:

```bash
CUDA_VISIBLE_DEVICES=0 python train_latent_diffusion.py -t \
  --base ./configs/stage2_saii_cldm_marmousi.yaml ./configs/training_lightning.yaml
```

To continue Stage 2 training from the released checkpoint, add `--resume`:

```bash
CUDA_VISIBLE_DEVICES=0 python train_latent_diffusion.py -t \
  --resume ./models/stage2_ldm.ckpt \
  --base ./configs/stage2_saii_cldm_marmousi.yaml ./configs/training_lightning.yaml
```

## Paper

Seismic Acoustic Impedance Inversion Framework Based on Conditional Latent Generative Diffusion Model.
arXiv: [2506.13529](https://arxiv.org/abs/2506.13529)