Update app.py

app.py

@@ -4,77 +4,60 @@ from neuralop.models import FNO
 import matplotlib.pyplot as plt
 import numpy as np
 import os
-import requests
-from tqdm import tqdm
-from huggingface_hub import HfApi, HfFolder, Repository, create_repo # <--- ADD THIS IMPORT
-import spaces # <--- ADD THIS IMPORT
+# 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
 
 # --- Configuration ---
-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"
+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"
 
 # --- Global Variables for Model and Data (loaded once) ---
 MODEL = None
 FULL_DATASET_X = None
 
-# --- Function to Download Dataset ---
-def download_file(url, local_filename):
-    """Downloads a file from a URL to a local path with a progress bar."""
-    if os.path.exists(local_filename):
-        print(f"{local_filename} already exists. Skipping download.")
-        return
-
-    print(f"Downloading {url} to {local_filename}...")
+# --- Function to Download Dataset (MODIFIED to use hf_hub_download) ---
+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:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-
-        total_size = int(response.headers.get('content-length', 0))
-        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:
-                for chunk in response.iter_content(chunk_size=block_size):
-                    if chunk:
-                        f.write(chunk)
-                        pbar.update(len(chunk))
-        print(f"Downloaded {local_filename} successfully.")
-    except requests.exceptions.RequestException as e:
-        print(f"Error downloading file: {e}")
-        raise gr.Error(f"Failed to download dataset from Zenodo: {e}")
-
-# --- 1. Model Loading Function ---
+        # 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 (No change from last successful CUDA fix) ---
 def load_model():
-    """Loads the pre-trained FNO model."""
+    """Loads the pre-trained FNO model to CPU."""
     global MODEL
     if MODEL is None:
-        print("Loading FNO model...")
+        print("Loading FNO model to CPU...")
         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.")
+            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 ---
+# --- 2. Dataset Loading Function (MODIFIED) ---
 def load_dataset():
-    """Downloads and loads the initial conditions dataset."""
+    """Downloads and loads the initial conditions dataset from HF Hub."""
     global FULL_DATASET_X
     if FULL_DATASET_X is None:
-        download_file(DATASET_URL, LOCAL_DATASET_PATH)
+        # 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:
-            data = torch.load(LOCAL_DATASET_PATH, map_location='cpu')
+            data = torch.load(local_dataset_path, map_location='cpu')
             if isinstance(data, dict) and 'x' in data:
                 FULL_DATASET_X = data['x']
             elif isinstance(data, torch.Tensor):
@@ -87,8 +70,8 @@ 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 with @spaces.GPU()) ---
-@spaces.GPU() # <--- ADD THIS DECORATOR
+# --- 3. Inference Function for Gradio (No changes needed here) ---
+@spaces.GPU()
 def run_inference(sample_index: int):
     """
     Performs inference for a selected sample index from the dataset.
@@ -97,30 +80,28 @@ def run_inference(sample_index: int):
     model = load_model()
     dataset = load_dataset()
 
+    if torch.cuda.is_available() and next(model.parameters()).device == torch.device('cpu'):
+        model.cuda()
+        print("Model moved to GPU within run_inference.")
+
     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})")
@@ -166,10 +147,8 @@ 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()
-        # 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])