# Save this file as in the root of the cloned Surya repository import gradio as gr import torch from huggingface_hub import snapshot_download import yaml import numpy as np from PIL import Image import sunpy.map import sunpy.net.attrs as a from sunpy.net import Fido from astropy.wcs import WCS import astropy.units as u from reproject import reproject_interp import os import warnings import logging import datetime import matplotlib.pyplot as plt import sunpy.visualization.colormaps as sunpy_cm import traceback # --- Use the official Surya modules --- from surya.models.helio_spectformer import HelioSpectFormer from surya.utils.data import build_scalers from surya.datasets.helio import inverse_transform_single_channel # --- Configuration --- warnings.filterwarnings("ignore", category=UserWarning, module='sunpy') warnings.filterwarnings("ignore", category=FutureWarning) logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Global cache for model, config, etc. APP_CACHE = {} SDO_CHANNELS_MAP = { "aia94": (a.Wavelength(94 * u.angstrom), a.Sample(12 * u.s)), "aia131": (a.Wavelength(131 * u.angstrom), a.Sample(12 * u.s)), "aia171": (a.Wavelength(171 * u.angstrom), a.Sample(12 * u.s)), "aia193": (a.Wavelength(193 * u.angstrom), a.Sample(12 * u.s)), "aia211": (a.Wavelength(211 * u.angstrom), a.Sample(12 * u.s)), "aia304": (a.Wavelength(304 * u.angstrom), a.Sample(12 * u.s)), "aia335": (a.Wavelength(335 * u.angstrom), a.Sample(12 * u.s)), "aia1600": (a.Wavelength(1600 * u.angstrom), a.Sample(24 * u.s)), "hmi_m": (a.Physobs("intensity"), a.Sample(45 * u.s)), "hmi_bx": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)), "hmi_by": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)), # Placeholder "hmi_bz": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)), # Placeholder "hmi_v": (a.Physobs("los_velocity"), a.Sample(45 * u.s)), } SDO_CHANNELS = list(SDO_CHANNELS_MAP.keys()) # --- 1. Model Loading and Setup --- def setup_and_load_model(): if "model" in APP_CACHE: yield "Model already loaded. Skipping setup." return yield "Downloading model files (first run only)..." snapshot_download(repo_id="nasa-ibm-ai4science/Surya-1.0", local_dir="data/Surya-1.0", allow_patterns=["config.yaml", "scalers.yaml", "surya.366m.v1.pt"]) yield "Loading configuration and data scalers..." with open("data/Surya-1.0/config.yaml") as fp: config = yaml.safe_load(fp) APP_CACHE["config"] = config scalers_info = yaml.safe_load(open("data/Surya-1.0/scalers.yaml", "r")) APP_CACHE["scalers"] = build_scalers(info=scalers_info) yield "Initializing model architecture..." model_config = config["model"] model = HelioSpectFormer(...) # Full model definition device = "cuda" if torch.cuda.is_available() else "cpu" APP_CACHE["device"] = device yield f"Loading model weights to {device}..." weights = torch.load(f"data/Surya-1.0/surya.366m.v1.pt", map_location=torch.device(device)) model.load_state_dict(weights, strict=True) model.to(device) model.eval() APP_CACHE["model"] = model yield "✅ Model setup complete." # --- 2. Live Data Fetching and Preprocessing (as a generator) --- def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes): config = APP_CACHE["config"] img_size = config["model"]["img_size"] input_deltas = config["data"]["time_delta_input_minutes"] target_delta = forecast_horizon_minutes # Use user-provided horizon input_times = [target_dt + datetime.timedelta(minutes=m) for m in input_deltas] target_time = target_dt + datetime.timedelta(minutes=target_delta) all_times = sorted(list(set(input_times + [target_time]))) data_maps = {} total_downloads = len(all_times) * len(SDO_CHANNELS) downloads_done = 0 yield f"Starting download of {total_downloads} data files..." for t in all_times: data_maps[t] = {} for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()): downloads_done += 1 yield f"Downloading [{downloads_done}/{total_downloads}]: {channel} for {t.strftime('%Y-%m-%d %H:%M')}..." if channel in ["hmi_by", "hmi_bz"]: if data_maps[t].get("hmi_bx"): data_maps[t][channel] = data_maps[t]["hmi_bx"] continue # *** FIX: Use a.Time.nearest=True for robust fetching instead of a time window *** instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia query = Fido.search(a.Time(t), instrument, physobs, sample, a.Time.nearest==True) if not query: raise ValueError(f"No data found for {channel} near {t}") files = Fido.fetch(query, path="./data/sdo_cache") # Fetch the entire result data_maps[t][channel] = sunpy.map.Map(files[0]) yield "✅ All files downloaded. Starting preprocessing..." output_wcs = WCS(naxis=2) output_wcs.wcs.crpix = [(img_size + 1) / 2, (img_size + 1) / 2] output_wcs.wcs.cdelt = np.array([-1.2, 1.2]) * u.arcsec output_wcs.wcs.crval = [0, 0] * u.arcsec output_wcs.wcs.ctype = ['HPLN-TAN', 'HPLT-TAN'] scaler = APP_CACHE["scalers"] processed_tensors = {} for t, channel_maps in data_maps.items(): channel_tensors = [] for i, channel in enumerate(SDO_CHANNELS): smap = channel_maps[channel] reprojected_data, _ = reproject_interp(smap, output_wcs, shape_out=(img_size, img_size)) exp_time = smap.meta.get('exptime', 1.0) if exp_time is None or exp_time <= 0: exp_time = 1.0 norm_data = reprojected_data / exp_time scaled_data = scaler.transform(norm_data, c_idx=i) channel_tensors.append(torch.from_numpy(scaled_data.astype(np.float32))) processed_tensors[t] = torch.stack(channel_tensors) yield "✅ Preprocessing complete." input_tensor_list = [processed_tensors[t] for t in input_times] input_tensor = torch.stack(input_tensor_list, dim=1).unsqueeze(0) target_map = data_maps[target_time] last_input_map = data_maps[input_times[-1]] yield (input_tensor, last_input_map, target_map) # --- 3. Inference and Visualization --- def run_inference(input_tensor): # This function remains the same ... def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name): # This function remains the same ... # --- 4. Gradio UI and Controllers --- def forecast_controller(date_str, hour, minute, forecast_horizon): yield { log_box: gr.update(value="Starting forecast...", visible=True), run_button: gr.update(interactive=False), # Also disable the other controls date_input: gr.update(interactive=False), hour_slider: gr.update(interactive=False), minute_slider: gr.update(interactive=False), horizon_slider: gr.update(interactive=False), results_group: gr.update(visible=False) } try: if not date_str: raise gr.Error("Please select a date.") for status in setup_and_load_model(): yield { log_box: status } # Construct datetime from the new UI components target_dt = datetime.datetime.fromisoformat(f"{date_str}T{int(hour):02d}:{int(minute):02d}:00") data_pipeline = fetch_and_process_sdo_data(target_dt, forecast_horizon) # The rest of the generator logic remains the same... while True: try: status = next(data_pipeline) if isinstance(status, tuple): input_tensor, last_input_map, target_map = status break yield { log_box: status } except StopIteration: raise gr.Error("Data processing pipeline finished unexpectedly.") yield { log_box: "Running AI model inference..." } prediction_tensor = run_inference(input_tensor) yield { log_box: "Generating final visualizations..." } img_in, img_pred, img_target = generate_visualization(last_input_map, prediction_tensor, target_map, "aia171") yield { log_box: f"✅ Forecast complete for {target_dt.isoformat()} (+{forecast_horizon} mins).", results_group: gr.update(visible=True), # ... update states and images } except Exception as e: # ... error handling finally: # Re-enable all controls yield { run_button: gr.update(interactive=True), date_input: gr.update(interactive=True), hour_slider: gr.update(interactive=True), minute_slider: gr.update(interactive=True), horizon_slider: gr.update(interactive=True), } # --- 5. Gradio UI Definition --- with gr.Blocks(theme=gr.themes.Soft()) as demo: # State objects remain the same ... gr.Markdown(...) # Title remains the same # --- NEW: Controls Section --- with gr.Accordion("Step 1: Configure Forecast", open=True): with gr.Row(): date_input = gr.Textbox( label="Date", value=datetime.date.today().strftime("%Y-%m-%d") ) hour_slider = gr.Slider(label="Hour (UTC)", minimum=0, maximum=23, step=1, value=datetime.datetime.utcnow().hour - 3) minute_slider = gr.Slider(label="Minute", minimum=0, maximum=59, step=1, value=datetime.datetime.utcnow().minute) horizon_slider = gr.Slider( label="Forecast Horizon (minutes ahead)", minimum=12, maximum=120, step=12, value=12 ) run_button = gr.Button("🔮 Generate Forecast", variant="primary") # --- NEW: Moved log box to its own section --- with gr.Accordion("Step 2: View Log", open=False) as log_accordion: log_box = gr.Textbox(label="Log", interactive=False, visible=True, lines=5, max_lines=10) # --- Results section is now Step 3 --- with gr.Group(visible=False) as results_group: gr.Markdown("### Step 3: Explore Results") channel_selector = gr.Dropdown(...) with gr.Row(): input_display = gr.Image(...) prediction_display = gr.Image(...) target_display = gr.Image(...) # --- Event Handlers --- run_button.click( fn=forecast_controller, inputs=[date_input, hour_slider, minute_slider, horizon_slider], outputs=[ log_box, run_button, date_input, hour_slider, minute_slider, horizon_slider, results_group, state_last_input, state_prediction, state_target, input_display, prediction_display, target_display ] ) channel_selector.change(...) # This remains the same if __name__ == "__main__": # Fill in the missing ... from previous versions for the full script # This is a condensed version showing only the key changes demo.launch(debug=True)