update app.py, add choice button for VisionTSpp base and large

app.py

@@ -221,10 +221,6 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
         # pred_range = np.arange(context_len, context_len + pred_len)
         pred_range = np.arange(context_len-1, context_len + pred_len)
 
-        print(true_data[:, i].shape)
-        print(pred_median[:, i].shape)
-        print(pred_range.shape)
-
         pred_median_visual = np.concatenate([true_data[context_len-1:context_len, i], pred_median[:, i]])
         print(pred_median_visual.shape)
 
@@ -236,11 +232,6 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
             lower_quantile_pred_visual = np.concatenate([true_data[context_len-1:context_len, i], lower_quantile_pred])
             upper_quantile_pred_visual = np.concatenate([true_data[context_len-1:context_len, i], upper_quantile_pred])
 
-            print(lower_quantile_pred.shape)
-            print(upper_quantile_pred.shape)
-            print(lower_quantile_pred_visual.shape)
-            print(upper_quantile_pred_visual.shape)
-
             q_low, q_high = quantile_vals[j], quantile_vals[-(j+1)]
             ax.fill_between(pred_range, lower_quantile_pred_visual, upper_quantile_pred_visual, color=quantile_colors[j], alpha=0.7, label=f'{int(q_low*100)}-{int(q_high*100)}% Quantile')
             # ax.fill_between(pred_range, lower_quantile_pred, upper_quantile_pred, color=quantile_colors[j], alpha=0.7, label=f'{int(q_low*100)}-{int(q_high*100)}% Quantile')
@@ -387,40 +378,39 @@ def get_session_dir(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:
-            raise gr.Error("Please upload a CSV file when 'Upload CSV' is selected.")
-        uploaded_file_path = Path(session_dir) / Path(upload_file.name).name
-        shutil.copy(upload_file.name, uploaded_file_path)
-        df = pd.read_csv(uploaded_file_path)
-    else:
-        df = load_preset_data(data_source)
-
-    index, context_len, pred_len = int(index), int(context_len), int(pred_len)
-    # --- 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,
-        result.recon_img_fig,
-        result.csv_path,
-        gr.update(maximum=result.total_samples - 1, value=final_index),
-        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))
-    #     plt.text(0.5, 0.5, f"An error occurred:\n{str(e)}", ha='center', va='center', wrap=True, color='red', fontsize=12)
-    #     plt.axis('off')
-    #     plt.close(error_fig)
-    #     return error_fig, None, None, None, gr.update(), gr.update(value="Error"), session_id
+    try:
+        if data_source == "Upload CSV":
+            if upload_file is None:
+                raise gr.Error("Please upload a CSV file when 'Upload CSV' is selected.")
+            uploaded_file_path = Path(session_dir) / Path(upload_file.name).name
+            shutil.copy(upload_file.name, uploaded_file_path)
+            df = pd.read_csv(uploaded_file_path)
+        else:
+            df = load_preset_data(data_source)
+
+        index, context_len, pred_len = int(index), int(context_len), int(pred_len)
+        # --- 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,
+            result.recon_img_fig,
+            result.csv_path,
+            gr.update(maximum=result.total_samples - 1, value=final_index),
+            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))
+        plt.text(0.5, 0.5, f"An error occurred:\n{str(e)}", ha='center', va='center', wrap=True, color='red', fontsize=12)
+        plt.axis('off')
+        plt.close(error_fig)
+        return error_fig, None, None, None, gr.update(), gr.update(value="Error"), session_id
 
 
 with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes.Soft()) as demo: