Upload 3 files

Dockerfile

@@ -0,0 +1,47 @@
+# 1. Use Python 3.10 (Standard for HF Spaces)
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# 2. Install System Dependencies
+# Required for Cartopy (maps) and Git (for installing from source)
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    git \
+    libgeos-dev \
+    libproj-dev \
+    proj-bin \
+    && rm -rf /var/lib/apt/lists/*
+
+# 3. CRITICAL: Install NATTEN from Pre-built Wheels
+# If we don't do this, pip tries to compile it from source (takes 20+ mins) and fails.
+# We match the CUDA version (cu121) and Python version (cp310).
+RUN pip install natten==0.17.3 -f https://shi-labs.com/natten/wheels/cu121/
+
+# 4. CRITICAL: Install Earth2Studio from GitHub
+# We use the [stormscope] tag to tell it we have the dependencies ready.
+RUN pip install "earth2studio[stormscope] @ git+https://github.com/NVIDIA/earth2studio.git"
+
+# 5. Install Other Python Dependencies
+# (Streamlit, Maps, etc.)
+RUN pip install \
+    streamlit \
+    torch \
+    torchvision \
+    numpy \
+    matplotlib \
+    cartopy \
+    huggingface_hub \
+    scipy
+
+# 6. Copy App Code
+COPY . .
+
+# 7. Launch App
+CMD ["streamlit", "run", "app.py", \
+     "--server.port", "7860", \
+     "--server.address", "0.0.0.0", \
+     "--server.enableCORS", "false", \
+     "--server.enableXsrfProtection", "false", \
+     "--server.fileWatcherType", "none"]

app.py

@@ -0,0 +1,177 @@
+import streamlit as st
+import time
+import sys
+import os
+import gc
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
+from datetime import datetime, timedelta, timezone
+
+# --- PAGE CONFIG ---
+st.set_page_config(page_title="Canada Generative Radar (Earth2Studio)", layout="wide")
+
+# --- VISIBLE LOGGING ---
+status_container = st.empty()
+
+def log_to_ui(msg, type="info"):
+    """Helper to print logs to the UI and console."""
+    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
+    if type == "info": status_container.info(f"📋 {msg}")
+    elif type == "success": status_container.success(f"✅ {msg}")
+    elif type == "error": status_container.error(f"❌ {msg}")
+    elif type == "warning": status_container.warning(f"⚠️ {msg}")
+
+log_to_ui("🚀 Initializing Radar App...")
+
+# --- EARTH2STUDIO IMPORTS ---
+try:
+    # 1. Import Earth2Studio
+    from earth2studio.models.px import StormScopeMRMS
+    from earth2studio.utils.time import to_time_array
+    
+    # 2. Setup Device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"✅ Setup Complete. Device: {device}")
+    
+except ImportError as e:
+    st.error(f"CRITICAL ERROR: {e}")
+    st.info("The app is likely missing the 'StormScope' update. Check the Dockerfile installation logs.")
+    st.stop()
+
+# --- CONFIG ---
+# Canadian Regions of Interest
+REGIONS = {
+    "Toronto / Southern Ontario": {"lat": 43.7, "lon": -79.4, "zoom": 4},
+    "Montreal / Quebec":          {"lat": 45.5, "lon": -73.6, "zoom": 4},
+    "Vancouver / BC":             {"lat": 49.3, "lon": -123.1, "zoom": 4},
+    "Calgary / Alberta":          {"lat": 51.0, "lon": -114.1, "zoom": 5},
+}
+
+# --- MODEL MANAGERS (CACHED) ---
+
+@st.cache_resource(show_spinner=False)
+def load_radar_model():
+    """
+    Loads the StormScopeMRMS model. 
+    Cached so we don't reload 4GB+ weights on every interaction.
+    """
+    gc.collect()
+    torch.cuda.empty_cache()
+    print("Loading StormScope Model...", flush=True)
+    
+    # Load Model Package
+    package = StormScopeMRMS.load_default_package()
+    model = StormScopeMRMS.load_model(package)
+    model.to(device)
+    model.eval()
+    return model
+
+@st.cache_data(show_spinner=False, ttl=1800) # Cache for 30 mins
+def run_radar_generation(region_name, lat_center, lon_center):
+    """
+    Runs the generative radar inference for a specific location.
+    Returns the image figure (matplotlib) to display.
+    """
+    model = load_radar_model()
+    
+    # 1. Setup Time (Use a recent past time to ensure data availability)
+    # Real-time radar data often lags by 30-60 mins in public buckets
+    now = datetime.now(timezone.utc) - timedelta(hours=1) 
+    time_str = now.strftime("%Y-%m-%dT%H:00:00")
+    time_obj = to_time_array([time_str])
+
+    # 2. Fetch Input Data
+    # StormScope has a built-in fetch_data method that grabs the necessary
+    # initial conditions (usually previous radar frames) from the internet.
+    try:
+        # Note: This connects to NASA/NOAA servers. If they are down, this fails.
+        x, coords = model.fetch_data(time_obj)
+        x = x.to(device)
+    except Exception as e:
+        return None, f"Data Fetch Error (External Source): {str(e)}"
+
+    # 3. Run Inference (Generate Next Frame)
+    with torch.no_grad():
+        out, out_coords = model(x, coords)
+
+    # 4. Extract Data (Reflectivity)
+    # Output shape: [Batch, Time, Lat, Lon] or [Batch, Channel, Lat, Lon]
+    # StormScopeMRMS output channel 0 is typically reflectivity
+    radar_data = out[0, 0, :, :].cpu().numpy()
+    
+    lats = out_coords['lat']
+    lons = out_coords['lon']
+
+    # 5. Crop / Focus on Canada Region requested
+    # We simply return the full array and coords, and let the UI handle zooming via plotting
+    return (radar_data, lats, lons), "Success"
+
+# --- UI LAYOUT ---
+st.title("🇨🇦 Canada Generative Radar (Earth2Studio)")
+st.markdown("""
+Using **NVIDIA Earth-2 StormScope** to generate high-resolution radar reflectivity.
+*Note: This model is trained on US data but covers Southern Canada.*
+""")
+
+# Sidebar Controls
+with st.sidebar:
+    st.header("Settings")
+    selected_region = st.selectbox("Choose Region", list(REGIONS.keys()))
+    
+    region_info = REGIONS[selected_region]
+    lat_center = region_info["lat"]
+    lon_center = region_info["lon"]
+    zoom_deg = st.slider("Zoom (Degrees Radius)", 1.0, 10.0, 4.0)
+
+# Main Execution Button
+if st.button("📡 Generate Radar Forecast", type="primary", use_container_width=True):
+    
+    log_to_ui(f"Fetching data and generating radar for {selected_region}...", type="info")
+    
+    t0 = time.time()
+    
+    # Run Inference
+    result, msg = run_radar_generation(selected_region, lat_center, lon_center)
+    
+    if result is None:
+        log_to_ui(msg, type="error")
+    else:
+        radar_data, lats, lons = result
+        elapsed = time.time() - t0
+        log_to_ui(f"Generation Complete in {elapsed:.2f}s", type="success")
+        
+        # Plotting
+        st.subheader(f"Radar Reflectivity: {selected_region}")
+        
+        fig = plt.figure(figsize=(10, 8))
+        ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
+        
+        # Map Features
+        ax.add_feature(cfeature.COASTLINE, linewidth=1)
+        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=1)
+        ax.add_feature(cfeature.LAKES, alpha=0.3, color='blue')
+        
+        # Set Extent (Zoom)
+        extent = [lon_center - zoom_deg, lon_center + zoom_deg, 
+                  lat_center - zoom_deg, lat_center + zoom_deg]
+        ax.set_extent(extent, crs=ccrs.PlateCarree())
+        
+        # Plot Radar Data
+        # Use a transparent colormap for low values (0 reflectivity)
+        mesh = ax.pcolormesh(lons, lats, radar_data, 
+                             transform=ccrs.PlateCarree(),
+                             cmap='nipy_spectral', 
+                             vmin=0, vmax=70,  # dBZ range
+                             shading='auto')
+        
+        plt.colorbar(mesh, ax=ax, label='Reflectivity (dBZ)', shrink=0.7)
+        plt.title(f"Simulated Radar | Center: {lat_center}, {lon_center}")
+        
+        st.pyplot(fig)
+
+# Footer
+st.markdown("---")
+st.caption("Powered by NVIDIA Earth2Studio • Runs on Hugging Face GPU Spaces")

requirements.txt

@@ -0,0 +1,15 @@
+streamlit
+pandas
+numpy
+requests
+altair
+earth2studio
+torch
+matplotlib
+cartopy
+h5netcdf
+zarr
+onnx
+onnxruntime-gpu
+einops
+nvidia-physicsnemo