app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import json
import time
import struct
import tempfile
from pathlib import Path

import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr
from huggingface_hub import hf_hub_download

# NOTE: `cra5` / `compressai` are imported lazily inside the @spaces.GPU worker
# (see _get_net). The prebuilt `compressai` C++ extension poisons CUDA in the
# ZeroGPU main process fork (worker_init -> torch.init fails with "No CUDA GPUs
# are available"), so it MUST NOT be imported at module scope. Importing it only
# inside the GPU worker — where real CUDA is already live — works cleanly.

# --------------------------------------------------------------------------------------
# CRA5 — Extreme Compression of ERA5
#
# CRA5 compresses the ECMWF ERA5 reanalysis (268 atmospheric variables on a
# 721 x 1440 global grid, ~1.1 GB of float32 per hourly snapshot) into a ~2.4 MB
# binary stream using an efficient Variational Transformer (VAEformer) — a ~460x
# compression ratio while preserving the physical fields.
#
# This Space loads one of those tiny .bin streams from the CRA5 dataset and
# decompresses it back into a full global weather field, then visualises the
# reconstructed variables. It demonstrates the *decoder* side of the codec:
# tiny stream in -> full physical weather field out.
# --------------------------------------------------------------------------------------

MODEL_REPO = "taohan10200/CRA5-model"
MODEL_FILE = "cra5_268v_300k.pth"
DATASET_REPO = "taohan10200/CRA5-Dataset"

DEVICE = "cuda"
CACHE_VERSION = "1"

# ---- variable / level definition (from cra5/api/cra5_268v_config.py) ------------------
PRESSURE_VNAMES = ["z", "q", "u", "v", "t", "r", "w"]
SINGLE_VNAMES = ["v10", "u10", "v100", "u100", "t2m", "tcc", "sp", "tp", "msl"]
PRESSURE_LEVELS = [
    1000, 975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 700, 650, 600, 550,
    500, 450, 400, 350, 300, 250, 225, 200, 175, 150, 125, 100, 70, 50, 30, 20,
    10, 7, 5, 3, 2, 1,
]

# Human-readable descriptions for the pressure-level variables.
VAR_LONGNAME = {
    "z": "Geopotential",
    "q": "Specific humidity",
    "u": "U wind component",
    "v": "V wind component",
    "t": "Temperature",
    "r": "Relative humidity",
    "w": "Vertical velocity",
    "v10": "10m V wind", "u10": "10m U wind",
    "v100": "100m V wind", "u100": "100m U wind",
    "t2m": "2m temperature", "tcc": "Total cloud cover",
    "sp": "Surface pressure", "tp": "Total precipitation",
    "msl": "Mean sea-level pressure",
}


def _build_channel_mapping():
    """channel index -> readable variable name, matching cra5_api ordering."""
    ch2name, name2ch = {}, {}
    idx = 0
    for v in PRESSURE_VNAMES:
        for lvl in PRESSURE_LEVELS:
            key = f"{v}_{lvl}"
            ch2name[idx] = key
            name2ch[key] = idx
            idx += 1
    for v in SINGLE_VNAMES:
        ch2name[idx] = v
        name2ch[v] = idx
        idx += 1
    return ch2name, name2ch


CH2NAME, NAME2CH = _build_channel_mapping()
N_CHANNELS = len(CH2NAME)  # 268
GRID_H, GRID_W = 721, 1440
UNCOMPRESSED_BYTES = N_CHANNELS * GRID_H * GRID_W * 4  # float32 snapshot size

# A curated menu of variables the user can visualise (label -> channel key).
DISPLAY_CHOICES = {
    "Geopotential @ 500 hPa (z500)": "z_500",
    "Temperature @ 850 hPa (t850)": "t_850",
    "Specific humidity @ 500 hPa (q500)": "q_500",
    "U wind @ 500 hPa (u500)": "u_500",
    "V wind @ 500 hPa (v500)": "v_500",
    "Relative humidity @ 500 hPa (r500)": "r_500",
    "Vertical velocity @ 500 hPa (w500)": "w_500",
    "2m temperature (t2m)": "t2m",
    "Mean sea-level pressure (msl)": "msl",
    "Total cloud cover (tcc)": "tcc",
}
DEFAULT_DISPLAY = [
    "Geopotential @ 500 hPa (z500)",
    "Temperature @ 850 hPa (t850)",
    "2m temperature (t2m)",
    "Mean sea-level pressure (msl)",
]


# ---- mean / std (for de-normalisation) ------------------------------------------------
def _load_mean_std():
    api_dir = Path(__file__).parent / "cra5" / "api"
    with open(api_dir / "mean_std.json") as f:
        ms = json.load(f)
    with open(api_dir / "mean_std_single.json") as f:
        ms_single = json.load(f)

    # level_mapping: index of each requested pressure level within the model's
    # 37-level table (here identical, all 37 levels used in order).
    mean_list, std_list = [], []
    for v in PRESSURE_VNAMES:
        for i in range(len(PRESSURE_LEVELS)):
            mean_list.append(ms["mean"][v][i])
            std_list.append(ms["std"][v][i])
    for v in SINGLE_VNAMES:
        mean_list.append(ms_single["mean"][v])
        std_list.append(ms_single["std"][v])
    mean = np.array(mean_list, dtype=np.float32)[:, None, None]
    std = np.array(std_list, dtype=np.float32)[:, None, None]
    return mean, std


MEAN_NP, STD_NP = _load_mean_std()

# ---- state-dict key renaming (verbatim from cra5/models/compressai/zoo/pretrained.py) -
def _rename_key(key):
    if key.startswith("module."):
        key = key[7:]
    if ".downsample." in key:
        return key.replace("downsample", "skip")
    if key.startswith("entropy_bottleneck."):
        if key.startswith("entropy_bottleneck._biases."):
            return f"entropy_bottleneck._bias{key[-1]}"
        if key.startswith("entropy_bottleneck._matrices."):
            return f"entropy_bottleneck._matrix{key[-1]}"
        if key.startswith("entropy_bottleneck._factors."):
            return f"entropy_bottleneck._factor{key[-1]}"
    return key


def _load_pretrained(state_dict):
    return {_rename_key(k): v for k, v in state_dict.items()}


# ---- weights (downloaded at module scope; safe — no CUDA / no compressai) -------------
print("Downloading CRA5 VAEformer weights ...")
CKPT_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
print(f"Weights cached at {CKPT_PATH}")

# The model + de-normalisation tensors are built lazily inside the GPU worker.
_NET = None          # VAEformer on cuda
_MEAN = None         # mean tensor on cuda
_STD = None          # std tensor on cuda


def _get_net():
    """Build (once) and return the VAEformer on the GPU.

    Imports of `cra5` / `compressai` happen HERE, inside the @spaces.GPU worker,
    because the compressai C++ extension breaks CUDA init if imported in the
    ZeroGPU main process (see note at top of file).
    """
    global _NET, _MEAN, _STD
    if _NET is not None:
        return _NET, _MEAN, _STD

    from cra5.models.vaeformer.vaeformer import VAEformer

    state_dict = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
    if isinstance(state_dict, dict) and "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]
    # Match the reference load path (zoo._load_model): rename keys before building.
    state_dict = _load_pretrained(state_dict)

    net = VAEformer.from_state_dict(state_dict).eval()
    net.update(force=True)  # build entropy-coder CDF tables for (de)compression
    net = net.to(DEVICE)

    _NET = net
    _MEAN = torch.from_numpy(MEAN_NP).to(DEVICE)
    _STD = torch.from_numpy(STD_NP).to(DEVICE)
    print("Model ready on GPU.")
    return _NET, _MEAN, _STD


# ---- binary stream I/O (matches cra5/api/utils.py + cra5_api.decode_from_bin) ----------
def _read_uints(fd, n, fmt=">{:d}I"):
    sz = struct.calcsize("I")
    return struct.unpack(fmt.format(n), fd.read(n * sz))


def _read_bytes(fd, n, fmt=">{:d}s"):
    sz = struct.calcsize("s")
    return struct.unpack(fmt.format(n), fd.read(n * sz))[0]


def _parse_bin(bin_path):
    """Read a CRA5 .bin stream into (strings, latent_shape)."""
    with Path(bin_path).open("rb") as f:
        shape = _read_uints(f, 2)  # z spatial shape (H', W')
        n_strings = _read_uints(f, 1)[0]
        lstrings = []
        for _ in range(n_strings):
            s = _read_bytes(f, _read_uints(f, 1)[0])
            lstrings.append([s])
    return lstrings, shape


# --------------------------------------------------------------------------------------
# Fetching sample streams from the CRA5 dataset
# --------------------------------------------------------------------------------------
def _bin_path_for(timestamp: str) -> str:
    year = timestamp[:4]
    return hf_hub_download(
        repo_id=DATASET_REPO,
        filename=f"{year}/{timestamp}.bin",
        repo_type="dataset",
    )


# --------------------------------------------------------------------------------------
# Visualisation
# --------------------------------------------------------------------------------------
def _render(x_hat_np, timestamp, display_labels):
    keys = [DISPLAY_CHOICES[l] for l in display_labels]
    n = len(keys)
    fig, axs = plt.subplots(n, 1, figsize=(11, 3.1 * n), squeeze=False)
    extent = [-180, 180, -90, 90]
    for i, key in enumerate(keys):
        ch = NAME2CH[key]
        field = x_hat_np[ch]
        ax = axs[i, 0]
        im = ax.imshow(field, cmap="jet", extent=extent, aspect="auto")
        base = key.split("_")[0]
        lname = VAR_LONGNAME.get(base, key)
        lvl = key.split("_")[1] + " hPa" if "_" in key else "surface"
        ax.set_title(f"{lname}  ({key})   —   reconstructed from {timestamp}", fontsize=11)
        ax.set_xlabel("longitude"); ax.set_ylabel("latitude")
        fig.colorbar(im, ax=ax, fraction=0.025, pad=0.02)
    plt.tight_layout()
    out = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    fig.savefig(out.name, dpi=90, bbox_inches="tight")
    plt.close(fig)
    return out.name


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=60)
def decompress_and_visualize(timestamp: str, display_labels: list, progress=gr.Progress()):
    """Decompress a CRA5 binary stream into a global weather field and visualise it.

    Args:
        timestamp: An ERA5 timestamp available in the CRA5 dataset, e.g.
            "2023-06-01T00:00:00". The matching ~2.4 MB .bin stream is fetched
            from the CRA5 dataset and decoded by the VAEformer.
        display_labels: Which reconstructed weather variables to plot.

    Returns:
        A rendered figure of the reconstructed variables, a compression-stats
        summary, and a per-variable statistics table.
    """
    if not display_labels:
        display_labels = DEFAULT_DISPLAY

    progress(0.05, desc="Loading model onto GPU ...")
    net, mean, std = _get_net()  # builds + moves to GPU on first call (inside worker)

    progress(0.2, desc="Fetching compressed stream ...")
    bin_path = _bin_path_for(timestamp)
    bin_bytes = os.path.getsize(bin_path)

    progress(0.35, desc="Parsing binary stream ...")
    lstrings, shape = _parse_bin(bin_path)

    progress(0.5, desc="Decoding with VAEformer ...")
    t0 = time.time()
    with torch.no_grad():
        out = net.decompress(lstrings, shape)  # {"x_hat": normalized field}
        x_hat = out["x_hat"].squeeze(0)            # (268, 721, 1440) normalized
        x_hat = x_hat * std + mean                 # de-normalise to physical units
    x_hat_np = x_hat.detach().float().cpu().numpy()
    decode_time = time.time() - t0

    progress(0.85, desc="Rendering maps ...")
    fig_path = _render(x_hat_np, timestamp, display_labels)

    ratio = UNCOMPRESSED_BYTES / bin_bytes
    summary = (
        f"### Compression summary — `{timestamp}`\n"
        f"- **Compressed stream:** {bin_bytes/1e6:.2f} MB\n"
        f"- **Uncompressed snapshot:** {UNCOMPRESSED_BYTES/1e9:.2f} GB "
        f"({N_CHANNELS} variables × {GRID_H} × {GRID_W} float32)\n"
        f"- **Compression ratio:** ≈ **{ratio:.0f}×**\n"
        f"- **Latent grid:** {shape[0]} × {shape[1]}\n"
        f"- **Decode time (GPU):** {decode_time:.2f} s"
    )

    # small per-variable stats table
    rows = []
    for label in display_labels:
        key = DISPLAY_CHOICES[label]
        f = x_hat_np[NAME2CH[key]]
        rows.append([key, f"{f.min():.2f}", f"{f.mean():.2f}", f"{f.max():.2f}"])

    return fig_path, summary, rows


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1150px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

SAMPLE_TIMESTAMPS = [
    "2023-06-01T00:00:00",
    "2023-01-01T00:00:00",
    "2023-09-15T12:00:00",
    "2023-12-25T00:00:00",
    "1979-01-01T00:00:00",
]

with gr.Blocks(title="CRA5 — Extreme ERA5 Compression") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🌍 CRA5 — Extreme Compression of ERA5\n"
            "**VAEformer** compresses a full global ERA5 weather snapshot "
            "(268 variables on a 721×1440 grid, ~1.1 GB) into a **~2.4 MB** binary "
            "stream — a **~460× compression ratio**. This demo fetches one of those "
            "tiny streams from the [CRA5 dataset](https://huggingface.co/datasets/taohan10200/CRA5-Dataset), "
            "**decompresses it on the GPU**, and visualises the reconstructed fields.\n\n"
            "Model: [taohan10200/CRA5-model](https://huggingface.co/taohan10200/CRA5-model) · "
            "Paper: [arXiv:2405.03376](https://arxiv.org/abs/2405.03376) · "
            "Code: [github.com/taohan10200/CRA5](https://github.com/taohan10200/CRA5)"
        )

        with gr.Row():
            timestamp = gr.Dropdown(
                choices=SAMPLE_TIMESTAMPS,
                value=SAMPLE_TIMESTAMPS[0],
                label="ERA5 timestamp (from the CRA5 dataset)",
                allow_custom_value=True,
                scale=3,
            )
            run = gr.Button("Decompress & visualize", variant="primary", scale=1)

        display = gr.CheckboxGroup(
            choices=list(DISPLAY_CHOICES.keys()),
            value=DEFAULT_DISPLAY,
            label="Variables to visualise",
        )

        out_image = gr.Image(label="Reconstructed global fields", type="filepath")
        out_summary = gr.Markdown()
        out_table = gr.Dataframe(
            headers=["variable", "min", "mean", "max"],
            label="Reconstructed variable statistics (physical units)",
            interactive=False,
        )

        gr.Markdown(
            "ℹ️ Custom timestamps: any hourly ERA5 time from **1979–2023** present in the "
            "CRA5 dataset works, formatted like `2023-06-01T00:00:00`."
        )

        gr.Examples(
            examples=[
                ["2023-06-01T00:00:00", DEFAULT_DISPLAY],
                ["2023-12-25T00:00:00", ["2m temperature (t2m)", "Mean sea-level pressure (msl)", "Total cloud cover (tcc)"]],
                ["1979-01-01T00:00:00", ["Geopotential @ 500 hPa (z500)", "Temperature @ 850 hPa (t850)"]],
            ],
            inputs=[timestamp, display],
            outputs=[out_image, out_summary, out_table],
            fn=decompress_and_visualize,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        decompress_and_visualize,
        inputs=[timestamp, display],
        outputs=[out_image, out_summary, out_table],
        api_name="decompress",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)