Update app.py

app.py

@@ -4,14 +4,15 @@ from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
 import os
-import requests # <--- ADD THIS IMPORT for downloading files
-from tqdm import tqdm # Optional: for a progress bar during download
+import requests
+from tqdm import tqdm
+from huggingface_hub import HfApi, HfFolder, Repository, create_repo # <--- ADD THIS IMPORT
+import spaces # <--- ADD THIS IMPORT
 
 # --- Configuration ---
-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
+MODEL_PATH = "fno_ckpt_single_res"
 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
+LOCAL_DATASET_PATH = "navier_stokes_2d.pt"
 
 # --- Global Variables for Model and Data (loaded once) ---
 MODEL = None
@@ -27,10 +28,10 @@ def download_file(url, local_filename):
     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)
+        response.raise_for_status()
 
         total_size = int(response.headers.get('content-length', 0))
-        block_size = 1024 # 1 KB
+        block_size = 1024
 
         with open(local_filename, 'wb') as f:
             with tqdm(total=total_size, unit='iB', unit_scale=True, desc=local_filename) as pbar:
@@ -43,15 +44,21 @@ def download_file(url, local_filename):
         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) ---
+# --- 1. Model Loading Function ---
 def load_model():
     """Loads the pre-trained FNO model."""
     global MODEL
     if MODEL is None:
         print("Loading FNO model...")
         try:
+            # Load to CPU, then move to GPU if available and needed
             MODEL = torch.load(MODEL_PATH, weights_only=False, map_location='cpu')
+            # Move model to GPU if available
+            if torch.cuda.is_available():
+                MODEL.cuda()
+                print("Model moved to GPU.")
+            else:
+                print("CUDA not available. Model will run on CPU.")
             MODEL.eval()
             print("Model loaded successfully.")
         except Exception as e:
@@ -59,12 +66,12 @@ def load_model():
             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."""
     global FULL_DATASET_X
     if FULL_DATASET_X is None:
-        download_file(DATASET_URL, LOCAL_DATASET_PATH) # <--- Download here!
+        download_file(DATASET_URL, LOCAL_DATASET_PATH)
         print("Loading dataset from local file...")
         try:
             data = torch.load(LOCAL_DATASET_PATH, map_location='cpu')
@@ -80,27 +87,40 @@ 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 change) ---
+# --- 3. Inference Function for Gradio (MODIFIED with @spaces.GPU()) ---
+@spaces.GPU() # <--- ADD THIS DECORATOR
 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() # This will trigger download and load if not already done
+    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}.")
 
+    # 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)
 
+    # Move input tensor to GPU if model is on GPU
+    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.")
+
+
     print(f"Running inference for sample index {sample_index}...")
     with torch.no_grad():
         predicted_solution = model(single_initial_condition)
 
+    # 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})")
@@ -127,8 +147,6 @@ 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,
@@ -148,8 +166,10 @@ with gr.Blocks() as demo:
     )
 
     def load_initial_data_and_predict():
+        # Ensure model and dataset are loaded when the space starts
         load_model()
-        load_dataset() # This will now download if not present
+        load_dataset()
+        # Run inference for the default value (index 0)
         return run_inference(0)
 
     demo.load(load_initial_data_and_predict, inputs=None, outputs=[input_image_plot, output_image_plot])