Update app.py

app.py

@@ -1,43 +1,73 @@
 import gradio as gr
 import torch
-from neuralop.models import FNO # Make sure this import path matches your environment
+from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
-import os # For path handling
+import os
+import requests # <--- ADD THIS IMPORT for downloading files
+from tqdm import tqdm # Optional: for a progress bar during download
 
 # --- 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
+MODEL_PATH = "fno_ckpt_single_res" # This model file still needs to be in your repo
+# Zenodo direct download URL for the Navier-Stokes 2D dataset
+DATASET_URL = "https://zenodo.org/record/12825163/files/navier_stokes_2d.pt?download=1"
+LOCAL_DATASET_PATH = "navier_stokes_2d.pt" # Where the file will be saved locally in the Space
 
 # --- Global Variables for Model and Data (loaded once) ---
 MODEL = None
-FULL_DATASET_X = None # To store all initial conditions
-
-# --- 1. Model Loading Function ---
+FULL_DATASET_X = None
+
+# --- Function to Download Dataset ---
+def download_file(url, local_filename):
+    """Downloads a file from a URL to a local path with a progress bar."""
+    if os.path.exists(local_filename):
+        print(f"{local_filename} already exists. Skipping download.")
+        return
+
+    print(f"Downloading {url} to {local_filename}...")
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status() # Raise an HTTPError for bad responses (4xx or 5xx)
+
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 1024 # 1 KB
+
+        with open(local_filename, 'wb') as f:
+            with tqdm(total=total_size, unit='iB', unit_scale=True, desc=local_filename) as pbar:
+                for chunk in response.iter_content(chunk_size=block_size):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+        print(f"Downloaded {local_filename} successfully.")
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading file: {e}")
+        raise gr.Error(f"Failed to download dataset from Zenodo: {e}")
+
+
+# --- 1. Model Loading Function (No change here for model) ---
 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
+            MODEL.eval()
             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 ---
+# --- 2. Dataset Loading Function (MODIFIED) ---
 def load_dataset():
-    """Loads the initial conditions dataset."""
+    """Downloads and loads the initial conditions dataset."""
     global FULL_DATASET_X
     if FULL_DATASET_X is None:
-        print("Loading dataset...")
+        download_file(DATASET_URL, LOCAL_DATASET_PATH) # <--- Download here!
+        print("Loading dataset from local file...")
         try:
-            data = torch.load(DATASET_PATH, map_location='cpu')
+            data = torch.load(LOCAL_DATASET_PATH, map_location='cpu')
             if isinstance(data, dict) and 'x' in data:
                 FULL_DATASET_X = data['x']
             elif isinstance(data, torch.Tensor):
@@ -45,54 +75,47 @@ def load_dataset():
             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}")
+            raise gr.Error(f"Failed to load dataset from local file: {e}")
     return FULL_DATASET_X
 
-# --- 3. Inference Function for Gradio ---
+# --- 3. Inference Function for Gradio (No change) ---
 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()
+    dataset = load_dataset() # This will trigger download and load if not already done
 
     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
+    plt.close(fig_input)
 
     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
+    plt.close(fig_output)
 
     return fig_input, fig_output
 
-# --- Gradio Interface Setup ---
+# --- Gradio Interface Setup (No change) ---
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -104,9 +127,11 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
+            # Max value can be dynamic based on dataset size if needed,
+            # but 9999 for 10,000 samples is correct.
             sample_input_slider = gr.Slider(
                 minimum=0,
-                maximum=9999, # Assuming 10,000 samples based on your dataset shape
+                maximum=9999,
                 value=0,
                 step=1,
                 label="Select Sample Index"
@@ -116,23 +141,18 @@ with gr.Blocks() as demo:
             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)
+        load_dataset() # This will now download if not present
         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()