Update app.py

app.py

@@ -1,10 +1,15 @@
+You 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.
+
+```python
 import gradio as gr
 import torch
 from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
 import os
-import spaces
+# import spaces # No longer needed if running purely on CPU and not using @spaces.GPU()
 from huggingface_hub import hf_hub_download
 
 # --- Configuration ---
@@ -66,39 +71,37 @@ def load_dataset():
             raise gr.Error(f"Failed to load dataset from local file: {e}")
     return FULL_DATASET_X
 
-# --- 3. Inference Function for Gradio (MODIFIED: Explicit device handling) ---
-@spaces.GPU()
+# --- 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.
-    Ensures model and input are on the correct device (GPU).
+    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 (GPU if available, else CPU)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Determine the target device (always CPU for ZeroCPU space)
+    device = torch.device("cpu") # Explicitly set to CPU as you're on ZeroCPU
 
     model = load_model() # Model is initially loaded to CPU
 
-    # Move model to the correct device ONLY when inside the @spaces.GPU() decorated function
-    # and only if it's not already on the target device.
+    # 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.")
+        print(f"Model moved to {device} within run_inference.") # Will now print 'Model moved to cpu...'
 
     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 directly
+    # 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}.")
+    print(f"Input moved to {device}.") # Will now print 'Input moved to cpu.'
 
     print(f"Running inference for sample index {sample_index}...")
     with torch.no_grad(): # Disable gradient calculations for inference
-        predicted_solution = model(single_initial_condition) # This is where the error occurred before
+        predicted_solution = model(single_initial_condition)
 
-    # Move results back to CPU for plotting with Matplotlib
+    # 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()
 
@@ -117,7 +120,7 @@ def run_inference(sample_index: int):
 
     return fig_input, fig_output
 
-# --- Gradio Interface Setup (No change) ---
+# --- Gradio Interface Setup (MODIFIED to add blurb) ---
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -137,6 +140,16 @@ with gr.Blocks() as demo:
                 label="Select Sample Index"
             )
             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")
             output_image_plot = gr.Plot(label="Predicted Solution")
@@ -151,10 +164,12 @@ with gr.Blocks() as demo:
         # These functions are called during main process startup (CPU)
         load_model()
         load_dataset()
-        # The actual inference call here will ensure GPU utilization via @spaces.GPU()
+        # 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])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+
+```