Update app.py

app.py

@@ -4,24 +4,19 @@ from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
 import os
-# import spaces # No longer needed if running purely on CPU and not using @spaces.GPU()
 from huggingface_hub import hf_hub_download
 
-# --- Configuration ---
-MODEL_PATH = "fno_ckpt_single_res" # This model file still needs to be in your Space's repo
-HF_DATASET_REPO_ID = "ajsbsd/navier-stokes-2d-dataset" # Your new repo ID
+MODEL_PATH = "fno_ckpt_single_res"
+HF_DATASET_REPO_ID = "ajsbsd/navier-stokes-2d-dataset"
 HF_DATASET_FILENAME = "navier_stokes_2d.pt"
 
-# --- Global Variables for Model and Data (loaded once) ---
 MODEL = None
 FULL_DATASET_X = None
 
-# --- Function to Download Dataset from HF Hub ---
 def download_file_from_hf_hub(repo_id, filename):
     """Downloads a file from Hugging Face Hub."""
     print(f"Downloading {filename} from {repo_id} on Hugging Face Hub...")
     try:
-        # hf_hub_download returns the local path to the downloaded file
         local_path = hf_hub_download(repo_id=repo_id, filename=filename)
         print(f"Downloaded {filename} to {local_path} successfully.")
         return local_path
@@ -29,8 +24,6 @@ def download_file_from_hf_hub(repo_id, filename):
         print(f"Error downloading file from HF Hub: {e}")
         raise gr.Error(f"Failed to download dataset from Hugging Face Hub: {e}")
 
-
-# --- 1. Model Loading Function (Loads to CPU, device transfer handled in run_inference) ---
 def load_model():
     """Loads the pre-trained FNO model to CPU."""
     global MODEL
@@ -45,7 +38,6 @@ def load_model():
             raise gr.Error(f"Failed to load model: {e}")
     return MODEL
 
-# --- 2. Dataset Loading Function ---
 def load_dataset():
     """Downloads and loads the initial conditions dataset from HF Hub."""
     global FULL_DATASET_X
@@ -56,11 +48,7 @@ def load_dataset():
             data = torch.load(local_dataset_path, map_location='cpu')
             if isinstance(data, dict) and 'x' in data:
                 FULL_DATASET_X = data['x']
-            elif isinstance(dYou can easily add that blurb by inserting a `gr.Markdown()` component within the same `gr.Column()` as your `sample_input_slider` and `run_button`. This effectively places it within Gradio's "flexbox" layout, ensuring it's always visible below the slider and button.
-
-Here's your `app.py` code with the blurb added in the correct place. I've also updated the `run_inference` function to explicitly target `torch.device("cpu")` and removed the `@spaces.GPU()` decorator, which aligns with your successful run on ZeroCPU.
-
-```pythonata, torch.Tensor):
+            elif isinstance(data, torch.Tensor):
                 FULL_DATASET_X = data
             else:
                 raise ValueError("Unknown dataset format or 'x' key missing.")
@@ -70,41 +58,34 @@ Here's your `app.py` code with the blurb added in the correct place. I've also u
             raise gr.Error(f"Failed to load dataset from local file: {e}")
     return FULL_DATASET_X
 
-# --- 3. Inference Function for Gradio ---
-# Removed @spaces.GPU() decorator as you're running on ZeroCPU
 def run_inference(sample_index: int):
     """
     Performs inference for a selected sample index from the dataset on CPU.
     Returns two Matplotlib figures: one for input, one for output.
     """
-    # Determine the target device (always CPU for ZeroCPU space)
-    device = torch.device("cpu") # Explicitly set to CPU as you're on ZeroCPU
+    device = torch.device("cpu")
 
-    model = load_model() # Model is initially loaded to CPU
+    model = load_model()
 
-    # Model device check is still good practice, even if always CPU here
     if next(model.parameters()).device != device:
         model.to(device)
-        print(f"Model moved to {device} within run_inference.") # Will now print 'Model moved to cpu...'
+        print(f"Model moved to {device} within run_inference.")
 
     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}.")
 
-    # Move input tensor to the correct device
     single_initial_condition = dataset[sample_index:sample_index+1, :, :].unsqueeze(1).to(device)
-    print(f"Input moved to {device}.") # Will now print 'Input moved to cpu.'
+    print(f"Input moved to {device}.")
 
     print(f"Running inference for sample index {sample_index}...")
-    with torch.no_grad(): # Disable gradient calculations for inference
+    with torch.no_grad():
         predicted_solution = model(single_initial_condition)
 
-    # Move results back to CPU for plotting with Matplotlib (already on CPU now)
     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})")
@@ -119,7 +100,6 @@ def run_inference(sample_index: int):
 
     return fig_input, fig_output
 
-# --- Gradio Interface Setup (MODIFIED to add blurb) ---
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -140,14 +120,12 @@ with gr.Blocks() as demo:
             )
             run_button = gr.Button("Generate Solution")
 
-            # --- ADDED BLURB HERE ---
             gr.Markdown(
                 """
                 ### Project Inspiration
                 This Hugging Face Space demonstrates the concepts and models from the research paper **'Principled approaches for extending neural architectures to function spaces for operator learning'** (available as a preprint on [arXiv](https://arxiv.org/abs/2506.10973)). The underlying code for the neural operators and the experiments can be explored further in the associated [GitHub repository](https://github.com/neuraloperator/NNs-to-NOs). The Navier-Stokes dataset used for training and inference, crucial for these fluid dynamics simulations, is openly accessible and citable via [Zenodo](https://zenodo.org/records/12825163).
                 """
             )
-            # --- END ADDED BLURB ---
 
         with gr.Column():
             input_image_plot = gr.Plot(label="Selected Initial Condition")
@@ -160,10 +138,8 @@ with gr.Blocks() as demo:
     )
 
     def load_initial_data_and_predict():
-        # These functions are called during main process startup (CPU)
         load_model()
         load_dataset()
-        # The actual inference call here will now run on CPU
         return run_inference(0)
 
     demo.load(load_initial_data_and_predict, inputs=None, outputs=[input_image_plot, output_image_plot])