Update app.py

app.py

@@ -1,172 +1,316 @@
 import gradio as gr
 import torch
-from huggingface_hub import hf_hub_download
-from surya.model import Surya # This now works because of the file structure
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from huggingface_hub import snapshot_download
+import yaml
 import numpy as np
 from PIL import Image
+import sunpy.visualization.colormaps as sunpy_cm
 import os
+import glob
 import warnings
+import logging
 
-# Suppress warnings for a cleaner demo experience
+# --- Suppress verbose logging and warnings for a cleaner UI ---
 warnings.filterwarnings("ignore")
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# --- 1. Define Constants and Data Channels ---
-# Based on the Surya project's data preprocessing
-AIA_CHANNELS = ["94", "131", "171", "193", "211", "304", "335", "1600"]
-HMI_CHANNELS = ["bx", "by", "bz", "by_abs", "bz_abs"]
-ALL_CHANNELS = [f"AIA {ch} Å" for ch in AIA_CHANNELS] + [f"HMI {ch}" for ch in HMI_CHANNELS]
+# --- Dependencies from the Surya Repository ---
+# NOTE: To make this script self-contained, the required classes and functions
+# from the 'surya' library are included directly here.
+# In a full installation, these would be imported.
 
-# --- 2. Caching and Loading the Model and Data ---
+from surya_dependencies import (
+    HelioSpectFormer,
+    HelioNetCDFDataset,
+    build_scalers,
+    custom_collate_fn,
+    inverse_transform_single_channel,
+    SDO_CHANNELS,
+    AIA_CHANNELS,
+    HMI_CHANNELS
+)
+
+# --- Global Cache for Model and Data ---
+# We use a simple dictionary to act as a cache to avoid reloading.
+APP_CACHE = {
+    "model": None,
+    "config": None,
+    "scalers": None,
+    "dataset": None,
+    "dataloader": None,
+    "device": "cuda" if torch.cuda.is_available() else "cpu",
+}
+
+# --- 1. Setup and Data Download ---
 @gr.cache
-def load_model_and_data():
+def setup_environment_and_download_data():
     """
-    Downloads the pre-trained Surya model, the test data, and initializes the model.
-    This function is cached so this happens only once.
+    Downloads all necessary files from Hugging Face: model, config, scalers, and validation data.
+    Also creates the necessary index file for the dataset loader.
+    This function is cached by Gradio to run only once.
     """
-    print("Downloading model and test data... This may take a moment.")
-    # Define local directories for caching
-    model_dir = "./surya_model"
-    data_dir = "./surya_data"
-    os.makedirs(model_dir, exist_ok=True)
-    os.makedirs(data_dir, exist_ok=True)
-
-    # Download the model weights and test data from Hugging Face
-    checkpoint_path = hf_hub_download(
+    logger.info("Setting up environment. This will run only once.")
+    local_dir = "data/Surya-1.0"
+    # Download model, config, and scalers
+    snapshot_download(
         repo_id="nasa-ibm-ai4science/Surya-1.0",
-        filename="surya.366m.v1.pt",
-        local_dir=model_dir
+        local_dir=local_dir,
+        allow_patterns=["config.yaml", "scalers.yaml", "surya.366m.v1.pt"],
     )
-    test_data_path = hf_hub_download(
-        repo_id="nasa-ibm-ai4science/Surya-1.0",
-        filename="test_data.pt",
-        local_dir=data_dir
-    )
-    print("Downloads complete.")
-
-    # Initialize the model architecture
-    model = Surya(
-        img_size=4096,
-        patch_size=16,
-        in_chans=13,
-        embed_dim=1280,
-        spectral_blocks=2,
-        attention_blocks=8,
+
+    # Download validation data
+    data_dir = "data/Surya-1.0_validation_data"
+    snapshot_download(
+        repo_id="nasa-ibm-ai4science/Surya-1.0_validation_data",
+        repo_type="dataset",
+        local_dir=data_dir,
+        allow_patterns="20140107_1[5-9]??.nc",
     )
 
-    # Load the weights into the model
-    print("Loading model weights...")
-    model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
-    model.eval()
-    print("Model loaded successfully.")
+    # The test script requires an index file. We'll create it dynamically.
+    index_dir = "data/test_indices"
+    os.makedirs(index_dir, exist_ok=True)
+    index_file_path = os.path.join(index_dir, "test_surya_index.csv")
+    
+    with open(index_file_path, "w") as f:
+        f.write("path\n")
+        # Find the downloaded NetCDF files and write their paths to the index
+        search_path = os.path.join(data_dir, "**", "*.nc")
+        for nc_file in sorted(glob.glob(search_path, recursive=True)):
+            f.write(f"{nc_file}\n")
+    logger.info(f"Created index file at {index_file_path}")
+    return index_file_path, local_dir
+
+# --- 2. Model and Data Loading ---
+def load_essentials(model_dir):
+    """Loads config, scalers, and the model into the APP_CACHE."""
+    if APP_CACHE["model"] is None:
+        logger.info("Loading config, scalers, and model for the first time...")
+        # Load config
+        with open(os.path.join(model_dir, "config.yaml")) as fp:
+            config = yaml.safe_load(fp)
+        APP_CACHE["config"] = config
 
-    # Load the test data
-    test_data = torch.load(test_data_path)
-    test_input = test_data["input"]   # Input tensor for the model
-    test_label = test_data["label"]   # Ground truth for comparison
+        # Build scalers for data normalization
+        scalers_info = yaml.safe_load(open(os.path.join(model_dir, "scalers.yaml"), "r"))
+        APP_CACHE["scalers"] = build_scalers(info=scalers_info)
 
-    return model, test_input, test_label
+        # Initialize model from config
+        model = HelioSpectFormer(
+            img_size=config["model"]["img_size"],
+            patch_size=config["model"]["patch_size"],
+            in_chans=len(config["data"]["sdo_channels"]),
+            embed_dim=config["model"]["embed_dim"],
+            time_embedding={"type": "linear", "time_dim": len(config["data"]["time_delta_input_minutes"])},
+            depth=config["model"]["depth"],
+            n_spectral_blocks=config["model"]["n_spectral_blocks"],
+            num_heads=config["model"]["num_heads"],
+            mlp_ratio=config["model"]["mlp_ratio"],
+            drop_rate=config["model"]["drop_rate"],
+            dtype=torch.bfloat16,
+            window_size=config["model"]["window_size"],
+            dp_rank=config["model"]["dp_rank"],
+            learned_flow=config["model"]["learned_flow"],
+            use_latitude_in_learned_flow=config["model"]["learned_flow"],
+            init_weights=False,
+            checkpoint_layers=[i for i in range(config["model"]["depth"])],
+            rpe=config["model"]["rpe"],
+            ensemble=config["model"]["ensemble"],
+            finetune=config["model"]["finetune"],
+        )
+        
+        # Load pre-trained weights
+        path_weights = os.path.join(model_dir, "surya.366m.v1.pt")
+        weights = torch.load(path_weights, map_location=torch.device(APP_CACHE["device"]))
+        model.load_state_dict(weights, strict=True)
+        model.to(APP_CACHE["device"])
+        model.eval()
+        
+        n_params = sum(p.numel() for p in model.parameters()) / 1e6
+        logger.info(f"Surya FM: {n_params:.2f}M parameters loaded to {APP_CACHE['device']}.")
+        APP_CACHE["model"] = model
 
-# --- 3. Helper function for Image Conversion ---
-def tensor_to_image(tensor_slice):
+def get_dataloader(index_path):
+    """Initializes and returns a DataLoader for the validation data."""
+    if APP_CACHE["dataloader"] is None:
+        logger.info("Initializing dataset and dataloader...")
+        config = APP_CACHE["config"]
+        dataset = HelioNetCDFDataset(
+            index_path=index_path,
+            time_delta_input_minutes=config["data"]["time_delta_input_minutes"],
+            time_delta_target_minutes=config["data"]["time_delta_target_minutes"],
+            n_input_timestamps=len(config["data"]["time_delta_input_minutes"]),
+            rollout_steps=1,
+            channels=config["data"]["sdo_channels"],
+            scalers=APP_CACHE["scalers"],
+            phase="valid", # Important: ensure no random augmentations
+        )
+        dataloader = DataLoader(
+            dataset, shuffle=False, batch_size=1, num_workers=2,
+            pin_memory=True, drop_last=False, collate_fn=custom_collate_fn,
+        )
+        APP_CACHE["dataloader"] = dataloader
+        APP_CACHE["dataset"] = dataset # Also cache dataset for transformation info
+    return APP_CACHE["dataloader"]
+
+
+# --- 3. Core Inference and Visualization Logic ---
+def run_model_inference():
     """
-    Normalizes a 2D tensor slice and converts it to a PIL Image for display.
+    Performs a single prediction step using the loaded model and dataloader.
+    Returns the raw input, prediction, and ground truth tensors.
     """
-    # Detach tensor from graph, move to CPU, and convert to numpy
-    img_np = tensor_slice.detach().cpu().numpy()
+    model = APP_CACHE["model"]
+    dataloader = APP_CACHE["dataloader"]
+    device = APP_CACHE["device"]
+
+    # Get the first (and only) batch of data from the validation set
+    batch_data, batch_metadata = next(iter(dataloader))
     
-    # Normalize the tensor to a 0-255 range for image display
-    min_val, max_val = np.min(img_np), np.max(img_np)
-    if max_val > min_val:
-        img_np = (img_np - min_val) / (max_val - min_val)
+    logger.info("Running inference on the validation batch...")
+    with torch.no_grad():
+        # Prepare input batch for the model
+        input_batch = {key: batch_data[key].to(device) for key in ["ts", "time_delta_input"]}
+        # Run model prediction
+        with torch.autocast(device_type=device.split(':')[0], dtype=torch.bfloat16):
+            prediction_tensor = model(input_batch)
     
-    img_array = (img_np * 255).astype(np.uint8)
-    return Image.fromarray(img_array)
-
+    # Get the input and target tensors for comparison
+    input_tensor = input_batch["ts"].to(dtype=torch.float32).cpu()
+    target_tensor = batch_data["forecast"].cpu()
+    prediction_tensor = prediction_tensor.to(dtype=torch.float32).cpu()
+    
+    logger.info("Inference complete.")
+    return input_tensor, prediction_tensor, target_tensor
 
-# --- 4. Main Prediction and Visualization Function ---
-def run_forecast(channel_name, progress=gr.Progress()):
+def create_visualizations(channel_name, input_tensor, prediction_tensor, target_tensor):
     """
-    This function is triggered by the button click in the Gradio interface.
-    It runs the model prediction and generates the images for display.
+    Takes raw tensors and a channel name, applies inverse transformation,
+    and converts them to displayable PIL Images.
     """
-    progress(0, desc="Loading model and data (first run may be slow)...")
-    # Load the model and data (will be fast after the first run due to caching)
-    model, test_input, test_label = load_model_and_data()
+    if input_tensor is None:
+        return None, None, None, "Please run the forecast first."
+
+    logger.info(f"Creating visualization for channel: {channel_name}")
+    c_idx = SDO_CHANNELS.index(channel_name)
+    dataset = APP_CACHE["dataset"]
+    means, stds, epsilons, sl_scale_factors = dataset.transformation_inputs()
     
-    progress(0.5, desc="Running inference on the model...")
-    # Perform the forecast
-    with torch.no_grad():
-        prediction = model(test_input)
+    # --- Denormalize data for visualization ---
+    # Final input image given to the model (last in sequence)
+    input_slice = inverse_transform_single_channel(
+        input_tensor[0, c_idx, -1, :, :].numpy(),
+        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
+    )
+    # Model's prediction
+    pred_slice = inverse_transform_single_channel(
+        prediction_tensor[0, c_idx, :, :].numpy(),
+        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
+    )
+    # Ground truth image
+    target_slice = inverse_transform_single_channel(
+        target_tensor[0, c_idx, 0, :, :].numpy(),
+        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
+    )
+    
+    # --- Convert to images ---
+    # Use a shared color scale for better comparison, clipped at 99.5th percentile
+    vmax = np.quantile(np.concatenate([input_slice, pred_slice, target_slice]), 0.995)
+    
+    # Determine colormap from channel name
+    cmap_name = f"sdoaia{channel_name.replace('aia', '')}" if 'aia' in channel_name else 'hmimag'
+    cmap = plt.get_cmap(sunpy_cm.cmlist.get(cmap_name, 'gray'))
 
-    progress(0.8, desc="Generating visualizations...")
-    # Get the index of the selected channel
-    channel_index = ALL_CHANNELS.index(channel_name)
+    def to_pil(data, vmin=0, vmax=vmax, cmap=cmap):
+        data_clipped = np.clip(data, vmin, vmax)
+        data_norm = (data_clipped - vmin) / (vmax - vmin)
+        return Image.fromarray((cmap(data_norm)[:, :, :3] * 255).astype(np.uint8))
 
-    # Extract the last time step from the input sequence for display
-    # Shape: [batch, channels, time, height, width] -> select channel, last time step
-    input_slice = test_input[0, channel_index, -1, :, :]
-    input_image = tensor_to_image(input_slice)
+    return to_pil(input_slice), to_pil(pred_slice), to_pil(target_slice), f"Displaying forecast for {channel_name}"
 
-    # Extract the corresponding slice from the model's prediction
-    # Shape: [batch, channels, time, height, width] -> select channel, first predicted step
-    predicted_slice = prediction[0, channel_index, 0, :, :]
-    predicted_image = tensor_to_image(predicted_slice)
+# --- 4. Gradio Controller Functions ---
+def forecast_controller(channel_name, progress=gr.Progress()):
+    """
+    Main function triggered by the 'Generate' button. Orchestrates the entire pipeline.
+    """
+    progress(0, desc="Downloading model and data (first launch only)...")
+    index_path, model_dir = setup_environment_and_download_data()
+
+    progress(0.4, desc="Loading model and building data pipeline...")
+    load_essentials(model_dir)
+    get_dataloader(index_path)
+    
+    progress(0.7, desc=f"Running inference on {APP_CACHE['device']}...")
+    input_t, pred_t, target_t = run_model_inference()
 
-    # Extract the corresponding slice from the ground truth label
-    label_slice = test_label[0, channel_index, 0, :, :]
-    label_image = tensor_to_image(label_slice)
+    progress(0.9, desc="Creating visualizations...")
+    img_in, img_pred, img_target, status = create_visualizations(channel_name, input_t, pred_t, target_t)
     
-    print(f"Forecast generated for channel: {channel_name}")
-    return input_image, predicted_image, label_image
+    return img_in, img_pred, img_target, status, input_t, pred_t, target_t
+
 
-# --- 5. Building the Gradio Interface ---
+# --- 5. Gradio UI Layout ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    # Hidden state variables to store the raw tensors after inference
+    state_input = gr.State()
+    state_prediction = gr.State()
+    state_target = gr.State()
+    
     gr.Markdown(
         """
         <div align="center">
-        # ☀️ Surya: A Live Demonstration of NASA's Heliophysics Foundation Model ☀️
-        This demo runs the actual Surya model to forecast solar activity. It uses the official test data for **2014-01-07**,
-        allowing a direct comparison between the model's prediction and the real ground truth.
+        # ☀️ Surya: Live Model Demo ☀️
+        ### An Interactive Interface for NASA's Heliophysics Foundation Model
+        This demo runs the **actual** Surya model on its official validation data for **2014-01-07**.
+        Click the button to generate a forecast, then use the dropdown to explore the results across different SDO instrument channels.
         </div>
         """
     )
-
+    
     with gr.Row():
         channel_selector = gr.Dropdown(
-            choices=ALL_CHANNELS,
-            value=ALL_CHANNELS[2], # Default to "AIA 171 Å"
+            choices=SDO_CHANNELS,
+            value="aia171",
             label="🛰️ Select SDO Instrument Channel",
             info="Choose which solar observation channel to visualize."
         )
+        run_button = gr.Button("🔮 Generate Forecast for 2014-01-07", variant="primary", scale=2)
 
-    run_button = gr.Button("🔮 Generate Forecast for 2014-01-07", variant="primary")
-
+    status_box = gr.Textbox(label="Status", interactive=False, value="Ready. Press 'Generate Forecast' to start.")
+    
     with gr.Row():
         with gr.Column():
             gr.Markdown("### ⬅️ Final Input Image")
-            gr.Markdown("The last image shown to the model before it makes a prediction.")
-            input_display = gr.Image(label="Input Observation", height=400, width=400)
+            gr.Markdown("The last observation shown to the model (T-1).")
+            input_display = gr.Image(label="Input", height=512, width=512, interactive=False)
         with gr.Column():
             gr.Markdown("### 🔮 Model's Forecast")
-            gr.Markdown("What the Surya model predicted the Sun would look like.")
-            prediction_display = gr.Image(label="Surya Prediction", height=400, width=400)
+            gr.Markdown("Surya's prediction for the next timestep (T+0).")
+            prediction_display = gr.Image(label="Prediction", height=512, width=512, interactive=False)
         with gr.Column():
             gr.Markdown("### ✅ Ground Truth")
-            gr.Markdown("What the Sun *actually* looked like at the forecast time.")
-            label_display = gr.Image(label="Actual Observation", height=400, width=400)
-            
-    gr.Markdown(
-        "--- \n"
-        "**Note:** The first time you run a forecast, the app will download the 366M-parameter model (~1.4 GB) and test data. Subsequent runs will be much faster. "
-        "The images are downscaled for display in this demo. "
-        "For more information, visit the [Surya Hugging Face Repository](https://huggingface.co/nasa-ibm-ai4science/Surya-1.0)."
-    )
+            gr.Markdown("What the Sun *actually* looked like at T+0.")
+            label_display = gr.Image(label="Ground Truth", height=512, width=512, interactive=False)
     
+    # --- Event Handlers ---
     run_button.click(
-        fn=run_forecast,
+        fn=forecast_controller,
         inputs=[channel_selector],
-        outputs=[input_display, prediction_display, label_display]
+        outputs=[input_display, prediction_display, label_display, status_box, state_input, state_prediction, state_target]
+    )
+    
+    channel_selector.change(
+        fn=create_visualizations,
+        inputs=[channel_selector, state_input, state_prediction, state_target],
+        outputs=[input_display, prediction_display, label_display, status_box]
     )
 
 if __name__ == "__main__":
+    # The 'surya_dependencies.py' file must be in the same directory as this script.
+    # Create the placeholder file if it doesn't exist.
+    if not os.path.exists("surya_dependencies.py"):
+        raise FileNotFoundError("The required 'surya_dependencies.py' file is missing. Please download it from the provided source.")
     demo.launch(debug=True)