import gradio as gr import torch 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 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 except Exception as e: 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 if MODEL is None: print("Loading FNO model to CPU...") try: MODEL = torch.load(MODEL_PATH, weights_only=False, map_location='cpu') MODEL.eval() # Set to evaluation mode print("Model loaded successfully to CPU.") 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(): """Downloads and loads the initial conditions dataset from HF Hub.""" global FULL_DATASET_X if FULL_DATASET_X is None: local_dataset_path = download_file_from_hf_hub(HF_DATASET_REPO_ID, HF_DATASET_FILENAME) print("Loading dataset from local file...") try: 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): 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 Exception as e: print(f"Error loading dataset: {e}") 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 model = load_model() # Model is initially loaded to CPU # 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...' 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"Running inference for sample index {sample_index}...") with torch.no_grad(): # Disable gradient calculations for inference 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})") fig_input.colorbar(im_input, ax=ax_input, label="Vorticity") 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) return fig_input, fig_output # --- Gradio Interface Setup (MODIFIED to add blurb) --- 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, value=0, step=1, 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") run_button.click( fn=run_inference, inputs=[sample_input_slider], outputs=[input_image_plot, output_image_plot] ) 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]) if __name__ == "__main__": demo.launch()