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

app.py

@@ -221,7 +221,13 @@ 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 = true_data[context_len-1:context_len, i] + pred_median[:, i]
+        print(pred_median_visual.shape)
+
         ax.plot(pred_range, pred_median_visual, label='Prediction (Median)', color='red', linewidth=1.5)
         # ax.plot(pred_range, pred_median[:, i], label='Prediction (Median)', color='red', linewidth=1.5)
 
@@ -229,7 +235,12 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
             lower_quantile_pred, upper_quantile_pred = quantile_preds[j][:, i], quantile_preds[-(j+1)][:, i]
             lower_quantile_pred_visual = true_data[context_len-1:context_len, i] + lower_quantile_pred
             upper_quantile_pred_visual = 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')
@@ -336,6 +347,7 @@ def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
     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()),
@@ -375,39 +387,40 @@ 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
-        )
+    # try:
 
-    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
+    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: