Add Overthrust inference benchmark and model card

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/demo.png filter=lfs diff=lfs merge=lfs -text
+data/Overthrust_trueimp.mat filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,116 @@
+---
+library_name: diffusers
+pipeline_tag: image-to-image
+tags:
+- seismic-inversion
+- impedance-inversion
+- diffusion
+- ddpm
+- overthrust
+---
+
+# Seismic-LDDPM
+
+Seismic-LDDPM is a latent DDPM pipeline for seismic impedance inversion. The
+pipeline takes a low-frequency impedance image (`dipin`) and a synthetic seismic
+record (`record`) and predicts the impedance image.
+
+This repository includes:
+
+- Diffusers-format model components: `vq_model`, `unet`, `scheduler`, and
+  `condition_encoder`.
+- `SeismicImpInvLDDPMPipeline` in `pipeline.py`.
+- A complete Overthrust benchmark sample at `data/Overthrust_trueimp.mat`.
+- Inference scripts under `inference/`.
+
+## Installation
+
+```bash
+git clone https://huggingface.co/mally-2000/seismic-lddpm
+cd seismic-lddpm
+pip install -r requirements.txt
+```
+
+## Overthrust Evaluation
+
+The Overthrust evaluation script is intentionally fixed to the bundled
+`data/Overthrust_trueimp.mat`. It cuts the full model into six `256 x 256`
+patches, synthesizes the seismic records and low-frequency impedance inputs,
+runs inference, stitches the six predictions back together, and computes the
+metrics.
+
+```bash
+python inference/eval.py \
+  --model . \
+  --output outputs/overthrust \
+  --num-inference-steps 1000
+```
+
+Outputs:
+
+- `outputs/overthrust/full_target.npy`
+- `outputs/overthrust/full_prediction.npy`
+- `outputs/overthrust/full_reconstruction.npy`
+- `outputs/overthrust/comparison_impedance.png`
+- `outputs/overthrust/metrics_summary.json`
+
+## Benchmark Result
+
+Evaluated locally on the bundled Overthrust benchmark with 1000 DDPM steps,
+`noise_snr=15`, `dipin_v=0.012`, `f0=30`, `phase=0`, `seed=1234`, and patch
+indices `[0, 1, 2, 3, 4, 5]`.
+
+| Space | PSNR | SSIM | PCC | RRE | NMSE |
+|---|---:|---:|---:|---:|---:|
+| Normalized | 30.7698 | 0.9339 | 0.9963 | 0.0435 | 0.001894 |
+| Impedance | 33.4413 | 0.9554 | 0.9957 | 0.0324 | 0.001050 |
+| VQ reconstruction | 37.7954 | 0.9677 | 0.9983 | 0.0209 | 0.000435 |
+
+![Overthrust evaluation](assets/demo.png)
+
+## Single-Sample Inference
+
+For a single prepared sample:
+
+```bash
+python inference/infer.py \
+  --dipin path/to/dipin.npy \
+  --record path/to/record.npy \
+  --model . \
+  --output outputs/single
+```
+
+The input arrays may be `H x W`, `C x H x W`, or `B x C x H x W`. The script
+converts them to BCHW tensors and saves `prediction.npy` and `prediction.png`.
+
+## Python Usage
+
+```python
+import torch
+from pipeline import SeismicImpInvLDDPMPipeline
+
+pipe = SeismicImpInvLDDPMPipeline.from_pretrained(
+    "mally-2000/seismic-lddpm",
+    torch_dtype=torch.float32,
+    trust_remote_code=True,
+).to("cuda")
+
+result = pipe(
+    dipin=dipin,      # torch.Tensor, BCHW
+    record=record,    # torch.Tensor, BCHW
+    num_inference_steps=1000,
+    seed=1234,
+)
+
+prediction = result.impedance_samples
+```
+
+## Notes
+
+- `inference/dataset.py` contains a lightweight `SeismicBase` and
+  `OverthrustTrueimpDataset`; it does not depend on the original training
+  repository's `ldm.data.seisimic`.
+- Synthetic record generation is seeded through the benchmark configuration so
+  the published Overthrust evaluation is reproducible.
+- The bundled Overthrust file is used only as a compact benchmark input for
+  reproducing this model's inference pipeline.

data/Overthrust_trueimp.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59345b022a90f174efd444004192f9edf93ef5f65c556a84f52e9596f1695bd5
+size 334424

examples/expected_metrics.json

@@ -0,0 +1,46 @@
+{
+  "config": {
+    "size": 256,
+    "patch_indices": [
+      0,
+      1,
+      2,
+      3,
+      4,
+      5
+    ],
+    "noise_snr": 15,
+    "dipin_v": 0.012,
+    "f0": 30,
+    "f0_phase": 0,
+    "seed": 1234,
+    "zhengyan_type": "nonlinear",
+    "normalize": "minmax",
+    "batch_size": 3,
+    "num_inference_steps": 1000
+  },
+  "normalized": {
+    "PSNR": 30.76981345155257,
+    "rre": 0.043521951884031296,
+    "SSIM": 0.9339061199595424,
+    "PCC": 0.9963035366574778,
+    "nmse": 0.001894210814498365,
+    "mse": 3.811614279402499e-09
+  },
+  "impedance": {
+    "PSNR": 33.44134288739278,
+    "rre": 0.03240736946463585,
+    "SSIM": 0.955363744873021,
+    "PCC": 0.9957485049549735,
+    "nmse": 0.0010502231307327747,
+    "mse": 0.11484166561534005
+  },
+  "encode_impedance": {
+    "PSNR": 37.79544219163976,
+    "rre": 0.020859846845269203,
+    "SSIM": 0.9676508176373475,
+    "PCC": 0.9982799028636675,
+    "nmse": 0.0004351270035840571,
+    "mse": 0.04758103889550172
+  }
+}

inference/__init__.py

@@ -0,0 +1 @@
+"""Inference helpers for seismic-lddpm."""

inference/dataset.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import pylops
+import scipy.io
+import torch
+from scipy.fftpack import fft, ifft
+from scipy.signal import butter, filtfilt
+from torch.utils.data import Dataset
+
+
+class SeismicBase:
+    @staticmethod
+    def phaseshift(w: np.ndarray, d: float) -> np.ndarray:
+        if d == 0:
+            return w
+        wf_shift = fft(w) * np.exp(1j * (np.pi * d / 180.0))
+        return np.real(ifft(wf_shift))
+
+    @staticmethod
+    def add_gaussian_band_noise(
+        target_snr: float,
+        data: np.ndarray,
+        rng: np.random.Generator | None = None,
+    ) -> tuple[np.ndarray, float]:
+        if target_snr == 0:
+            return data, 0.0
+        rng = rng or np.random.default_rng()
+        signal_energy = np.linalg.norm(data) ** 2
+        noise_energy = signal_energy / (10 ** (target_snr / 10))
+        initial_noise = rng.normal(loc=0, scale=1, size=data.shape)
+        noise = filtfilt(
+            np.ones(3) / 3,
+            1,
+            filtfilt(np.ones(3) / 3, 1, initial_noise.T, method="gust").T,
+            method="gust",
+        )
+        noise = noise * np.sqrt(noise_energy / np.linalg.norm(noise) ** 2)
+        noisy_data = data + noise
+        actual_snr = 10 * np.log10(signal_energy / np.linalg.norm(noise) ** 2)
+        return noisy_data, float(actual_snr)
+
+    @staticmethod
+    def add_gaussian_noise(
+        target_snr: float,
+        data: np.ndarray,
+        rng: np.random.Generator | None = None,
+    ) -> tuple[np.ndarray, float]:
+        if target_snr == 0:
+            return data, 0.0
+        rng = rng or np.random.default_rng()
+        signal_energy = np.linalg.norm(data) ** 2
+        noise_energy = signal_energy / (10 ** (target_snr / 10))
+        noise_std = np.sqrt(noise_energy / data.size)
+        noise = rng.normal(0, noise_std, data.shape)
+        noisy_data = data + noise
+        actual_snr = 10 * np.log10(signal_energy / np.linalg.norm(noise) ** 2)
+        return noisy_data, float(actual_snr)
+
+
+class OverthrustTrueimpDataset(SeismicBase, Dataset):
+    """Overthrust benchmark dataset used by seismic-lddpm evaluation."""
+
+    def __init__(
+        self,
+        size: int = 256,
+        interval: int = 1,
+        special_splits: bool = False,
+        use_mask: bool = False,
+        record_noraml: bool = True,
+        normalize: str = "minmax",
+        zhengyan_type: str = "linear",
+        train_keys: tuple[str, ...] | list[str] = ("image", "dipin", "record"),
+        ricks: tuple[int, ...] | list[int] = (30,),
+        ricks_phase: tuple[int, ...] | list[int] = (0,),
+        noise_snr: tuple[int, ...] | list[int] = (0,),
+        noise_type: str = "guassian_band",
+        dipins: tuple[float, ...] | list[float] = (0.012,),
+        dipin_nsmoothz: int = 20,
+        dipin_nsmoothx: int = 20,
+        patch_indices: tuple[int, ...] | list[int] | None = None,
+        base_seed: int = 1234,
+        data_dir: str | Path | None = None,
+        cache_dir: str | Path = "outputs/cache",
+        fixed_f0: int | None = None,
+        fixed_dipin_v: float | None = None,
+        fixed_noise_snr: int | None = None,
+        fixed_f0_phase: int | None = None,
+    ):
+        self.name = "Overthrust_trueimp"
+        self.size = size
+        self.interval = interval
+        self.special_splits = special_splits
+        self.use_mask = use_mask
+        self.record_noraml = record_noraml
+        self.normalize = normalize
+        self.zhengyan_type = zhengyan_type
+        self.train_keys = list(train_keys)
+        self.ricks = list(ricks)
+        self.ricks_phase = list(ricks_phase)
+        self.noise_snr = list(noise_snr)
+        self.noise_type = noise_type
+        self.dipins = list(dipins)
+        self.dipin_nsmoothz = dipin_nsmoothz
+        self.dipin_nsmoothx = dipin_nsmoothx
+        self.base_seed = base_seed
+        self.have_exp = False
+        self.info: dict[str, float | str] = {}
+        self.fixed_f0 = fixed_f0
+        self.fixed_dipin_v = fixed_dipin_v
+        self.fixed_noise_snr = fixed_noise_snr
+        self.fixed_f0_phase = fixed_f0_phase
+        self.data_dir = Path(data_dir or os.getenv("DATASET_DIR", "data"))
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+        self._load_big_impedance()
+        self._build_splits_and_patches(special_splits=special_splits)
+        self._build_wavelets()
+        self.big_reflect = self._load_or_build_reflect()
+        self.record_data = {
+            f0: {
+                phase: {
+                    snr: self._patches_from_big_image(
+                        self._load_or_build_record(f0=f0, phase=phase, noise_snr=snr)
+                    )
+                    for snr in self.noise_snr
+                }
+                for phase in self.ricks_phase
+            }
+            for f0 in self.ricks
+        }
+        self.dipin_datas = {
+            dipin_v: self._patches_from_big_image(self._load_or_build_dipin(dipin_v))
+            for dipin_v in self.dipins
+        }
+        all_indices = list(range(len(self.splits)))
+        self.patch_indices = all_indices if patch_indices is None else list(patch_indices)
+
+    def __len__(self) -> int:
+        return len(self.patch_indices)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        patch_idx = self.patch_indices[index]
+        f0 = self.fixed_f0 if self.fixed_f0 is not None else random.choice(self.ricks)
+        dipin_v = (
+            self.fixed_dipin_v
+            if self.fixed_dipin_v is not None
+            else random.choice(self.dipins)
+        )
+        noise_snr = (
+            self.fixed_noise_snr
+            if self.fixed_noise_snr is not None
+            else random.choice(self.noise_snr)
+        )
+        f0_phase = (
+            self.fixed_f0_phase
+            if self.fixed_f0_phase is not None
+            else random.choice(self.ricks_phase)
+        )
+        sample = {
+            "patch_idx": torch.tensor(patch_idx, dtype=torch.long),
+            "seed": torch.tensor(
+                self.base_seed + index + int(noise_snr) * 100, dtype=torch.long
+            ),
+        }
+        if "image" in self.train_keys:
+            sample["image"] = torch.from_numpy(self.file_data[patch_idx]).float()
+        if "dipin" in self.train_keys:
+            sample["dipin"] = torch.from_numpy(self.dipin_datas[dipin_v][patch_idx]).float()
+            sample["dipin_v"] = torch.tensor(dipin_v, dtype=torch.float32).reshape(1, 1, 1)
+        if "record" in self.train_keys:
+            sample["record"] = torch.from_numpy(
+                self.record_data[f0][f0_phase][noise_snr][patch_idx]
+            ).float()
+            sample["rick_v"] = torch.tensor(f0, dtype=torch.float32).reshape(1, 1, 1)
+            sample["rick_phase"] = torch.tensor(f0_phase, dtype=torch.float32).reshape(1, 1, 1)
+            sample["snr_v"] = torch.tensor(noise_snr, dtype=torch.float32).reshape(1, 1, 1)
+        if "reflection" in self.train_keys:
+            sample["reflection"] = torch.from_numpy(self.reflect_data[patch_idx]).float()
+        if "mask_speed" in self.train_keys:
+            sample["mask_speed"] = torch.from_numpy(
+                self.mask_data[patch_idx] * self.file_data[patch_idx]
+            ).float()
+        if self.use_mask:
+            sample["mask"] = torch.from_numpy(self.mask_data[patch_idx]).float()
+        return sample
+
+    def fan(self, x: np.ndarray) -> np.ndarray:
+        minn = 5.0931
+        maxn = 6.501110975896774
+        return np.exp(x * (maxn - minn) + minn) * 10.9 + 200
+
+    def inv_normal(self, x: np.ndarray) -> np.ndarray:
+        vmin = float(self.info["normal_min"])
+        vmax = float(self.info["normal_max"])
+        if self.normalize == "minmax":
+            return x * (vmax - vmin) + vmin
+        return x * vmax
+
+    def _load_big_impedance(self) -> None:
+        file_path = self.data_dir / "Overthrust_trueimp.mat"
+        if not file_path.exists():
+            raise FileNotFoundError(f"Overthrust data not found: {file_path}")
+        wave = scipy.io.loadmat(file_path)["Overthrust_trueimp"].T
+        wave = np.log(wave)
+        normal_min = wave.min()
+        normal_max = wave.max()
+        self.info.update(
+            {"normal_min": normal_min, "normal_max": normal_max, "normal": "max"}
+        )
+        self.big_img_unnorm = wave
+        self.big_speedimg = wave
+        if self.normalize == "max":
+            wave = wave / normal_max
+        elif self.normalize == "minmax":
+            wave = (wave - normal_min) / (normal_max - normal_min)
+        else:
+            raise ValueError(f"Unsupported normalize: {self.normalize}")
+        self.big_img = wave.astype(np.float32)
+
+    def _build_splits_and_patches(self, special_splits: bool = False) -> None:
+        self.big_mask = np.zeros(self.big_img.shape, dtype=np.float32)
+        for col in (100, 200, 300):
+            if col < self.big_mask.shape[1]:
+                self.big_mask[:, col : col + 1] = 1
+        if special_splits:
+            splits = []
+            for x in range(0, 551 - self.size, 20):
+                for y in range(0, 551 - self.size, 20):
+                    splits.append((x, y))
+            for y in range(0, 551 - self.size, 9):
+                splits.extend([(30, y), (90, y), (140, y)])
+        elif self.size == 256:
+            splits = [
+                (0, 0),
+                (146, 0),
+                (551 - 256, 0),
+                (0, 145),
+                (146, 145),
+                (551 - 256, 145),
+            ]
+        else:
+            splits = []
+            interval_size = self.size - 1
+            for r in range(0, self.big_img.shape[0] - self.size, interval_size):
+                for c in range(0, self.big_img.shape[1] - self.size, interval_size):
+                    splits.append((r, c))
+                splits.append((r, self.big_img.shape[1] - self.size))
+            for c in range(0, self.big_img.shape[1] - self.size, interval_size):
+                splits.append((self.big_img.shape[0] - self.size, c))
+            splits.append(
+                (self.big_img.shape[0] - self.size, self.big_img.shape[1] - self.size)
+            )
+
+        self.splits = []
+        patches = []
+        masks = []
+        for x, y in splits:
+            x2 = x + self.size
+            y2 = y + self.size
+            if x2 > self.big_img.shape[0] or y2 > self.big_img.shape[1]:
+                continue
+            self.splits.append((x, y))
+            patches.append(self.big_img[x:x2, y:y2].reshape(1, self.size, self.size))
+            masks.append(self.big_mask[x:x2, y:y2].reshape(1, self.size, self.size))
+        self.file_data = np.stack(patches, axis=0).astype(np.float32)[:: self.interval]
+        self.mask_data = np.stack(masks, axis=0).astype(np.float32)[:: self.interval]
+        self.splits = self.splits[:: self.interval]
+
+    def _build_wavelets(self) -> None:
+        nt0 = 256
+        dt0 = 0.002
+        self.wavelets = {}
+        for f0 in self.ricks:
+            self.wavelets[f0] = {}
+            wav = pylops.utils.wavelets.ricker(np.arange(nt0 // 2) * dt0, f0)[0]
+            for phase in self.ricks_phase:
+                self.wavelets[f0][phase] = self.phaseshift(wav, phase)
+
+    def _cache_path(self, name: str) -> Path:
+        return self.cache_dir / name
+
+    def _load_or_build_reflect(self) -> np.ndarray:
+        cache_path = self._cache_path(
+            f"Overthrust_trueimpBig_sesimic_reflect_{self.zhengyan_type}.npy"
+        )
+        if not cache_path.exists():
+            size = self.big_img.shape[0]
+            if self.zhengyan_type == "linear":
+                s1 = np.diag(0.5 * np.ones(size - 1, dtype="float32"), k=1) - np.diag(
+                    0.5 * np.ones(size - 1, dtype="float32"), k=-1
+                )
+                s1[-1] = s1[0] = 0
+                reflect = s1 @ self.big_img
+            elif self.zhengyan_type == "nonlinear":
+                expspeed = (
+                    np.exp(self.big_img_unnorm)
+                    if self.have_exp is False
+                    else self.big_img_unnorm
+                )
+                s1 = np.eye(size, k=1) - np.eye(size, k=0)
+                s2 = np.eye(size, k=1) + np.eye(size, k=0)
+                s1[-1] = 0
+                s2[-1] = 0
+                numerator = s1 @ expspeed
+                denominator = s2 @ expspeed
+                denominator = np.where(denominator < 1e-6, 1e-6, denominator)
+                reflect = numerator / denominator
+            else:
+                raise ValueError(f"Unsupported zhengyan_type: {self.zhengyan_type}")
+            np.save(cache_path, reflect)
+        reflect = np.load(cache_path).astype(np.float32)
+        self.reflect_data = self._patches_from_big_image(reflect)
+        return reflect
+
+    def _load_or_build_record(self, f0: int, phase: int, noise_snr: int) -> np.ndarray:
+        cache_path = self._cache_path(
+            f"Overthrust_trueimpBig_sesimic_record__{self.zhengyan_type}"
+            f"_ricker={f0:02d}-{phase:03d}_{self.noise_type}={noise_snr:02d}"
+            f"_seed={self.base_seed}.npy"
+        )
+        if not cache_path.exists():
+            wav = self.wavelets[f0][phase]
+            w_mat = pylops.utils.signalprocessing.convmtx(
+                wav, self.big_reflect.shape[0], len(wav) // 2
+            )[: self.big_reflect.shape[0]]
+            records_clear = w_mat @ self.big_reflect
+            rng = np.random.default_rng(self.base_seed + f0 * 1000 + phase * 10 + noise_snr)
+            if self.noise_type == "guassian_band":
+                record, _ = self.add_gaussian_band_noise(noise_snr, records_clear, rng=rng)
+            elif self.noise_type == "guassian":
+                record, _ = self.add_gaussian_noise(noise_snr, records_clear, rng=rng)
+            else:
+                raise ValueError(f"Unsupported noise_type: {self.noise_type}")
+            np.save(cache_path, record)
+        record = np.load(cache_path).astype(np.float32)
+        self.info.update(
+            {
+                "record_minn": min(float(self.info.get("record_minn", 10)), float(record.min())),
+                "record_maxn": max(float(self.info.get("record_maxn", -10)), float(record.max())),
+                "record_normal": "max",
+            }
+        )
+        if self.record_noraml:
+            record = record / 0.3215932963300079
+            self.info["record_maxn"] = 0.3215932963300079
+        return record
+
+    def _load_or_build_dipin(self, dipin_v: float) -> np.ndarray:
+        cache_path = self._cache_path(
+            f"Overthrust_trueimpBig_sesimic_dipin={dipin_v:.03f}.npy"
+        )
+        if not cache_path.exists():
+            bb, aa = butter(2, dipin_v, "low")
+            smooth_filter_z = np.ones(self.dipin_nsmoothz) / float(self.dipin_nsmoothz)
+            smooth_filter_x = np.ones(self.dipin_nsmoothx) / float(self.dipin_nsmoothx)
+            mback = filtfilt(bb, aa, self.big_img.T).T
+            mback = filtfilt(smooth_filter_z, 1, mback, axis=0)
+            mback = filtfilt(smooth_filter_x, 1, mback, axis=1)
+            np.save(cache_path, mback)
+        return np.load(cache_path).astype(np.float32)
+
+    def _patches_from_big_image(self, big_image: np.ndarray) -> np.ndarray:
+        patches = []
+        for x, y in self.splits:
+            patches.append(
+                big_image[x : x + self.size, y : y + self.size].reshape(
+                    1, self.size, self.size
+                )
+            )
+        return np.stack(patches, axis=0).astype(np.float32)

inference/eval.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from scipy.stats import pearsonr
+from skimage.metrics import structural_similarity
+from torch.utils.data import DataLoader
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from inference.dataset import OverthrustTrueimpDataset
+from pipeline import SeismicImpInvLDDPMPipeline
+
+
+OVERTHRUST_CONFIG = {
+    "size": 256,
+    "patch_indices": [0, 1, 2, 3, 4, 5],
+    "noise_snr": 15,
+    "dipin_v": 0.012,
+    "f0": 30,
+    "f0_phase": 0,
+    "seed": 1234,
+    "zhengyan_type": "nonlinear",
+    "normalize": "minmax",
+    "batch_size": 3,
+}
+
+
+def stitch_patches(
+    patches: list[np.ndarray], splits: list[tuple[int, int]], big_shape: tuple[int, int], img_size: int
+) -> np.ndarray:
+    rec = np.zeros(big_shape, dtype=np.float32)
+    cnt = np.zeros(big_shape, dtype=np.float32)
+    for idx, (x, y) in enumerate(splits):
+        rec[x : x + img_size, y : y + img_size] += patches[idx]
+        cnt[x : x + img_size, y : y + img_size] += 1
+    return rec / np.maximum(cnt, 1)
+
+
+def compute_metrics(prediction: np.ndarray, target: np.ndarray) -> dict[str, float]:
+    diff = prediction - target
+    denom = np.linalg.norm(diff.ravel()) ** 2
+    psnr = float("inf") if denom == 0 else float(
+        10.0 * np.log10(len(prediction.ravel()) * np.max(prediction.ravel()) ** 2 / denom)
+    )
+    return {
+        "PSNR": psnr,
+        "rre": float(np.linalg.norm(diff.ravel()) / np.linalg.norm(target.ravel())),
+        "SSIM": float(structural_similarity(target, prediction, data_range=target.max())),
+        "PCC": float(pearsonr(prediction.ravel(), target.ravel()).statistic),
+        "nmse": float(np.sum(diff ** 2) / np.sum(target ** 2)),
+        "mse": float(np.mean(diff ** 2) / prediction.size),
+    }
+
+
+def save_comparison(
+    target_impedance: np.ndarray,
+    prediction_impedance: np.ndarray,
+    output_path: Path,
+) -> None:
+    error = np.abs(target_impedance - prediction_impedance)
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    vmin_imp = min(target_impedance.min(), prediction_impedance.min())
+    vmax_imp = max(target_impedance.max(), prediction_impedance.max())
+    for ax, arr, title in zip(
+        axes,
+        [target_impedance, prediction_impedance, error],
+        ["Target (Impedance)", "Prediction (Impedance)", "Error (Impedance)"],
+    ):
+        if "Error" in title:
+            im = ax.imshow(arr, cmap="hot", vmin=0, vmax=error.max())
+        else:
+            im = ax.imshow(arr, cmap="jet", vmin=vmin_imp, vmax=vmax_imp)
+        ax.set_title(title)
+        ax.axis("off")
+        plt.colorbar(im, ax=ax, fraction=0.046)
+    plt.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+
+
+def evaluate_overthrust(
+    pipe: SeismicImpInvLDDPMPipeline,
+    output_dir: str | Path = "outputs/overthrust",
+    num_inference_steps: int = 1000,
+    device: str | torch.device | None = None,
+) -> dict[str, object]:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    pipe = pipe.to(device)
+
+    dataset = OverthrustTrueimpDataset(
+        size=OVERTHRUST_CONFIG["size"],
+        normalize=OVERTHRUST_CONFIG["normalize"],
+        zhengyan_type=OVERTHRUST_CONFIG["zhengyan_type"],
+        ricks=[OVERTHRUST_CONFIG["f0"]],
+        ricks_phase=[OVERTHRUST_CONFIG["f0_phase"]],
+        noise_snr=[OVERTHRUST_CONFIG["noise_snr"]],
+        dipins=[OVERTHRUST_CONFIG["dipin_v"]],
+        record_noraml=True,
+        train_keys=["image", "dipin", "record"],
+        patch_indices=OVERTHRUST_CONFIG["patch_indices"],
+        base_seed=OVERTHRUST_CONFIG["seed"],
+        data_dir=REPO_ROOT / "data",
+        cache_dir=output_dir / "cache",
+        fixed_f0=OVERTHRUST_CONFIG["f0"],
+        fixed_dipin_v=OVERTHRUST_CONFIG["dipin_v"],
+        fixed_noise_snr=OVERTHRUST_CONFIG["noise_snr"],
+        fixed_f0_phase=OVERTHRUST_CONFIG["f0_phase"],
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=OVERTHRUST_CONFIG["batch_size"],
+        shuffle=False,
+        num_workers=0,
+    )
+
+    all_predictions: list[np.ndarray] = []
+    all_targets: list[np.ndarray] = []
+    all_reconstructions: list[np.ndarray] = []
+    for batch in loader:
+        seeds = batch["seed"].tolist()
+        dipin = batch["dipin"].to(device)
+        record = batch["record"].to(device)
+        image = batch["image"].to(device)
+        output = pipe(
+            dipin=dipin,
+            record=record,
+            image=image,
+            num_inference_steps=num_inference_steps,
+            seeds=seeds,
+        )
+        prediction = output.impedance_samples
+        reconstruction = output.impedance_reconstructed
+        for local_idx in range(prediction.shape[0]):
+            all_predictions.append(prediction[local_idx, 0].detach().cpu().numpy())
+            all_targets.append(image[local_idx, 0].detach().cpu().numpy())
+            all_reconstructions.append(reconstruction[local_idx, 0].detach().cpu().numpy())
+
+    full_target = stitch_patches(
+        all_targets, dataset.splits, dataset.big_img.shape, OVERTHRUST_CONFIG["size"]
+    )
+    full_prediction = stitch_patches(
+        all_predictions, dataset.splits, dataset.big_img.shape, OVERTHRUST_CONFIG["size"]
+    )
+    full_reconstruction = stitch_patches(
+        all_reconstructions, dataset.splits, dataset.big_img.shape, OVERTHRUST_CONFIG["size"]
+    )
+
+    full_target_impedance = dataset.fan(full_target)
+    full_prediction_impedance = dataset.fan(full_prediction)
+    full_reconstruction_impedance = dataset.fan(full_reconstruction)
+
+    metrics_summary = {
+        "config": {**OVERTHRUST_CONFIG, "num_inference_steps": num_inference_steps},
+        "normalized": compute_metrics(full_prediction, full_target),
+        "impedance": compute_metrics(full_prediction_impedance, full_target_impedance),
+        "encode_impedance": compute_metrics(
+            full_reconstruction_impedance, full_target_impedance
+        ),
+    }
+
+    paths = {
+        "full_target": output_dir / "full_target.npy",
+        "full_prediction": output_dir / "full_prediction.npy",
+        "full_reconstruction": output_dir / "full_reconstruction.npy",
+        "comparison": output_dir / "comparison_impedance.png",
+        "metrics": output_dir / "metrics_summary.json",
+    }
+    np.save(paths["full_target"], full_target)
+    np.save(paths["full_prediction"], full_prediction)
+    np.save(paths["full_reconstruction"], full_reconstruction)
+    save_comparison(full_target_impedance, full_prediction_impedance, paths["comparison"])
+    paths["metrics"].write_text(json.dumps(metrics_summary, indent=2), encoding="utf-8")
+    return {
+        "metrics": metrics_summary,
+        "paths": {key: str(value) for key, value in paths.items()},
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Evaluate seismic-lddpm on Overthrust.")
+    parser.add_argument("--model", default="mally-2000/seismic-lddpm")
+    parser.add_argument("--output", default="outputs/overthrust")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--num-inference-steps", type=int, default=1000)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    pipe = SeismicImpInvLDDPMPipeline.from_pretrained(
+        args.model,
+        torch_dtype=torch.float32,
+        trust_remote_code=True,
+    )
+    result = evaluate_overthrust(
+        pipe,
+        output_dir=args.output,
+        num_inference_steps=args.num_inference_steps,
+        device=args.device,
+    )
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

inference/infer.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from pipeline import SeismicImpInvLDDPMPipeline
+
+
+def load_bchw_npy(path: str | Path) -> torch.Tensor:
+    arr = np.load(path).astype(np.float32)
+    if arr.ndim == 2:
+        arr = arr[None, None, :, :]
+    elif arr.ndim == 3:
+        arr = arr[None, :, :, :]
+    elif arr.ndim != 4:
+        raise ValueError(f"Expected 2D, 3D, or 4D array at {path}, got shape {arr.shape}")
+    return torch.from_numpy(arr)
+
+
+def save_prediction_png(prediction: np.ndarray, output_path: Path) -> None:
+    fig, ax = plt.subplots(figsize=(5, 5))
+    im = ax.imshow(prediction, cmap="jet")
+    ax.axis("off")
+    plt.colorbar(im, ax=ax, fraction=0.046)
+    plt.tight_layout()
+    fig.savefig(output_path, dpi=150)
+    plt.close(fig)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run seismic-lddpm inference on one sample.")
+    parser.add_argument("--dipin", required=True, help="Path to low-frequency impedance .npy")
+    parser.add_argument("--record", required=True, help="Path to seismic record .npy")
+    parser.add_argument("--model", default="mally-2000/seismic-lddpm")
+    parser.add_argument("--output", default="outputs/single")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--num-inference-steps", type=int, default=1000)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    pipe = SeismicImpInvLDDPMPipeline.from_pretrained(
+        args.model,
+        torch_dtype=torch.float32,
+        trust_remote_code=True,
+    ).to(device)
+    dipin = load_bchw_npy(args.dipin).to(device)
+    record = load_bchw_npy(args.record).to(device)
+    output = pipe(
+        dipin=dipin,
+        record=record,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+    )
+    prediction = output.impedance_samples.detach().cpu().numpy()
+    np.save(output_dir / "prediction.npy", prediction)
+    save_prediction_png(prediction[0, 0], output_dir / "prediction.png")
+    print(f"Saved: {output_dir / 'prediction.npy'}")
+    print(f"Saved: {output_dir / 'prediction.png'}")
+
+
+if __name__ == "__main__":
+    main()

pipeline.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from diffusers import DDPMScheduler, DiffusionPipeline, UNet2DModel, VQModel
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SeismicImpInvLDDPMPipelineOutput(BaseOutput):
+    impedance_samples: torch.Tensor | np.ndarray
+    impedance_latents: torch.Tensor | np.ndarray
+    impedance_dipin: torch.Tensor | np.ndarray
+    impedance_reconstructed: torch.Tensor | np.ndarray | None = None
+    record_features: torch.Tensor | np.ndarray | None = None
+
+
+class SeismicImpInvLDDPMPipeline(DiffusionPipeline):
+    """SAII-LDDPM impedance inversion pipeline."""
+
+    def __init__(
+        self,
+        vq_model: VQModel,
+        condition_encoder: torch.nn.Module,
+        unet: UNet2DModel,
+        scheduler: DDPMScheduler,
+    ):
+        super().__init__()
+        self.register_modules(
+            vq_model=vq_model,
+            condition_encoder=condition_encoder,
+            unet=unet,
+            scheduler=scheduler,
+        )
+
+    def _encode_conditioning(
+        self, dipin: torch.Tensor, record: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        dipin_latents = self.vq_model.encode(dipin).latents
+        if hasattr(self.condition_encoder, "encode") and callable(
+            self.condition_encoder.encode
+        ):
+            record_features = self.condition_encoder.encode(record)
+        else:
+            record_features = self.condition_encoder(record)
+        return (
+            dipin_latents.to(dtype=self.unet.dtype),
+            record_features.to(dtype=self.unet.dtype),
+        )
+
+    @staticmethod
+    def _extract_into_tensor(
+        arr: torch.Tensor, timesteps: torch.Tensor, broadcast_shape: torch.Size
+    ) -> torch.Tensor:
+        values = arr.to(device=timesteps.device, dtype=torch.float32).gather(0, timesteps)
+        return values.reshape(timesteps.shape[0], *((1,) * (len(broadcast_shape) - 1)))
+
+    @staticmethod
+    def _build_legacy_ddpm_buffers(
+        scheduler: DDPMScheduler, device: torch.device
+    ) -> dict[str, torch.Tensor]:
+        betas = scheduler.betas.to(device=device, dtype=torch.float32)
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = torch.cat(
+            [torch.ones(1, device=device), alphas_cumprod[:-1]], dim=0
+        )
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        posterior_log_variance_clipped = torch.log(
+            torch.clamp(posterior_variance, min=1e-20)
+        )
+        return {
+            "sqrt_recip_alphas_cumprod": torch.sqrt(1.0 / alphas_cumprod),
+            "sqrt_recipm1_alphas_cumprod": torch.sqrt(1.0 / alphas_cumprod - 1),
+            "posterior_mean_coef1": betas
+            * torch.sqrt(alphas_cumprod_prev)
+            / (1.0 - alphas_cumprod),
+            "posterior_mean_coef2": (1.0 - alphas_cumprod_prev)
+            * torch.sqrt(alphas)
+            / (1.0 - alphas_cumprod),
+            "posterior_log_variance_clipped": posterior_log_variance_clipped,
+        }
+
+    @staticmethod
+    def _randn_like_sample(
+        sample: torch.Tensor, generator: torch.Generator | list[torch.Generator] | None
+    ) -> torch.Tensor:
+        if isinstance(generator, list):
+            if len(generator) != sample.shape[0]:
+                raise ValueError(
+                    f"Expected {sample.shape[0]} generators, got {len(generator)}"
+                )
+            return torch.cat(
+                [
+                    torch.randn(
+                        sample[i : i + 1].shape,
+                        generator=sample_generator,
+                        device=sample.device,
+                        dtype=sample.dtype,
+                    )
+                    for i, sample_generator in enumerate(generator)
+                ],
+                dim=0,
+            )
+        return torch.randn(
+            sample.shape, generator=generator, device=sample.device, dtype=sample.dtype
+        )
+
+    def _ddpm_step(
+        self,
+        latents: torch.Tensor,
+        conditioning: torch.Tensor,
+        timestep: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+        buffers: dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        model_input = torch.cat([latents, conditioning], dim=1)
+        noise_pred = self.unet(model_input, timestep).sample
+        pred_x0 = (
+            self._extract_into_tensor(
+                buffers["sqrt_recip_alphas_cumprod"], timestep, latents.shape
+            )
+            * latents
+            - self._extract_into_tensor(
+                buffers["sqrt_recipm1_alphas_cumprod"], timestep, latents.shape
+            )
+            * noise_pred
+        )
+        pred_x0 = self.vq_model.quantize(pred_x0)[0]
+        model_mean = (
+            self._extract_into_tensor(
+                buffers["posterior_mean_coef1"], timestep, latents.shape
+            )
+            * pred_x0
+            + self._extract_into_tensor(
+                buffers["posterior_mean_coef2"], timestep, latents.shape
+            )
+            * latents
+        )
+        noise = self._randn_like_sample(latents, generator)
+        nonzero_mask = (1 - (timestep == 0).float()).reshape(
+            latents.shape[0], *((1,) * (len(latents.shape) - 1))
+        )
+        return model_mean + nonzero_mask * (
+            0.5
+            * self._extract_into_tensor(
+                buffers["posterior_log_variance_clipped"], timestep, latents.shape
+            )
+        ).exp() * noise
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        dipin: torch.Tensor,
+        record: torch.Tensor,
+        image: torch.Tensor | None = None,
+        num_inference_steps: int = 1000,
+        seed: int | None = None,
+        seeds: list[int] | tuple[int, ...] | torch.Tensor | None = None,
+        generator: torch.Generator | None = None,
+        output_type: str = "tensor",
+    ) -> SeismicImpInvLDDPMPipelineOutput:
+        device = self.unet.device
+        if seeds is not None:
+            if isinstance(seeds, torch.Tensor):
+                seeds = seeds.detach().cpu().tolist()
+            seeds = [int(value) for value in seeds]
+            if len(seeds) != dipin.shape[0]:
+                raise ValueError(f"Expected {dipin.shape[0]} seeds, got {len(seeds)}")
+            generator = [
+                torch.Generator(device=device).manual_seed(value) for value in seeds
+            ]
+        elif seed is not None:
+            generator = torch.Generator(device=device).manual_seed(seed)
+        elif generator is None:
+            generator = torch.Generator(device=device)
+
+        dipin = dipin.to(device=device, dtype=self.vq_model.dtype)
+        record = record.to(device=device, dtype=self.unet.dtype)
+        impedance_dipin, record_features = self._encode_conditioning(dipin, record)
+        conditioning = torch.cat([impedance_dipin, record_features], dim=1)
+        impedance_latents = self._randn_like_sample(
+            torch.empty(
+                impedance_dipin.shape,
+                device=device,
+                dtype=self.unet.dtype,
+            ),
+            generator,
+        )
+        buffers = self._build_legacy_ddpm_buffers(self.scheduler, device)
+        for t in reversed(range(num_inference_steps)):
+            timestep = torch.full(
+                (impedance_latents.shape[0],), t, device=device, dtype=torch.long
+            )
+            impedance_latents = self._ddpm_step(
+                impedance_latents, conditioning, timestep, generator, buffers
+            )
+
+        impedance_samples = self.vq_model.decode(
+            impedance_latents.to(dtype=self.vq_model.dtype)
+        ).sample
+        impedance_reconstructed = None
+        if image is not None:
+            image = image.to(device=device, dtype=self.vq_model.dtype)
+            image_latents = self.vq_model.encode(image).latents
+            impedance_reconstructed = self.vq_model.decode(image_latents).sample
+
+        if output_type == "np":
+            impedance_samples = impedance_samples.detach().cpu().numpy()
+            impedance_latents = impedance_latents.detach().cpu().numpy()
+            impedance_dipin = impedance_dipin.detach().cpu().numpy()
+            record_features = record_features.detach().cpu().numpy()
+            if impedance_reconstructed is not None:
+                impedance_reconstructed = impedance_reconstructed.detach().cpu().numpy()
+
+        return SeismicImpInvLDDPMPipelineOutput(
+            impedance_samples=impedance_samples,
+            impedance_latents=impedance_latents,
+            impedance_dipin=impedance_dipin,
+            impedance_reconstructed=impedance_reconstructed,
+            record_features=record_features,
+        )
+
+    @torch.no_grad()
+    def encode_decode(
+        self, image: torch.Tensor, output_type: str = "tensor"
+    ) -> torch.Tensor | np.ndarray:
+        image = image.to(device=self.vq_model.device, dtype=self.vq_model.dtype)
+        reconstruction = self.vq_model.decode(self.vq_model.encode(image).latents).sample
+        if output_type == "np":
+            return reconstruction.detach().cpu().numpy()
+        return reconstruction

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.6.0
+diffusers>=0.25.0,<0.31.0
+accelerate>=0.25.0
+safetensors>=0.4.0
+numpy>=1.23.0,<2.0
+scipy>=1.10.1
+matplotlib>=3.9.4
+scikit-image>=0.24.0
+pylops==2.2.0
+pytorch-wavelets>=1.3.0