Clean public checkpoint config names

README.md

@@ -1,6 +1,6 @@
 # SAII-CLDM LDM Checkpoints
 
-This repository hosts the raw LDM-backend checkpoints for SAII-CLDM.
+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).
 
@@ -8,15 +8,14 @@ For the Diffusers-format release with bundled inference code, see
 
 | File | Description |
 | --- | --- |
-| `stage1_vqgan.ckpt` | Stage 1 VQGAN checkpoint (first stage autoencoder), 721 MB. |
-| `stage2_ldm.ckpt` | Stage 2 latent diffusion / CLDM checkpoint (original A101 run, epoch 212 / step 13,991), 4.3 GB. |
-| `ldm_backend_vqgan_marmousi.yaml` | Stage 1 VQGAN training config. |
-| `ldm_backend_a101_train_hwd.yaml` | Stage 2 CLDM training config, aligned with the released `stage2_ldm.ckpt`. |
-| `ldm_backend_lightning.yaml` | Lightning trainer / logger / Overthrust eval callback config (used for Stage 2). |
+| `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. |
 
-## Download the checkpoints
+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.
 
-With [huggingface_hub](https://huggingface.co/docs/huggingface_hub):
+## Download
 
 ```bash
 pip install huggingface_hub
@@ -30,6 +29,7 @@ Or in Python:
 
 ```python
 from huggingface_hub import snapshot_download
+
 snapshot_download(
     "mally-2000/saii-cldm-ldm-checkpoints",
     repo_type="model",
@@ -37,33 +37,36 @@ snapshot_download(
 )
 ```
 
-## Run inference on Overthrust
+## Use With The Code Repository
 
-Clone the training code repository and copy the checkpoints/configs into it:
+Clone the SAII-CLDM code repository and copy the checkpoints/configs into it:
 
 ```bash
-git clone <SAII-CLDM-training-repo>
-cd SAII-CLDM
+git clone https://github.com/Mally-cj/cldm-diffusers.git
+cd cldm-diffusers
 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/*.yaml configs/
+cp saii-cldm-ldm-checkpoints/stage1_vqgan_marmousi.yaml configs/
+cp saii-cldm-ldm-checkpoints/stage2_saii_cldm_marmousi.yaml configs/
 ```
 
-Set the required paths in `.env` (copy from `.env.example`):
+Set the required local paths in `.env`:
 
 ```bash
-FIRST_STAGE_CKPT=./models/stage1_vqgan.ckpt
-OVERTHRUST_DATA_DIR=./data/overthrust
+cp .env.example .env
+# edit FIRST_STAGE_CKPT, MARMousi_NPZ, and OVERTHRUST_DATA_DIR as needed
 ```
 
-Then run CLDM inference on the Overthrust benchmark. `eval_overthrust.py` defaults to `configs/ldm_backend_a101_train_hwd.yaml`, which is now aligned with the released Stage 2 checkpoint:
+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
+  --device cuda \
+  --steps 30
 ```
 
 If GPU 0 is occupied or you hit OOM, switch to another GPU:
@@ -71,30 +74,21 @@ 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
+  --device cuda \
+  --steps 30
 ```
 
-The script will save `metrics_summary.json`, numpy arrays, and a comparison figure under `runs/eval_cldm`.
-
-## Resume training
-
-Stage 2 training / resume from the released checkpoint:
+To continue Stage 2 training, use the code repository's `configs/training_lightning.yaml` together with the Stage 2 model/data config:
 
 ```bash
-CUDA_VISIBLE_DEVICES=0 python train_ldm_backend.py -t \
+CUDA_VISIBLE_DEVICES=0 python train_latent_diffusion.py -t \
   --resume ./models/stage2_ldm.ckpt \
-  --base ./configs/ldm_backend_a101_train_hwd.yaml ./configs/ldm_backend_lightning.yaml
-```
-
-Stage 1 VQGAN training from scratch:
-
-```bash
-CUDA_VISIBLE_DEVICES=0 python train_ldm_backend.py -t \
-  --base ./configs/ldm_backend_vqgan_marmousi.yaml
+  --base ./configs/stage2_saii_cldm_marmousi.yaml ./configs/training_lightning.yaml
 ```
 
 ## Paper
 
-SAII-CLDM: Seismic Acoustic Impedance Inversion through Conditional Latent Diffusion Model.
-arXiv: [2506.13529](https://arxiv.org/html/2506.13529v1)
+Seismic Acoustic Impedance Inversion Framework Based on Conditional Latent Generative Diffusion Model.
+arXiv: [2506.13529](https://arxiv.org/abs/2506.13529)

ldm_backend_lightning.yaml

@@ -1,25 +0,0 @@
-lightning:
-  callbacks:
-    overthrust_eval:
-      params:
-        dipin_v: 0.012
-        every_n_epochs: 10
-        f0: 30
-        f0_phase: 0
-        log_wandb: true
-        noise_snr: 15
-        num_timesteps: 1000
-        output_dir: runs/eval
-      target: ldm.callbacks.overthrust_eval.OverthrustEvalCallback
-  logger:
-    params:
-      entity: 601882280-xi-an-jiaotong-university-
-      name: ldm-backend-a101-train-hwd
-      offline: true
-      project: saii-cldm-reproduce
-    target: pytorch_lightning.loggers.WandbLogger
-  trainer:
-    accelerator: gpu
-    benchmark: true
-    gpus: 1
-    max_epochs: 500

ldm_backend_vqgan_marmousi.yaml → stage1_vqgan_marmousi.yaml

@@ -74,8 +74,8 @@ lightning:
     target: pytorch_lightning.loggers.WandbLogger
     params:
       entity: 601882280-xi-an-jiaotong-university-
-      project: saii-cldm-reproduce
-      name: ldm-backend-vqgan-marmousi
+      project: saii-cldm
+      name: saii-vqgan-marmousi
       offline: true
   trainer:
     benchmark: true