update app.py, add choice of base and large models

app.py

@@ -10,10 +10,11 @@ import copy
 import uuid
 import shutil
 import time
-import threading # <-- NEW: Import for background tasks
+import threading
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
+# Assuming visionts package is available
 from visionts import VisionTSpp, freq_to_seasonality_list
 
 # ========================
@@ -60,6 +61,7 @@ def periodic_cleanup_task():
         cleanup_old_sessions()
         time.sleep(CLEANUP_INTERVAL_SECONDS)
 
+
 # ========================
 # 1. Model Configuration
 # ========================
@@ -67,22 +69,68 @@ def periodic_cleanup_task():
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 REPO_ID = "Lefei/VisionTSpp"
 LOCAL_DIR = "./hf_models/VisionTSpp"
-CKPT_PATH = os.path.join(LOCAL_DIR, "visiontspp_model.ckpt")
-ARCH = 'mae_base'
 
-if not os.path.exists(CKPT_PATH):
-    from huggingface_hub import snapshot_download
-    print("Downloading model from Hugging Face Hub...")
-    snapshot_download(repo_id=REPO_ID, local_dir=LOCAL_DIR, local_dir_use_symlinks=False, resume_download=True)
+# --- Define model configurations ---
+MODEL_CONFIGS = {
+    "base": {
+        "arch": 'mae_base',
+        "ckpt_path": os.path.join(LOCAL_DIR, "visiontspp_base.ckpt")
+    },
+    "large": {
+        "arch": 'mae_large',
+        "ckpt_path": os.path.join(LOCAL_DIR, "visiontspp_large.ckpt")
+    }
+}
+
+# Download both checkpoints if they don't exist
+for model_size, config in MODEL_CONFIGS.items():
+    if not os.path.exists(config["ckpt_path"]):
+        print(f"Downloading {model_size} model from Hugging Face Hub...")
+        snapshot_download(
+            repo_id=REPO_ID,
+            local_dir=LOCAL_DIR,
+            local_dir_use_symlinks=False,
+            resume_download=True,
+            allow_patterns=[f"*{model_size}*"] # Download only relevant files if possible
+        )
 
 QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
-# Assuming VisionTSpp is defined in a separate file or installed package
-# from visionts import VisionTSpp, freq_to_seasonality_list # Placeholder for your model import
-model = VisionTSpp(ARCH, ckpt_path=CKPT_PATH, quantile=True, clip_input=True, complete_no_clip=False, color=True).to(DEVICE)
-print(f"Model loaded on {DEVICE}")
 
-imagenet_mean = np.array([0.485, 0.456, 0.406])
-imagenet_std = np.array([0.229, 0.224, 0.225])
+# --- NEW: Global variables to hold the currently loaded model and its size ---
+CURRENT_MODEL_SIZE = None
+CURRENT_MODEL = None
+
+def load_model_for_size(model_size: str):
+    """Loads the specified VisionTS++ model."""
+    global CURRENT_MODEL, CURRENT_MODEL_SIZE
+    if model_size not in MODEL_CONFIGS:
+        raise ValueError(f"Invalid model size: {model_size}. Available: {list(MODEL_CONFIGS.keys())}")
+
+    config = MODEL_CONFIGS[model_size]
+    print(f"Loading {model_size} model...")
+    model = VisionTSpp(
+        config["arch"],
+        ckpt_path=config["ckpt_path"],
+        quantile=True,
+        clip_input=True,
+        complete_no_clip=False,
+        color=True
+    ).to(DEVICE)
+    print(f"Model {model_size} loaded on {DEVICE}")
+
+    # Unload the previous model to free memory if it was loaded
+    if CURRENT_MODEL is not None:
+        print(f"Unloading previous model ({CURRENT_MODEL_SIZE})...")
+        del CURRENT_MODEL
+        torch.cuda.empty_cache() # Clear GPU cache if using CUDA
+
+    CURRENT_MODEL = model
+    CURRENT_MODEL_SIZE = model_size
+    return model
+
+# Load the default model (base) on startup
+if CURRENT_MODEL is None:
+    CURRENT_MODEL = load_model_for_size("base") # Or "large" as default
 
 
 # ========================
@@ -182,7 +230,7 @@ class PredictionResult:
         self.total_samples = total_samples
         self.inferred_freq = inferred_freq
 
-def predict_at_index(df, index, context_len, pred_len, session_dir):
+def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
     if 'date' not in df.columns:
         raise gr.Error("❌ Input CSV must contain a 'date' column.")
 
@@ -204,7 +252,7 @@ def predict_at_index(df, index, context_len, pred_len, session_dir):
     total_samples = len(data) - context_len - pred_len + 1
     if total_samples <= 0:
         raise gr.Error(f"Data is too short. It needs at least context_len + pred_len = {context_len + pred_len} rows, but has {len(data)}.")
-    
+
     index = max(0, min(index, total_samples - 1))
 
     train_len = int(len(data) * 0.7)
@@ -219,9 +267,17 @@ def predict_at_index(df, index, context_len, pred_len, session_dir):
 
     periodicity_list = freq_to_seasonality_list(inferred_freq)
     periodicity = periodicity_list[0] if periodicity_list else 1
-    
+
     color_list = [i % 3 for i in range(nvars)]
-    model.update_config(context_len=context_len, pred_len=pred_len, periodicity=periodicity, 
+
+    # --- NEW: Load the requested model if it's not the current one ---
+    if CURRENT_MODEL_SIZE != model_size:
+        print(f"Switching model from {CURRENT_MODEL_SIZE} to {model_size}")
+        load_model_for_size(model_size)
+
+    # --- Use the currently loaded model ---
+    model = CURRENT_MODEL
+    model.update_config(context_len=context_len, pred_len=pred_len, periodicity=periodicity,
                         num_patch_input=7, padding_mode='constant')
 
     with torch.no_grad():
@@ -242,7 +298,7 @@ def predict_at_index(df, index, context_len, pred_len, session_dir):
 
     full_true_context = data[start_idx : start_idx + context_len]
     full_true_series = np.concatenate([full_true_context, y_true], axis=0)
-    
+
     ts_fig = visual_ts_with_quantiles(
         true_data=full_true_series, pred_median=pred_median,
         pred_quantiles_list=pred_quantiles, model_quantiles=list(all_preds.keys()),
@@ -275,9 +331,9 @@ def get_session_dir(session_id: gr.State):
         session_id = str(session_dir)
     return session_id
 
-def run_forecast(data_source, upload_file, index, context_len, pred_len, session_id: gr.State):
+def run_forecast(data_source, upload_file, index, context_len, pred_len, model_size, session_id: gr.State):
     session_dir = get_session_dir(session_id)
-    
+
     try:
         if data_source == "Upload CSV":
             if upload_file is None:
@@ -289,10 +345,11 @@ def run_forecast(data_source, upload_file, index, context_len, pred_len, session
             df = load_preset_data(data_source)
 
         index, context_len, pred_len = int(index), int(context_len), int(pred_len)
-        result = predict_at_index(df, index, context_len, pred_len, session_dir)
-        
+        # --- Pass model_size to predict_at_index ---
+        result = predict_at_index(df, index, context_len, pred_len, session_dir, model_size)
+
         final_index = min(index, result.total_samples - 1)
-            
+
         return (
             result.ts_fig,
             result.input_img_fig,
@@ -302,7 +359,7 @@ def run_forecast(data_source, upload_file, index, context_len, pred_len, session
             gr.update(value=result.inferred_freq),
             session_dir
         )
-    
+
     except Exception as e:
         print(f"Error during forecast: {e}")
         error_fig = plt.figure(figsize=(10, 5))
@@ -325,12 +382,19 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
         - ✅ **Slide** through different samples of the dataset for real-time forecasting.
         - ✅ **Download** the prediction results as a CSV file.
         - ✅ **User Isolation**: Each user session has its own temporary storage to prevent file conflicts. Old files are automatically cleaned up.
+        - ✅ **Model Selection**: Choose between 'base' and 'large' VisionTS++ models.
         """
     )
-    
+
     with gr.Row():
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("### 1. Data & Model Configuration")
+            # --- NEW: Add model selection dropdown ---
+            model_size = gr.Dropdown(
+                label="Select Model Size",
+                choices=["base", "large"],
+                value="base" # Default to base
+            )
             data_source = gr.Dropdown(
                 label="Select Data Source",
                 choices=list(PRESET_DATASETS.keys()) + ["Upload CSV"],
@@ -344,13 +408,13 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
                 2. Must contain a time column named `date` with a consistent frequency.
                 """
             )
-            
+
             context_len = gr.Number(label="Context Length (History)", value=336)
             pred_len = gr.Number(label="Prediction Length (Future)", value=96)
             freq_display = gr.Textbox(label="Detected Frequency", interactive=False)
 
             run_btn = gr.Button("🚀 Run Forecast", variant="primary")
-            
+
             gr.Markdown("### 2. Sample Selection")
             sample_index = gr.Slider(
                 label="Sample Index",
@@ -395,7 +459,8 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
 
     data_source.change(fn=toggle_upload_visibility, inputs=data_source, outputs=upload_file)
 
-    inputs = [data_source, upload_file, sample_index, context_len, pred_len, session_id_state]
+    # --- NEW: Include model_size in the inputs list ---
+    inputs = [data_source, upload_file, sample_index, context_len, pred_len, model_size, session_id_state]
     outputs = [ts_plot, input_img_plot, recon_img_plot, download_csv, sample_index, freq_display, session_id_state]
 
     run_btn.click(fn=run_forecast, inputs=inputs, outputs=outputs, api_name="run_forecast")
@@ -408,10 +473,10 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
 if __name__ == "__main__":
     # --- Run initial cleanup on startup ---
     cleanup_old_sessions()
-    
+
     # --- NEW: Start the periodic cleanup in a background daemon thread ---
     cleanup_thread = threading.Thread(target=periodic_cleanup_task, daemon=True)
     cleanup_thread.start()
 
     # --- Launch the Gradio app ---
-    demo.launch(debug=True)
+    demo.launch(debug=True)