Update app.py

app.py

@@ -1,26 +1,49 @@
 import gradio as gr
 import torch
 from huggingface_hub import hf_hub_download
-from surya.model import Surya
+from surya.model import Surya # This now works because of the file structure
 import numpy as np
 from PIL import Image
+import os
 import warnings
 
 # Suppress warnings for a cleaner demo experience
 warnings.filterwarnings("ignore")
 
-# --- Model Loading ---
+# --- 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]
+
+# --- 2. Caching and Loading the Model and Data ---
 @gr.cache
-def load_model():
+def load_model_and_data():
     """
-    Downloads the pre-trained Surya model weights and initializes the model.
-    This function is cached so the model is only loaded once.
+    Downloads the pre-trained Surya model, the test data, and initializes the model.
+    This function is cached so this happens 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(
         repo_id="nasa-ibm-ai4science/Surya-1.0",
-        filename="surya.366m.v1.pt"
+        filename="surya.366m.v1.pt",
+        local_dir=model_dir
+    )
+    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,
@@ -30,106 +53,120 @@ def load_model():
         attention_blocks=8,
     )
 
+    # Load the weights into the model
+    print("Loading model weights...")
     model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
     model.eval()
-    return model
+    print("Model loaded successfully.")
+
+    # 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
 
-model = load_model()
+    return model, test_input, test_label
 
-# --- Core Prediction Logic ---
-def predict_solar_activity(time_steps, forecast_horizon):
+# --- 3. Helper function for Image Conversion ---
+def tensor_to_image(tensor_slice):
     """
-    Generates a forecast of solar activity using the Surya model.
-    For this demo, we use a dummy input tensor to simulate the model's input.
-    In a real-world scenario, this function would fetch and preprocess
-    actual SDO data for the given time steps.
+    Normalizes a 2D tensor slice and converts it to a PIL Image for display.
     """
-    # Create a dummy input tensor representing a sequence of solar observations
-    # Shape: [batch_size, channels, time_steps, height, width]
-    dummy_input = torch.randn(1, 13, time_steps, 4096, 4096)
+    # Detach tensor from graph, move to CPU, and convert to numpy
+    img_np = tensor_slice.detach().cpu().numpy()
+    
+    # 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)
+    
+    img_array = (img_np * 255).astype(np.uint8)
+    return Image.fromarray(img_array)
 
-    # In a real application, you would replace the dummy input with actual,
-    # preprocessed data from the Solar Dynamics Observatory (SDO).
-    # Preprocessing would involve alignment and normalization as described
-    # in the Surya paper.
 
+# --- 4. Main Prediction and Visualization Function ---
+def run_forecast(channel_name, progress=gr.Progress()):
+    """
+    This function is triggered by the button click in the Gradio interface.
+    It runs the model prediction and generates the images for display.
+    """
+    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()
+    
+    progress(0.5, desc="Running inference on the model...")
+    # Perform the forecast
     with torch.no_grad():
-        # The model's prediction would be based on the forecast_horizon
-        # For this demo, we simulate a prediction by selecting a slice of the input
-        prediction = model(dummy_input)
-
-    # --- Visualization ---
-    # For demonstration, we will visualize one of the output channels.
-    # We will take the last predicted time step.
-    predicted_image_tensor = prediction[0, 0, -1, :, :] # Visualizing the first channel
-
-    # Normalize the tensor to a 0-255 range for image display
-    normalized_tensor = (predicted_image_tensor - predicted_image_tensor.min()) / \
-                        (predicted_image_tensor.max() - predicted_image_tensor.min())
-    image_array = (normalized_tensor * 255).byte().cpu().numpy()
-    predicted_image = Image.fromarray(image_array)
-
-    # For the flare prediction, we'll generate a dummy probability
-    flare_probability = np.random.rand()
-    if flare_probability > 0.5:
-        flare_class = "M-class or X-class Flare"
-        confidence = flare_probability
-    else:
-        flare_class = "No significant flare"
-        confidence = 1 - flare_probability
-
-    return predicted_image, {flare_class: confidence}
-
-# --- Gradio Interface Definition ---
+        prediction = model(test_input)
+
+    progress(0.8, desc="Generating visualizations...")
+    # Get the index of the selected channel
+    channel_index = ALL_CHANNELS.index(channel_name)
+
+    # 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)
+
+    # 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)
+
+    # Extract the corresponding slice from the ground truth label
+    label_slice = test_label[0, channel_index, 0, :, :]
+    label_image = tensor_to_image(label_slice)
+    
+    print(f"Forecast generated for channel: {channel_name}")
+    return input_image, predicted_image, label_image
+
+# --- 5. Building the Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         <div align="center">
-        # ☀️ Surya: Foundation Model for Heliophysics ☀️
-        *A Gradio Demo for NASA's Solar Foundation Model*
+        # ☀️ 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.
         </div>
         """
     )
-    gr.Markdown(
-        "Surya is a 366M-parameter foundation model trained on full-resolution, multi-instrument SDO observations. "
-        "This demo showcases its capability to forecast solar dynamics."
-    )
 
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Prediction Parameters")
-            time_steps_slider = gr.Slider(
-                minimum=1, maximum=10, value=5, step=1,
-                label="Number of Input Time Steps (12-min cadence)",
-                info="Represents the sequence of past solar observations to feed the model."
-            )
-            forecast_horizon_slider = gr.Slider(
-                minimum=1, maximum=24, value=1,
-                label="Forecast Horizon (hours)",
-                info="How far into the future to predict."
-            )
-            predict_button = gr.Button("🔮 Generate Forecast", variant="primary")
-
-        with gr.Column(scale=2):
-            gr.Markdown("### 🛰️ Predicted Solar Image")
-            output_image = gr.Image(label="Forecasted SDO Image (AIA 171 Å)", height=512, width=512)
-            gr.Markdown("### 💥 Solar Flare Prediction")
-            output_flare = gr.Label(label="Flare Probability")
-
-    predict_button.click(
-        fn=predict_solar_activity,
-        inputs=[time_steps_slider, forecast_horizon_slider],
-        outputs=[output_image, output_flare]
-    )
+        channel_selector = gr.Dropdown(
+            choices=ALL_CHANNELS,
+            value=ALL_CHANNELS[2], # Default to "AIA 171 Å"
+            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")
 
-    gr.Markdown("---")
+    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)
+        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)
+        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(
-        "**Note:** This demo uses a placeholder for real-time data fetching and displays a simulated prediction. "
-        "The core of this application is the loaded Surya model from NASA and IBM."
+        "--- \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(
-        "For more information, visit the [Surya model card on Hugging Face](https://huggingface.co/nasa-ibm-ai4science/Surya-1.0)."
+    
+    run_button.click(
+        fn=run_forecast,
+        inputs=[channel_selector],
+        outputs=[input_display, prediction_display, label_display]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)