update app.py

app.py

@@ -16,7 +16,10 @@ from visionts import VisionTSpp, freq_to_seasonality_list
 # ========================
 # 1. Configuration
 # ========================
+
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+# DEVICE = 'cpu'
+
 REPO_ID = "Lefei/VisionTSpp"
 LOCAL_DIR = "./hf_models/VisionTSpp"
 CKPT_PATH = os.path.join(LOCAL_DIR, "visiontspp_model.ckpt")
@@ -121,9 +124,15 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
 
     print(f"{len(pred_quantiles_list) = }")
     print(f"{len(model_quantiles) = }")
+    print(f"{model_quantiles = }")
     print(f"{pred_quantiles_list[0].shape = }")
 
-    sorted_quantiles = sorted(zip(model_quantiles, pred_quantiles_list + [pred_median]), key=lambda x: x[0])
+    # sorted_quantiles = sorted(zip(model_quantiles, pred_quantiles_list + [pred_median]), key=lambda x: x[0])
+    # sorted_quantiles = sorted(zip(model_quantiles, pred_quantiles_list), key=lambda x: x[0])
+
+    pred_quantiles_list.insert(len(QUANTILES)//2, pred_median)
+    sorted_quantiles = sorted(zip(QUANTILES, pred_quantiles_list), key=lambda x: x[0])
+
     quantile_preds = [item[1] for item in sorted_quantiles if item[0] != 0.5]
     quantile_vals = [item[0] for item in sorted_quantiles if item[0] != 0.5]
     
@@ -149,6 +158,7 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
     handles, labels = axes[0].get_legend_handles_labels()
     unique_labels = dict(zip(labels, handles))
     fig.legend(unique_labels.values(), unique_labels.keys(), loc='upper center', bbox_to_anchor=(0.5, 1.05), ncol=num_bands + 2)
+    
     plt.tight_layout(rect=[0, 0, 1, 0.95])
     plt.close(fig)
     
@@ -227,12 +237,9 @@ def predict_at_index(df, index, context_len, pred_len):
         print(f"{reconstructed_image.shape = }")
         print(f"{len(y_pred_quantile_list) = }")
 
-        # print(f"{input_image = }")
-        print(f"{input_image[0,0,0, :, 0] = }")
-        print(f"{input_image[0,0,0, 50:70, 0] = }")
-        print(f"{input_image[0,0,0, 100:120, 0] = }")
-        # print(f"{input_image[0] = }")
-        # print(f"{reconstructed_image = }")
+        # print(f"{input_image[0,0,0, :, 0] = }")
+        # print(f"{input_image[0,0,0, 50:70, 0] = }")
+        # print(f"{input_image[0,0,0, 100:120, 0] = }")
 
         all_y_pred_list = copy.deepcopy(y_pred_quantile_list)