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)