Create app.py

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+import torch
+from neuralop.models import FNO # Make sure this import path matches your environment
+import matplotlib.pyplot as plt
+import numpy as np
+import os # For path handling
+
+# --- Configuration ---
+# Set paths relative to the root of your Hugging Face Space repository
+MODEL_PATH = "fno_ckpt_single_res"
+DATASET_PATH = "navier_stokes_2d.pt" # Ensure this file is in your repo root
+
+# --- Global Variables for Model and Data (loaded once) ---
+MODEL = None
+FULL_DATASET_X = None # To store all initial conditions
+
+# --- 1. Model Loading Function ---
+def load_model():
+    """Loads the pre-trained FNO model."""
+    global MODEL
+    if MODEL is None:
+        print("Loading FNO model...")
+        try:
+            # Ensure model is loaded to CPU for general compatibility in Spaces
+            MODEL = torch.load(MODEL_PATH, weights_only=False, map_location='cpu')
+            MODEL.eval() # Set to evaluation mode
+            print("Model loaded successfully.")
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise gr.Error(f"Failed to load model: {e}")
+    return MODEL
+
+# --- 2. Dataset Loading Function ---
+def load_dataset():
+    """Loads the initial conditions dataset."""
+    global FULL_DATASET_X
+    if FULL_DATASET_X is None:
+        print("Loading dataset...")
+        try:
+            data = torch.load(DATASET_PATH, map_location='cpu')
+            if isinstance(data, dict) and 'x' in data:
+                FULL_DATASET_X = data['x']
+            elif isinstance(data, torch.Tensor):
+                FULL_DATASET_X = data
+            else:
+                raise ValueError("Unknown dataset format or 'x' key missing.")
+            print(f"Dataset loaded. Total samples: {FULL_DATASET_X.shape[0]}")
+        except FileNotFoundError:
+            print(f"Dataset file not found at {DATASET_PATH}")
+            raise gr.Error(f"Dataset file not found. Please ensure '{DATASET_PATH}' is in your Space.")
+        except Exception as e:
+            print(f"Error loading dataset: {e}")
+            raise gr.Error(f"Failed to load dataset: {e}")
+    return FULL_DATASET_X
+
+# --- 3. Inference Function for Gradio ---
+def run_inference(sample_index: int):
+    """
+    Performs inference for a selected sample index from the dataset.
+    Returns two Matplotlib figures: one for input, one for output.
+    """
+    model = load_model()
+    dataset = load_dataset()
+
+    if not (0 <= sample_index < dataset.shape[0]):
+        raise gr.Error(f"Sample index out of range. Please choose between 0 and {dataset.shape[0]-1}.")
+
+    # Extract single initial condition and add channel dimension
+    # (shape: [1, H, W] -> [1, 1, H, W])
+    single_initial_condition = dataset[sample_index:sample_index+1, :, :].unsqueeze(1)
+
+    print(f"Running inference for sample index {sample_index}...")
+    with torch.no_grad():
+        predicted_solution = model(single_initial_condition)
+
+    # Convert tensors to numpy for plotting
+    input_numpy = single_initial_condition.squeeze().cpu().numpy()
+    output_numpy = predicted_solution.squeeze().cpu().numpy()
+
+    # Create Matplotlib figures
+    fig_input, ax_input = plt.subplots()
+    im_input = ax_input.imshow(input_numpy, cmap='viridis')
+    ax_input.set_title(f"Initial Condition (Sample {sample_index})")
+    fig_input.colorbar(im_input, ax=ax_input, label="Vorticity")
+    plt.close(fig_input) # Close to prevent display issues in Gradio
+
+    fig_output, ax_output = plt.subplots()
+    im_output = ax_output.imshow(output_numpy, cmap='viridis')
+    ax_output.set_title(f"Predicted Solution")
+    fig_output.colorbar(im_output, ax=ax_output, label="Vorticity")
+    plt.close(fig_output) # Close to prevent display issues in Gradio
+
+    return fig_input, fig_output
+
+# --- Gradio Interface Setup ---
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Fourier Neural Operator (FNO) for Navier-Stokes Equations
+        Select a sample index from the pre-loaded dataset to see the FNO's prediction
+        of the vorticity field evolution.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            sample_input_slider = gr.Slider(
+                minimum=0,
+                maximum=9999, # Assuming 10,000 samples based on your dataset shape
+                value=0,
+                step=1,
+                label="Select Sample Index"
+            )
+            run_button = gr.Button("Generate Solution")
+        with gr.Column():
+            input_image_plot = gr.Plot(label="Selected Initial Condition")
+            output_image_plot = gr.Plot(label="Predicted Solution")
+
+    # Bind the button click to the inference function
+    run_button.click(
+        fn=run_inference,
+        inputs=[sample_input_slider],
+        outputs=[input_image_plot, output_image_plot]
+    )
+
+    # Optional: Load initial data on startup for the first display
+    def load_initial_data_and_predict():
+        # Ensure model and dataset are loaded when the space starts
+        load_model()
+        load_dataset()
+        # Run inference for the default value (index 0)
+        return run_inference(0)
+
+    demo.load(load_initial_data_and_predict, inputs=None, outputs=[input_image_plot, output_image_plot])
+
+# Launch the Gradio app (only runs when you test locally)
+if __name__ == "__main__":
+    demo.launch()