Update app.py

app.py

@@ -4,14 +4,11 @@ from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
 import os
-# import requests # <--- NO LONGER NEEDED for Zenodo download
-# from tqdm import tqdm # <--- NO LONGER NEEDED for Zenodo download
 import spaces
-from huggingface_hub import hf_hub_download # <--- ADD THIS IMPORT
+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
-# Updated: Hugging Face Dataset/Model ID and filename
 HF_DATASET_REPO_ID = "ajsbsd/navier-stokes-2d-dataset" # Your new repo ID
 HF_DATASET_FILENAME = "navier_stokes_2d.pt"
 
@@ -19,7 +16,7 @@ HF_DATASET_FILENAME = "navier_stokes_2d.pt"
 MODEL = None
 FULL_DATASET_X = None
 
-# --- Function to Download Dataset (MODIFIED to use hf_hub_download) ---
+# --- 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...")
@@ -33,7 +30,7 @@ def download_file_from_hf_hub(repo_id, filename):
         raise gr.Error(f"Failed to download dataset from Hugging Face Hub: {e}")
 
 
-# --- 1. Model Loading Function (No change from last successful CUDA fix) ---
+# --- 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
@@ -41,19 +38,18 @@ def load_model():
         print("Loading FNO model to CPU...")
         try:
             MODEL = torch.load(MODEL_PATH, weights_only=False, map_location='cpu')
-            MODEL.eval()
+            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 (MODIFIED) ---
+# --- 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:
-        # Call the new HF Hub download function
         local_dataset_path = download_file_from_hf_hub(HF_DATASET_REPO_ID, HF_DATASET_FILENAME)
         print("Loading dataset from local file...")
         try:
@@ -70,38 +66,43 @@ def load_dataset():
             raise gr.Error(f"Failed to load dataset from local file: {e}")
     return FULL_DATASET_X
 
-# --- 3. Inference Function for Gradio (No changes needed here) ---
+# --- 3. Inference Function for Gradio (MODIFIED: Explicit device handling) ---
 @spaces.GPU()
 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).
     Returns two Matplotlib figures: one for input, one for output.
     """
-    model = load_model()
-    dataset = load_dataset()
+    # Determine the target device (GPU if available, else CPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = load_model() # Model is initially loaded to CPU
 
-    if torch.cuda.is_available() and next(model.parameters()).device == torch.device('cpu'):
-        model.cuda()
-        print("Model moved to GPU within run_inference.")
+    # 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.
+    if next(model.parameters()).device != device:
+        model.to(device)
+        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}.")
 
-    single_initial_condition = dataset[sample_index:sample_index+1, :, :].unsqueeze(1)
-
-    if torch.cuda.is_available():
-        single_initial_condition = single_initial_condition.cuda()
-        print("Input moved to GPU.")
-    else:
-        print("CUDA not available. Input remains on CPU.")
+    # Move input tensor to the correct device directly
+    single_initial_condition = dataset[sample_index:sample_index+1, :, :].unsqueeze(1).to(device)
+    print(f"Input moved to {device}.")
 
     print(f"Running inference for sample index {sample_index}...")
-    with torch.no_grad():
-        predicted_solution = model(single_initial_condition)
+    with torch.no_grad(): # Disable gradient calculations for inference
+        predicted_solution = model(single_initial_condition) # This is where the error occurred before
 
+    # Move results back to CPU for plotting with Matplotlib
     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})")
@@ -147,8 +148,10 @@ 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 ensure GPU utilization via @spaces.GPU()
         return run_inference(0)
 
     demo.load(load_initial_data_and_predict, inputs=None, outputs=[input_image_plot, output_image_plot])