{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# SAII-CLDM Seismic Inversion Demo\n",
        "\n",
        "This notebook demonstrates seismic impedance inversion using SAII-CLDM (Seismic-Amplitude-Guided Iterative Inversion with Cross-Domain Latent Diffusion Model).\n",
        "\n",
        "**Paper:** [arXiv:2506.13529](https://arxiv.org/html/2506.13529v1)\n",
        "\n",
        "**Repository:** https://huggingface.co/mally-2000/saii-cldm-synthetic\n",
        "\n",
        "## Two Methods:\n",
        "1. **SAII-LDDPM** - Latent Diffusion with Denoising Probabilistic Model (1000 steps)\n",
        "2. **SAII-CLDM** - Cross-Domain Latent Diffusion Model (30 steps, faster)\n",
        "\n",
        "Make sure to enable GPU runtime: Runtime → Change runtime type → Hardware accelerator → GPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 1: Clone the Repository"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!git clone https://huggingface.co/mally-2000/saii-cldm-synthetic\n",
        "%cd saii-cldm-synthetic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 2: Install Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install -r requirements.txt -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 3: Run Inference\n",
        "\n",
        "Choose which method to run:\n",
        "- `LDDPM` (default): Higher quality, slower (1000 steps)\n",
        "- `CLDM`: Faster inference (30 steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Option A: SAII-LDDPM (Slower, Higher Quality)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python infer.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Option B: SAII-CLDM (Faster)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python infer.py CLDM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 4: Visualize Results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from IPython.display import Image, display\n",
        "import os\n",
        "\n",
        "# Display comparison image from LDDPM run\n",
        "if os.path.exists(\"outputs/infer_LDDPM/comparison.png\"):\n",
        "    print(\"SAII-LDDPM Results:\")\n",
        "    display(Image(filename=\"outputs/infer_LDDPM/comparison.png\"))\n",
        "\n",
        "# Display comparison image from CLDM run\n",
        "if os.path.exists(\"outputs/infer_CLDM/comparison.png\"):\n",
        "    print(\"\\nSAII-CLDM Results:\")\n",
        "    display(Image(filename=\"outputs/infer_CLDM/comparison.png\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 5: Download Results (Optional)\n",
        "\n",
        "Download the output files to your local machine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "import glob\n",
        "\n",
        "# Download all output files\n",
        "for f in glob.glob(\"outputs/**/*\", recursive=True):\n",
        "    if os.path.isfile(f):\n",
        "        print(f\"Downloading: {f}\")\n",
        "        files.download(f)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Expected Results on Overthrust Sample\n",
        "\n",
        "| Method | Steps | PSNR | SSIM | PCC | RRE |\n",
        "|---|---:|---:|---:|---:|---:|\n",
        "| SAII-LDDPM | 1000 | 33.44 | 0.955 | 0.996 | 0.032 |\n",
        "| SAII-CLDM | 30 | 33.13 | 0.949 | 0.995 | 0.034 |\n",
        "\n",
        "The comparison image shows (left to right):\n",
        "1. Low-frequency impedance (initial model)\n",
        "2. Seismic record (observed data)\n",
        "3. Target (ground truth impedance)\n",
        "4. Prediction (inverted impedance)"
      ]
    }
  ]
}