import os import ssl import warnings import spaces import certifi from aifs.plot import plot_field os.environ["SSL_CERT_FILE"] = certifi.where() os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() ssl._create_default_https_context = ssl.create_default_context warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=UserWarning) import gradio as gr import tempfile from aifs.device import device_label import shutil # Clear corrupted IC cache on startup if os.path.exists("ic_cache"): shutil.rmtree("ic_cache") os.makedirs("ic_cache", exist_ok=True) PLOTTABLE_FIELDS = [ "100u", "100v", "10u", "10v", "2t", "msl", "sp", "tcw", "swh", "mwp", "t_850", "t_500", "u_850", "v_850", "z_500", "q_700", ] UNITS_MAP = { "2t": "K", "t_850": "K", "t_500": "K", "msl": "Pa", "sp": "Pa", "z_500": "m²/s²", "100u": "m/s", "100v": "m/s", "10u": "m/s", "10v": "m/s", "u_850": "m/s", "v_850": "m/s", "swh": "m", "mwp": "s", "tcw": "kg/m²", "q_700": "kg/kg", } # ── Forecast runner ─────────────────────────────────────────────────────────── _BTN_RUNNING = gr.update(interactive=False, value="⏳ Running…") _BTN_READY = gr.update(interactive=True, value="▶ Run Forecast") def _phase(text: str) -> str: return f"**{text}**" @spaces.GPU def _run_forecast_gpu(fields, date, lead_time, num_chunks): """Thin GPU-scoped wrapper — only the inference step runs inside the ZeroGPU allocation, so slow/retriable ECMWF downloads never eat into (or blow) the GPU duration budget.""" from aifs.forecast import run_forecast as _run_forecast yield from _run_forecast(fields, date, lead_time=lead_time, num_chunks=num_chunks) def run_forecast(lead_time: int, num_chunks: int): from aifs.initial_conditions import load_ics def emit(status, phase, dd=gr.update(choices=[]), states=[], btn=_BTN_RUNNING): return status, _phase(phase), dd, states, btn yield emit("Starting up…", "📥 Downloading initial conditions from ECMWF…") try: fields = date = None for kind, payload in load_ics(cache_dir="ic_cache"): if kind == "log": yield emit(payload, "📥 Downloading initial conditions from ECMWF…") else: fields, date = payload label = device_label() total_steps = lead_time // 6 yield emit( f"Data ready — {date}", f"🤖 Running {lead_time}h forecast ({total_steps} steps) on {label}…", ) states = [] for kind, payload in _run_forecast_gpu(fields, date, lead_time, num_chunks): if kind == "log": yield emit(payload, f"🤖 Running inference on {label}…", states=states) else: states = payload timestamps = [str(s["date"]) for s in states] yield ( f"✅ Forecast complete — {len(states)} steps ({lead_time}h) · {label}", _phase(f"✅ Done — {len(states)} forecast steps ready to plot"), gr.update(choices=timestamps, value=timestamps[-1]), states, _BTN_READY, ) except Exception as exc: yield ( f"❌ Error: {exc}", _phase("❌ Something went wrong — see the log above"), gr.update(choices=[]), [], _BTN_READY, ) # ── Plot ────────────────────────────────────────────────────────────────────── def plot_forecast_handler(field_name: str, timestamp: str, states: list): if not states: return None, "Run the forecast first." state = next( (s for s in states if str(s["date"]) == timestamp), states[-1], ) if field_name not in state["fields"]: return None, f"Field '{field_name}' not available in this forecast state." values = state["fields"][field_name] units = UNITS_MAP.get(field_name, "") fig = plot_field(state=state, variable=field_name) path = os.path.join(tempfile.gettempdir(), "aifs_plot.png") fig.savefig(path, dpi=150, bbox_inches="tight") stats = ( f"**{field_name}** at {timestamp}\n\n" f"- Min: `{values.min():.4g}` {units}\n" f"- Max: `{values.max():.4g}` {units}\n" f"- Mean: `{values.mean():.4g}` {units}\n" f"- Grid points: `{len(values):,}`" ) return path, stats # ── UI ──────────────────────────────────────────────────────────────────────── DARK_CSS = """ body, .gradio-container { background: #0d1117!important; color: #cdd9e5!important; font-family: 'Inter','Segoe UI',sans-serif; } h1 { color: #58a6ff!important; letter-spacing: -0.5px; } h3 { color: #79c0ff!important; } .panel { background: #161b22!important; border: 1px solid #30363d!important; border-radius: 8px; } button.primary { background: #1f6feb!important; border: none!important; color: white!important; } button.primary:hover { background: #388bfd!important; } button.primary:disabled { background: #30363d!important; color: #8b949e!important; cursor: not-allowed!important; } .label-wrap { color: #8b949e!important; } textarea, input, select { background: #1c2128!important; color: #cdd9e5!important; border-color: #30363d!important; } textarea::placeholder, input::placeholder { color: #a8b3c0!important; } .output-markdown { color: #cdd9e5!important; } #phase-box { background: #161b22!important; border: 1px solid #30363d!important; border-radius: 6px; padding: 10px 14px!important; margin-bottom: 8px; } #phase-box strong { color: #e3b341!important; font-size: 0.95em; } /* === Notes / Markdown prose === */ .prose, .prose p, .prose li, .prose ul, .prose ol, .prose strong { color: #a8b3c0!important; } .prose h1, .prose h2, .prose h3, .prose h4 { color: #a8b3c0!important; } footer { display: none!important; } """ with gr.Blocks(css=DARK_CSS, title="AIFS Single v2 Forecast") as demo: gr.Markdown( """ # 🌍 AIFS Single v2 — Weather Forecast **ECMWF's AI Integrated Forecasting System** | Runs on CUDA · Apple MPS · CPU """ ) forecast_state = gr.State([]) with gr.Row(): with gr.Column(scale=1, elem_classes="panel"): gr.Markdown("### ⚙️ Forecast Settings") lead_time_sl = gr.Slider( minimum=6, maximum=96, step=6, value=6, label="Lead time (hours)", info="On CPU/MPS keep this at 6h for testing — each step is slow without a CUDA GPU.", ) num_chunks_sl = gr.Slider( minimum=1, maximum=32, step=1, value=16, label="Memory chunks", info="Higher = less memory, slightly slower. Ignored on CPU.", ) run_btn = gr.Button("▶ Run Forecast", variant="primary", size="lg") phase_md = gr.Markdown( "Ready — click **▶ Run Forecast** to begin.", elem_id="phase-box", ) status_box = gr.Textbox( label="Detailed log", lines=6, interactive=False, placeholder="Progress details will appear here…", ) with gr.Column(scale=2, elem_classes="panel"): gr.Markdown("### 🗺️ Visualise Output") with gr.Row(): field_dd = gr.Dropdown( choices=PLOTTABLE_FIELDS, value="2t", label="Field", info="Select the variable to plot.", ) timestamp_dd = gr.Dropdown( choices=[], label="Forecast step", info="Available after running the forecast.", ) plot_btn = gr.Button("🖼 Plot Field", variant="secondary") map_img = gr.Image(label="Global Map", type="filepath") stats_md = gr.Markdown() run_btn.click( fn=run_forecast, inputs=[lead_time_sl, num_chunks_sl], outputs=[status_box, phase_md, timestamp_dd, forecast_state, run_btn], ) plot_btn.click( fn=plot_forecast_handler, inputs=[field_dd, timestamp_dd, forecast_state], outputs=[map_img, stats_md], ) gr.Markdown( """ --- **Notes** - No flash-attn required — uses PyTorch SDPA (works on CPU, MPS, and CUDA). - First run downloads the ~2 GB model checkpoint from Hugging Face and caches it. - Data: ECMWF Open Data. """ ) if __name__ == "__main__": demo.launch()