feat: e2s demo with FCN on HF ZeroGPU

.gitignore

@@ -0,0 +1,4 @@
+.venv/
+__pycache__/
+outputs/
+OLD_DOCKER/

Earth2Studio_P1.ipynb

@@ -0,0 +1,428 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Introduction to Earth2Studio \n",
+    "\n",
+    "In this notebook, we will introduce the Earth2Studio Python package and run a example that will demonstrate how to run a simple inference workflow to generate a basic determinstic forecast using one of the built in models of Earth-2 Inference Studio.\n",
+    "\n",
+    "\n",
+    "#### Contents of the Notebook\n",
+    "\n",
+    "- [Earth2Studio](#Earth2Studio)\n",
+    "- [Simple Deterministic Inference](#Simple-Deterministic-Inference)\n",
+    "    - [Set Up](#Set-Up)\n",
+    "    - [Execute the Workflow](#Execute-the-Workflow)\n",
+    "    - [Post Processing](#Post-Processing)\n",
+    "\n",
+    "#### Learning Outcomes\n",
+    "\n",
+    "- Earth2Studio Key Features\n",
+    "- How to instantiate a built in prognostic model\n",
+    "- Creating a data source and IO object\n",
+    "- Running a simple built in workflow\n",
+    "- Post-processing results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Earth2Studio\n",
+    "\n",
+    "Earth2Studio is a Python package built to empower researchers, scientists, and enthusiasts in the fields of weather and climate science with the latest artificial intelligence models/capabilities. With an intuitive design and comprehensive feature set, this package serves as a robust toolkit for exploring this AI revolution in the weather and climate science domain.\n",
+    "\n",
+    "### Package Design\n",
+    "\n",
+    "The goal of this package is to enable the use to extrapolate and build beyond what is implemented in it. The design philosophy of Earth2Studio embodies a modular architecture where the inference workflow acts as a flexible adhesive, seamlessly binding together various specialized software components with well-defined interfaces. Each component within the package serves a distinct purpose in typical inference workflows.\n",
+    "\n",
+    "<div style=\"display: flex; justify-content: center; gap: 10px;\">\n",
+    "  <figure style=\"text-align: center;\">\n",
+    "    <img src=\"https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/arch.png\" style=\"width: 100%; height: auto;\">\n",
+    "    <figcaption>Model architecture overview.</figcaption>\n",
+    "  </figure>\n",
+    "</div>\n",
+    "\n",
+    "By viewing the inference workflow as a dynamic connector, Earth2Studio facilitates effortless integration of these components, allowing researchers to easily swap out or augment functionalities to suit their specific needs. We recognize that many users will have their own custom workflow needs, thus encourage users to use the provided features as a starting point to build their own.\n",
+    "\n",
+    "<div style=\"display: flex; justify-content: center; gap: 10px;\">\n",
+    "  <figure style=\"text-align: center;\">\n",
+    "    <img src=\"https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/samples.png\" style=\"width: 100%; height: auto;\">\n",
+    "  </figure>\n",
+    "</div>\n",
+    "\n",
+    "Significant importance is placed on the interface that enables the connection between the components and the workflow. These are simple python protocols that all variants of a particular component must share. This not only enables a consistent API but also the generalization of workflows.\n",
+    "\n",
+    "### Key Features\n",
+    "\n",
+    "While Earth2Studio contains a large collection of general utilities, functions and tooling the following six are considered the core. For more information on these features, see the dedicated documentation for each.\n",
+    "\n",
+    "- **Built-in Workflows**: Multiple built-in inference workflows to accelerate your development and research.\n",
+    "- **Prognostic Models**: Support for the latest AI weather forecast models offered under a coherent interface.\n",
+    "- **Diagnostic Models**: Diagnostic models for mapping to other quantities of interest.\n",
+    "- **Datasources**: Datasources to connect on-prem and remote data stores to inference workflows.\n",
+    "- **IO**: Simple, yet powerful IO utilities to export data for post-processing.\n",
+    "- **Statistical Operators**: Statistical methods to fuse directly into your inference workflow for more complex uncertainty analysis.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Simple Deterministic Inference\n",
+    "\n",
+    "<div style=\"display: flex; justify-content: center; gap: 10px;\">\n",
+    "  <figure style=\"text-align: center;\">\n",
+    "    <img src=\"https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/deterministic.png\" style=\"width: 100%; height: auto;\">\n",
+    "  </figure>\n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Set Up**\n",
+    "All workflows inside Earth2Studio require constructed components to be  handed to them. In this example, let's take a look at the most basic: `earth2studio.run.deterministic`.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let us look at the built in Models, Datasource and IO Backends that are avaialable with Earth2Studio `0.2.0`release. \n",
+    "\n",
+    "\n",
+    "#### Prognostic Model: \n",
+    "\n",
+    "Prognostic models are a class of models that perform time-integration. Thus are typically used to generate forecast predictions.\n",
+    "\n",
+    "The list of Prognostic Models available as of `0.2.0` are: \n",
+    "- **models.px.DLWP** : Deep learning weather prediction (DLWP) prognostic model.\n",
+    "- **models.px.FCN** : FourCastNet global prognostic model.\n",
+    "- **models.px.FengWu** :  FengWu (operational) weather model consists of single auto-regressive model with a time-step size of 6 hours.\n",
+    "- **models.px.FuXi** : FuXi weather model consists of three auto-regressive U-net transfomer models with a time-step size of 6 hours.\n",
+    "- **models.px.Pangu24** : Pangu Weather 24 hour model.\n",
+    "- **models.px.Pangu6** : Pangu Weather 6 hour model.\n",
+    "- **models.px.Pangu3** : Pangu Weather 3 hour model.\n",
+    "- **models.px.Persistence** : Persistence model that generates a forecast by applying the identity operator on the initial condition and indexing the lead time by 6 hours.\n",
+    "- **models.px.SFNO** : Spherical Fourier Operator Network global prognostic model.\n",
+    "\n",
+    "\n",
+    "#### Data source : \n",
+    "\n",
+    "Data sources used for downloading, caching and reading different weather / climate data APIs into Xarray data arrays. Used for fetching initial conditions for inference and validation data for scoring.\n",
+    "\n",
+    "The list of Datasources available as of `0.2.0` are: \n",
+    "- **data.ARCO** : Analysis-Ready, Cloud Optimized (ARCO) is a data store of ERA5 re-analysis data currated by Google.\n",
+    "- **data.CDS** : The climate data source (CDS) serving ERA5 re-analysis data.\n",
+    "- **data.GFS** : The global forecast service (GFS) initial state data source provided on an equirectangular grid.\n",
+    "- **data.HRRR** : High-Resolution Rapid Refresh (HRRR) is a North-American weather forecast model with hourly data-assimilation developed by NOAA.\n",
+    "- **data.IFS** : The integrated forecast system (IFS) initial state data source provided on an equirectangular grid.\n",
+    "- **data.IMERG** : The Integrated Multi-satellitE Retrievals (IMERG) for GPM.\n",
+    "- **data.Random(domain_coords)** : A randomly generated normally distributed data.\n",
+    "- **data.WB2ERA5** : ERA5 reanalysis data with several derived variables on a 0.25 degree lat-lon grid from 1959 to 2023 (incl) to 6 hour intervals on 13 pressure levels.\n",
+    "- **data.WB2ERA5_121x240** : ERA5 reanalysis data with several derived variables down sampled to a 1.5 degree lat-lon grid from 1959 to 2023 (incl) to 6 hour intervals on 13 pressure levels.\n",
+    "- **data.WB2ERA5_32x64** : ERA5 reanalysis data with several derived variables down sampled to a 5.625 degree lat-lon grid from 1959 to 2023 (incl) to 6 hour intervals on 13 pressure levels.\n",
+    "- **data.WB2Climatology** : Climatology provided by WeatherBench2,\n",
+    "- **data.DataArrayFile** : A local xarray dataarray file data source.\n",
+    "- **data.DataSetFile** : A local xarray dataset file data source.\n",
+    "\n",
+    "#### IO Backend: \n",
+    "\n",
+    "The IO Backends for used for saving the inference results for further post processing.\n",
+    "\n",
+    "The list of IO Backends available as of `0.2.0` are: \n",
+    "- **io.KVBackend** : A key-value (dict) backend.\n",
+    "- **io.NetCDF4Backend** : A backend that supports the NetCDF4 format.\n",
+    "- **io.XarrayBackend** : An xarray backed IO object.\n",
+    "- **io.ZarrBackend** : A backend that supports the zarr format.\n",
+    "\n",
+    "\n",
+    "For this example, we will be using the following:\n",
+    "\n",
+    "- **Prognostic Model**: Use the built in FourCastNet Model :`earth2studio.models.px.FCN`.\n",
+    "- **Datasource**: Pull data from the GFS data api :`earth2studio.data.GFS`.\n",
+    "- **IO Backend**: Let's save the outputs into a Zarr store :`earth2studio.io.ZarrBackend`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ['EARTH2STUDIO_CACHE'] = os.getcwd() + \"/outputs/cache\"\n",
+    "os.makedirs(\"outputs\", exist_ok=True)\n",
+    "from dotenv import load_dotenv\n",
+    "load_dotenv()\n",
+    "\n",
+    "from earth2studio.data import GFS\n",
+    "from earth2studio.io import ZarrBackend\n",
+    "from earth2studio.models.px import FCN\n",
+    "\n",
+    "# Prognostic Model - Load the default model package which downloads the check point from NGC\n",
+    "package = FCN.load_default_package()\n",
+    "model = FCN.load_model(package)\n",
+    "\n",
+    "# Data Source - Create the data source\n",
+    "data = GFS()\n",
+    "\n",
+    "# IO Backend - Create the IO handler, store in memory\n",
+    "io = ZarrBackend()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Execute the Workflow**\n",
+    "\n",
+    "With all components initialized, running the workflow is a single line of Python code. \n",
+    "Workflow will return the provided IO object back to the user, which can be used to\n",
+    "then post process. Let us look at the API for Determinstic inference\n",
+    "\n",
+    "```python \n",
+    "\n",
+    "def deterministic(\n",
+    "    time: list[str] | list[datetime] | list[np.datetime64],\n",
+    "    nsteps: int,\n",
+    "    prognostic: PrognosticModel,\n",
+    "    data: DataSource,\n",
+    "    io: IOBackend,\n",
+    "    output_coords: CoordSystem = OrderedDict({}),\n",
+    "    device: torch.device | None = None,\n",
+    ") -> IOBackend:\n",
+    "    \"\"\"Built in deterministic workflow.\n",
+    "    This workflow creates a determinstic inference pipeline to produce a forecast\n",
+    "    prediction using a prognostic model.\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    time : list[str] | list[datetime] | list[np.datetime64]\n",
+    "        List of string, datetimes or np.datetime64\n",
+    "    nsteps : int\n",
+    "        Number of forecast steps\n",
+    "    prognostic : PrognosticModel\n",
+    "        Prognostic model\n",
+    "    data : DataSource\n",
+    "        Data source\n",
+    "    io : IOBackend\n",
+    "        IO object\n",
+    "    output_coords: CoordSystem, optional\n",
+    "        IO output coordinate system override, by default OrderedDict({})\n",
+    "    device : torch.device, optional\n",
+    "        Device to run inference on, by default None\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    IOBackend\n",
+    "        Output IO object\n",
+    "    \"\"\"\n",
+    "```\n",
+    "\n",
+    "For the forecast we will predict for 20 forecast steps which is 5 days."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import earth2studio.run as run\n",
+    "\n",
+    "nsteps = 20 # Each step has a lead time of 6 hours. \n",
+    "io = run.deterministic([\"2024-01-01\"], nsteps, model, data, io)\n",
+    "\n",
+    "print(io.root.tree())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Post Processing**\n",
+    "The last step is to post process our results. Cartopy is a great library for plotting\n",
+    "fields on projections of a sphere. Here we will just plot the temperature at 2 meters\n",
+    "(t2m) 1 day into the forecast.\n",
+    "\n",
+    "Notice that the Zarr IO function has additional APIs to interact with the stored data.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import cartopy.crs as ccrs\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "forecast = \"2024-01-01\"\n",
+    "variable = \"t2m\"\n",
+    "step = 4  # lead time = 4 x 6 = 24 hrs\n",
+    "\n",
+    "plt.close(\"all\")\n",
+    "# Create a Robinson projection\n",
+    "projection = ccrs.Robinson()\n",
+    "\n",
+    "# Create a figure and axes with the specified projection\n",
+    "fig, ax = plt.subplots(subplot_kw={\"projection\": projection}, figsize=(10, 6))\n",
+    "\n",
+    "# Plot the field using pcolormesh\n",
+    "im = ax.pcolormesh(\n",
+    "    io[\"lon\"][:],\n",
+    "    io[\"lat\"][:],\n",
+    "    io[variable][0, step],\n",
+    "    transform=ccrs.PlateCarree(),\n",
+    "    cmap=\"Spectral_r\",\n",
+    ")\n",
+    "\n",
+    "# Set title\n",
+    "ax.set_title(f\"{forecast} - Lead time: {6*step}hrs\")\n",
+    "\n",
+    "# Add coastlines and gridlines\n",
+    "ax.coastlines()\n",
+    "ax.gridlines()\n",
+    "plt.savefig(\"outputs/01_t2m_prediction.jpg\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let us now create a simple GIF that would go through all the steps using the below script.\n",
+    "\n",
+    "Kindly note, the below script would take approximately 10 minutes to show the output. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.animation as animation\n",
+    "from IPython.display import HTML\n",
+    "\n",
+    "forecast = \"2024-01-01\"\n",
+    "variable = \"t2m\"\n",
+    "num_timesteps = 20  # Number of time steps to create GIF\n",
+    "\n",
+    "# Create a Robinson projection\n",
+    "projection = ccrs.Robinson()\n",
+    "\n",
+    "# Create a figure and axes with the specified projection\n",
+    "fig, ax = plt.subplots(subplot_kw={\"projection\": projection}, figsize=(10, 6))\n",
+    "\n",
+    "# Create a function to update the frame\n",
+    "def update(frame):\n",
+    "    ax.clear()  # Clear the axis for new plot\n",
+    "    im = ax.pcolormesh(\n",
+    "        io[\"lon\"][:],\n",
+    "        io[\"lat\"][:],\n",
+    "        io[variable][0, frame],\n",
+    "        transform=ccrs.PlateCarree(),\n",
+    "        cmap=\"Spectral_r\",\n",
+    "    )\n",
+    "    ax.set_title(f\"{forecast} - Lead time: {6 * frame } hrs\")\n",
+    "    ax.coastlines()\n",
+    "    ax.gridlines()\n",
+    "    return im,\n",
+    "\n",
+    "# Create an animation\n",
+    "ani = animation.FuncAnimation(fig, update, frames=num_timesteps, blit=False)\n",
+    "\n",
+    "# Save as GIF & Display them - Kindly note, this cell takes around 10 minutes or more to display the output\n",
+    "ani.save('outputs/t2m_prediction_animation.gif')\n",
+    "print(\"Animation of 20 Timesteps\")\n",
+    "HTML(ani.to_html5_video())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Important: Free up GPU Memory!\n",
+    "\n",
+    "Run the below cell to free up GPU memory after training the model before moving to the next notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os._exit(00)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we had looked at the plot of t2m at the 4th step ( Each step is 6 hours in FCN Model ). In the Next notebook, let us extend this by adding a Diagnostic Model with FCN."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "--- \n",
+    "\n",
+    "Don't forget to check out additional [Open Hackathons Resources](https://www.openhackathons.org/s/technical-resources) and join our [OpenACC and Hackathons Slack Channel](https://www.openacc.org/community#slack) to share your experience and get more help from the community.\n",
+    "\n",
+    "---\n",
+    "\n",
+    "# Licensing\n",
+    "\n",
+    "Copyright © 2023 OpenACC-Standard.org.  This material is released by OpenACC-Standard.org, in collaboration with NVIDIA Corporation, under the Creative Commons Attribution 4.0 International (CC BY 4.0). These materials may include references to hardware and software developed by other entities; all applicable licensing and copyrights apply.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

README.md

@@ -0,0 +1,19 @@
+---
+title: Earth2 Inference Demo
+colorFrom: green
+colorTo: gray
+sdk: gradio
+sdk_version: 6.3.0
+python_version: "3.12.12"
+app_file: app.py
+app_port: 7860
+pinned: false
+tags:
+  - physics
+  - cfd
+  - machine-learning
+  - neural-operators
+  - fluid-dynamics
+  - scientific-computing
+  - earth2
+---

app.py

@@ -0,0 +1,598 @@
+import asyncio
+import os
+from typing import Optional, Tuple
+
+# Save original asyncio.run BEFORE any imports that might patch it (nest_asyncio)
+_ORIGINAL_ASYNCIO_RUN = asyncio.run
+
+# On ZeroGPU (shared A10G), TF32 matmul paths can occasionally trip cuBLAS errors in
+# some einsum-heavy models. Prefer full FP32 math for stability.
+os.environ.setdefault("NVIDIA_TF32_OVERRIDE", "0")
+# ZeroGPU H200-specific workarounds for cuBLAS strided-batch GEMM issues
+# H200 has 70GB VRAM, so memory isn't the issue - focus on CUDA context stability
+# - Force synchronous CUDA execution to avoid race conditions during dynamic GPU allocation
+# - Use deterministic cuBLAS workspace to ensure consistent behavior across GPU allocations
+os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":16:8")
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+
+# ZeroGPU decorator - only import on Hugging Face Spaces to avoid asyncio conflicts locally
+def _make_spaces_fallback():
+    class _SpacesFallback:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def _decorator(fn):
+                return fn
+            return _decorator
+    return _SpacesFallback()
+
+if os.environ.get("SPACE_ID"):
+    # Running on Hugging Face Spaces
+    try:
+        import spaces  # type: ignore
+    except Exception:
+        spaces = _make_spaces_fallback()  # type: ignore
+else:
+    # Local development - skip spaces import to avoid asyncio conflicts
+    spaces = _make_spaces_fallback()  # type: ignore
+
+
+def _ensure_cache_dirs() -> None:
+    os.makedirs("outputs", exist_ok=True)
+    os.makedirs(os.path.join("outputs", "cache"), exist_ok=True)
+    os.environ.setdefault("EARTH2STUDIO_CACHE", os.path.join(os.getcwd(), "outputs", "cache"))
+
+
+def _normalize_to_uint8(x: np.ndarray) -> np.ndarray:
+    x = np.asarray(x, dtype=np.float32)
+    finite = np.isfinite(x)
+    if not finite.any():
+        return np.zeros_like(x, dtype=np.uint8)
+    vmin = float(np.nanpercentile(x[finite], 2.0))
+    vmax = float(np.nanpercentile(x[finite], 98.0))
+    if vmax <= vmin:
+        return np.zeros_like(x, dtype=np.uint8)
+    y = (x - vmin) / (vmax - vmin)
+    y = np.clip(y, 0.0, 1.0)
+    return (y * 255.0).astype(np.uint8)
+
+
+def _apply_simple_colormap(u8: np.ndarray) -> np.ndarray:
+    """
+    Lightweight colormap without matplotlib:
+    map grayscale -> RGB using a simple blue->cyan->yellow->red ramp.
+    """
+    u = u8.astype(np.float32) / 255.0
+    r = np.clip(1.5 * u, 0.0, 1.0)
+    g = np.clip(1.5 * (1.0 - np.abs(u - 0.5) * 2.0), 0.0, 1.0)
+    b = np.clip(1.5 * (1.0 - u), 0.0, 1.0)
+    rgb = np.stack([r, g, b], axis=-1)
+    return (rgb * 255.0).astype(np.uint8)
+
+
+def _plot_latlon_field(lon: np.ndarray, lat: np.ndarray, field2d: np.ndarray, title: str) -> str:
+    """
+    Save a quick image to outputs/ and return the file path.
+    Avoids matplotlib/cartopy to keep system deps minimal on Spaces.
+    """
+    _ensure_cache_dirs()
+
+    out_path = os.path.join("outputs", "t2m.png")
+    gray = _normalize_to_uint8(field2d)
+    rgb = _apply_simple_colormap(gray)
+    img = Image.fromarray(rgb, mode="RGB").resize((1024, 512), resample=Image.BILINEAR)
+
+    draw = ImageDraw.Draw(img)
+    text = title
+    try:
+        font = ImageFont.load_default()
+    except Exception:
+        font = None
+    # simple text background for readability
+    pad = 6
+    tw, th = draw.textbbox((0, 0), text, font=font)[2:]
+    draw.rectangle((0, 0, tw + 2 * pad, th + 2 * pad), fill=(0, 0, 0))
+    draw.text((pad, pad), text, fill=(255, 255, 255), font=font)
+
+    img.save(out_path)
+    return out_path
+
+
+def _gpu_duration(nsteps: int) -> int:
+    """
+    Calculate GPU duration for inference only.
+    """
+    nsteps = max(1, int(nsteps))
+    # 30s base (model to GPU) + 15s per step
+    return int(min(300, 30 + nsteps * 15))
+
+
+@spaces.GPU(duration=lambda forecast_date, nsteps: _gpu_duration(int(nsteps)))
+def _run_inference(forecast_date: str, nsteps: int):
+    """
+    GPU-only function: load model, run inference, return extracted data.
+    
+    ZeroGPU uses multiprocessing so we can't pass unpicklable objects (GFS, model).
+    Everything must be created inside this function.
+    """
+    import torch
+    import earth2studio.run as run
+    from earth2studio.data import GFS
+    from earth2studio.io import ZarrBackend
+    
+    _ensure_cache_dirs()
+    
+    # Memory management for ZeroGPU
+    torch.backends.cudnn.benchmark = False  # More stable on shared GPU
+    # Prefer full FP32 math (avoid TF32) for stability on shared A10G
+    try:
+        torch.set_float32_matmul_precision("highest")
+    except Exception:
+        pass
+    try:
+        torch.backends.cuda.matmul.allow_tf32 = False
+    except Exception:
+        pass
+    try:
+        torch.backends.cudnn.allow_tf32 = False
+    except Exception:
+        pass
+    # Avoid reduced-precision reductions (guarded for older torch versions)
+    try:
+        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    except Exception:
+        pass
+    try:
+        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+    except Exception:
+        pass
+    try:
+        torch.cuda.set_device(0)
+    except Exception:
+        pass
+    torch.cuda.empty_cache()
+    
+    # Some cuBLAS "INVALID_VALUE" failures originate from non-contiguous einsum operands
+    # producing unsupported strides for batched GEMM. Force contiguity at the einsum boundary.
+    _orig_einsum = torch.einsum
+
+    def _einsum_contiguous(equation, *operands):
+        ops = []
+        for op in operands:
+            if isinstance(op, torch.Tensor) and not op.is_contiguous():
+                ops.append(op.contiguous())
+            else:
+                ops.append(op)
+        return _orig_einsum(equation, *ops)
+
+    torch.einsum = _einsum_contiguous  # type: ignore[assignment]
+
+    # Load model inside GPU function (ZeroGPU requirement)
+    try:
+        from earth2studio.models.px.fcn import FCN
+    except Exception:
+        from earth2studio.models.px import FCN
+    
+    package = FCN.load_default_package()
+    model = FCN.load_model(package)
+    
+    device = torch.device("cuda")
+    # Ensure FP32 weights for cuBLAS stability
+    try:
+        model = model.float()
+    except Exception:
+        pass
+    model = model.to(device)
+    model.eval()  # Ensure eval mode
+    
+    # Clear memory after model load
+    torch.cuda.empty_cache()
+    
+    # CRITICAL: Warmup CUDA/cuBLAS context on ZeroGPU's H200 before complex ops
+    # This ensures cuBLAS is fully initialized and strided-batch GEMM handlers are ready
+    try:
+        with torch.no_grad():
+            # Create dummy tensors matching FCN's expected input shape
+            # FCN expects (batch, channels, lat, lon) - use minimal batch/size for warmup
+            dummy_input = torch.randn(1, 73, 8, 8, device=device, dtype=torch.float32)
+            _ = model(dummy_input)
+            torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+    except Exception as warmup_err:
+        # If warmup fails, log but continue - the actual inference might still work
+        print(f"[Warning] CUDA warmup failed: {warmup_err}")
+    
+    data = GFS()
+    io = ZarrBackend()
+    
+    try:
+        with torch.no_grad():
+            io = run.deterministic([forecast_date], nsteps, model, data, io, device=device)
+        
+        # Extract ALL timesteps to numpy arrays (picklable) before returning
+        lon = np.asarray(io["lon"][:])
+        lat = np.asarray(io["lat"][:])
+        # Return all timesteps: shape (1, nsteps+1, lat, lon)
+        all_fields = np.asarray(io["t2m"][:])
+        
+        return lon, lat, all_fields
+    finally:
+        # Restore einsum in case this worker is reused
+        try:
+            torch.einsum = _orig_einsum  # type: ignore[assignment]
+        except Exception:
+            pass
+        # Free GPU memory aggressively
+        try:
+            model.to("cpu")
+        except Exception:
+            pass
+        try:
+            del model
+            del data
+            del io
+        except Exception:
+            pass
+        try:
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        except Exception:
+            pass
+
+
+def run_forecast(forecast_date: str, nsteps: int):
+    """
+    Run Earth2Studio deterministic inference and return cached results.
+    Returns: (forecast_date, nsteps, lon, lat, all_fields, status_msg)
+    """
+    _ensure_cache_dirs()
+
+    # Validate inputs
+    if not forecast_date:
+        return None, None, None, None, None, "ERROR: forecast_date is required (YYYY-MM-DD)."
+
+    nsteps = int(nsteps)
+    if nsteps < 1:
+        return None, None, None, None, None, "ERROR: nsteps must be >= 1"
+
+    # Run inference on GPU (model loaded inside due to ZeroGPU pickling)
+    try:
+        lon, lat, all_fields = _run_inference(forecast_date, nsteps)
+    except Exception as e:
+        return None, None, None, None, None, f"ERROR during inference: {type(e).__name__}: {e}"
+
+    # Return cached data for dynamic plot_step updates
+    status = f"SUCCESS: Computed {nsteps} forecast steps ({(nsteps+1)*6} hours total). Use plot_step slider to explore."
+    return forecast_date, nsteps, lon, lat, all_fields, status
+
+
+def update_plot_from_cache(forecast_date, nsteps, lon, lat, all_fields, plot_step):
+    """
+    Update the displayed plot from cached inference results (no GPU needed).
+    """
+    if lon is None or lat is None or all_fields is None:
+        return None, "No cached results. Click 'Run Inference' first."
+    
+    plot_step = int(plot_step)
+    nsteps = int(nsteps)
+    
+    # Validate plot_step
+    if plot_step < 0 or plot_step > nsteps:
+        return None, f"Invalid plot_step {plot_step} (must be 0-{nsteps})"
+    
+    # Extract the specific timestep
+    field = all_fields[0, plot_step]
+    
+    # Plot
+    img_path = _plot_latlon_field(
+        lon,
+        lat,
+        field,
+        title=f"{forecast_date} - t2m - lead={6 * plot_step}h",
+    )
+    return img_path, f"Displaying step {plot_step} (lead time: {6 * plot_step} hours)"
+
+
+def build_ui() -> gr.Blocks:
+    with gr.Blocks(title="Earth2Studio FCN (ZeroGPU)") as demo:
+        gr.Markdown(
+            """
+# Introduction to Earth2Studio
+
+Earth2Studio is a Python package built to empower researchers, scientists, and enthusiasts in the fields of weather and climate science with the latest artificial intelligence models and capabilities. With an intuitive design and a comprehensive feature set, it serves as a robust toolkit for exploring modern AI workflows for weather and climate.
+
+#### Learning Outcomes
+
+- Earth2Studio key features
+- How to instantiate a built-in prognostic model
+- Creating a data source and IO object
+- Running a simple built-in workflow
+- Post-processing results
+
+---
+
+## Package Design
+
+The goal of Earth2Studio is to enable users to extrapolate and build beyond what is implemented in it. The design philosophy embodies a **modular architecture** where the inference workflow acts as a flexible adhesive, seamlessly binding together various specialized software components with well-defined interfaces.
+
+<div style="display:flex; justify-content:center; gap: 10px;">
+  <figure style="text-align:center; max-width: 900px;">
+    <img src="https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/arch.png" style="width:100%; height:auto;">
+    <figcaption>Model architecture overview.</figcaption>
+  </figure>
+</div>
+
+By viewing the inference workflow as a dynamic connector, Earth2Studio facilitates effortless integration of these components, allowing researchers to easily swap out or augment functionalities to suit their specific needs.
+
+<div style="display:flex; justify-content:center; gap: 10px;">
+  <figure style="text-align:center; max-width: 900px;">
+    <img src="https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/samples.png" style="width:100%; height:auto;">
+  </figure>
+</div>
+
+### Key Features
+
+- **Built-in Workflows**: Multiple built-in inference workflows to accelerate your development and research.
+- **Prognostic Models**: Support for the latest AI weather forecast models offered under a coherent interface.
+- **Diagnostic Models**: Diagnostic models for mapping to other quantities of interest.
+- **Datasources**: Datasources to connect on-prem and remote data stores to inference workflows.
+- **IO**: Simple, yet powerful IO utilities to export data for post-processing.
+- **Statistical Operators**: Statistical methods to fuse directly into your inference workflow for more complex uncertainty analysis.
+
+---
+
+## Simple Deterministic Inference
+
+<div style="display:flex; justify-content:center; gap: 10px;">
+  <figure style="text-align:center; max-width: 900px;">
+    <img src="https://raw.githubusercontent.com/openhackathons-org/End-to-End-AI-for-Science/main/workspace/python/jupyter_notebook/Earth2Studio/images/deterministic.png" style="width:100%; height:auto;">
+  </figure>
+</div>
+
+All workflows inside Earth2Studio require constructed components to be handed to them. In this example, we use `earth2studio.run.deterministic`.
+
+### Prognostic Models
+
+Prognostic models are a class of models that perform time-integration. They are typically used to generate forecast predictions. Examples include:
+
+| Model | Description |
+|-------|-------------|
+| `models.px.FCN` | FourCastNet - AFNO-based global weather forecasting model (used in this demo) |
+| `models.px.SFNO` | Spherical Fourier Operator Network global prognostic model |
+| `models.px.Pangu24` | Pangu Weather 24 hour model |
+| `models.px.FuXi` | FuXi weather model with three auto-regressive U-net transformer models |
+| `models.px.Aurora` | Aurora transformer-based weather model |
+
+### Data Sources
+
+Data sources are used for downloading, caching and reading different weather/climate data APIs into Xarray data arrays. Used for fetching initial conditions for inference and validation data for scoring:
+
+| Data Source | Description |
+|-------------|-------------|
+| `data.GFS` | Global Forecast System initial state data source (used in this demo) |
+| `data.ARCO` | Analysis-Ready, Cloud Optimized ERA5 re-analysis data curated by Google |
+| `data.CDS` | Climate Data Store serving ERA5 re-analysis data |
+| `data.HRRR` | High-Resolution Rapid Refresh North-American weather forecast model |
+| `data.IFS` | Integrated Forecast System initial state data source |
+
+### IO Backends
+
+IO Backends are used for saving the inference results for further post-processing:
+
+| IO Backend | Description |
+|------------|-------------|
+| `io.ZarrBackend` | Zarr format backend (used in this demo) |
+| `io.NetCDF4Backend` | NetCDF4 format backend |
+| `io.XarrayBackend` | Xarray backed IO object |
+| `io.KVBackend` | Key-value (dict) backend |
+
+---
+
+## Code Overview
+
+### Set Up
+
+```python
+import os
+from earth2studio.data import GFS
+from earth2studio.io import ZarrBackend
+from earth2studio.models.px import FCN
+
+# Set cache directory
+os.environ['EARTH2STUDIO_CACHE'] = os.getcwd() + "/outputs/cache"
+
+# Prognostic Model - Load from NGC (ngc://models/nvidia/modulus/modulus_fcn@v0.2)
+package = FCN.load_default_package()
+model = FCN.load_model(package)
+
+# Data Source - Create the data source
+data = GFS()
+
+# IO Backend - Create the IO handler
+io = ZarrBackend()
+```
+
+### Execute the Workflow
+
+The `run.deterministic` function signature:
+
+```python
+def deterministic(
+    time: list[str] | list[datetime] | list[np.datetime64],
+    nsteps: int,
+    prognostic: PrognosticModel,
+    data: DataSource,
+    io: IOBackend,
+    output_coords: CoordSystem = OrderedDict({}),
+    device: torch.device | None = None,
+) -> IOBackend:
+    \"\"\"Built in deterministic workflow.
+
+    This workflow creates a deterministic inference pipeline to produce
+    a forecast prediction using a prognostic model.
+
+    Parameters
+    ----------
+    time : list[str] | list[datetime] | list[np.datetime64]
+        List of string, datetimes or np.datetime64
+    nsteps : int
+        Number of forecast steps
+    prognostic : PrognosticModel
+        Prognostic model
+    data : DataSource
+        Data source
+    io : IOBackend
+        IO object
+    output_coords: CoordSystem, optional
+        IO output coordinate system override
+    device : torch.device, optional
+        Device to run inference on
+
+    Returns
+    -------
+    IOBackend
+        Output IO object
+    \"\"\"
+```
+
+Running the forecast (each step is 6 hours for FCN, ~5-10 seconds/step on GPU):
+
+```python
+import earth2studio.run as run
+
+nsteps = 4  # 4 steps = 24 hours
+io = run.deterministic(["2024-01-01"], nsteps, model, data, io)
+
+print(io.root.tree())
+```
+
+### Post Processing
+
+```python
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+
+forecast = "2024-01-01"
+variable = "t2m"
+step = 1  # lead time = 1 x 6 = 6 hrs
+
+projection = ccrs.Robinson()
+fig, ax = plt.subplots(subplot_kw={"projection": projection}, figsize=(10, 6))
+
+im = ax.pcolormesh(
+    io["lon"][:],
+    io["lat"][:],
+    io[variable][0, step],
+    transform=ccrs.PlateCarree(),
+    cmap="Spectral_r",
+)
+
+ax.set_title(f"{forecast} - Lead time: {6*step}hrs")
+ax.coastlines()
+ax.gridlines()
+plt.savefig("outputs/t2m_prediction.jpg")
+```
+
+---
+
+## Interactive Demo
+
+This Space runs the deterministic workflow using **FCN** (FourCastNet, checkpoint from [NVIDIA NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/modulus/models/modulus_fcn)) and plots **t2m** (2-meter temperature) at your chosen lead time.
+
+FCN uses the AFNO (Adaptive Fourier Neural Operator) architecture and requires ~8GB VRAM.
+"""
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                forecast_date = gr.Textbox(
+                    label="Forecast Date",
+                    value="2024-01-01",
+                    placeholder="YYYY-MM-DD",
+                    info="GFS data available from ~2020-present",
+                    max_lines=1,
+                )
+            with gr.Column(scale=1):
+                nsteps = gr.Slider(
+                    minimum=1,
+                    maximum=5,
+                    step=1,
+                    value=5,
+                    label="Number of Forecast Steps",
+                    info="Each step = 6 hours (5 steps = 30 hours total)",
+                )
+        
+        run_btn = gr.Button("Run Inference on ZeroGPU H200", variant="primary")
+        
+        with gr.Row():
+            plot_step = gr.Slider(
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=2,
+                label="Display Timestep",
+                info="0=initial conditions, 1-N=forecast steps (updates instantly from cache)",
+            )
+        
+        status = gr.Textbox(label="Status", interactive=False)
+        out_img = gr.Image(label="2-meter Temperature (t2m)", type="filepath")
+        
+        # Hidden state to cache inference results
+        cached_date = gr.State(value=None)
+        cached_nsteps = gr.State(value=None)
+        cached_lon = gr.State(value=None)
+        cached_lat = gr.State(value=None)
+        cached_fields = gr.State(value=None)
+
+        def _sync_plot_step_max(n: int):
+            n = int(n)
+            # deterministic outputs n+1 time points, so max plot_step = n
+            new_max = max(1, n)
+            # Default to middle timestep for more interesting view
+            new_val = min(n // 2, new_max)
+            return gr.Slider(maximum=new_max, value=new_val)
+
+        # Update plot_step max when nsteps changes
+        nsteps.change(fn=_sync_plot_step_max, inputs=[nsteps], outputs=[plot_step])
+        
+        # Run inference and cache results
+        run_btn.click(
+            fn=run_forecast,
+            inputs=[forecast_date, nsteps],
+            outputs=[cached_date, cached_nsteps, cached_lon, cached_lat, cached_fields, status],
+        ).then(
+            fn=update_plot_from_cache,
+            inputs=[cached_date, cached_nsteps, cached_lon, cached_lat, cached_fields, plot_step],
+            outputs=[out_img, status],
+        )
+        
+        # Update plot when plot_step slider changes (instant, uses cache)
+        plot_step.change(
+            fn=update_plot_from_cache,
+            inputs=[cached_date, cached_nsteps, cached_lon, cached_lat, cached_fields, plot_step],
+            outputs=[out_img, status],
+        )
+
+    return demo
+
+
+# ============================================================
+# STARTUP
+# Note: Model is loaded inside @spaces.GPU function because
+# ZeroGPU uses multiprocessing and can't pickle the model.
+# ============================================================
+print("[App] Building Gradio UI...")
+
+# Create demo at module level so HF Spaces can find it
+demo = build_ui()
+
+if __name__ == "__main__":
+    # Fix for local testing: nest_asyncio patches asyncio.run in a way
+    # incompatible with uvicorn's loop_factory. Restore original.
+    asyncio.run = _ORIGINAL_ASYNCIO_RUN
+    
+    demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+# Pin torch to version required by earth2studio (>=2.5.0) that's also in ZeroGPU supported list
+# earth2studio 0.9.0+ requires torch>=2.5.0, and ZeroGPU docs list 2.5.1 as supported
+torch==2.5.1
+# Other dependencies
+gradio>=4.0.0
+spaces
+earth2studio[fcn]>=0.9.0
+numpy>=2.0.0
+zarr
+pillow