Upload folder using huggingface_hub

app.py

@@ -31,6 +31,7 @@ RESULTS_DIR = "generated_results"
 DEFAULT_ANALYSIS_OUT = "analysis_outputs/outputs_all_models_20260325/T4"
 os.makedirs(RESULTS_DIR, exist_ok=True)
 _BG_JOBS = {}
+_CHECKPOINT_CACHE = None
 
 try:
     import mlflow
@@ -193,6 +194,9 @@ def _download_hf_settings_for_repo(repo_id: str):
 
 
 def discover_checkpoints():
+    global _CHECKPOINT_CACHE
+    if _CHECKPOINT_CACHE is not None:
+        return list(_CHECKPOINT_CACHE)
     found = []
     local_roots = [
         ("ablation_results", "cross_attention"),
@@ -234,7 +238,8 @@ def discover_checkpoints():
                 "hf_settings": settings,
             }
         )
-    return found
+    _CHECKPOINT_CACHE = list(found)
+    return list(found)
 
 
 def _guess_analysis_dir(experiment: str, ckpt_path: str, family: str = "cross_attention", settings: dict | None = None) -> str:
@@ -1106,6 +1111,22 @@ def auto_load_default_with_outputs():
     return load_selected_model_with_outputs(default_checkpoint_label())
 
 
+def safe_load_selected_model_with_outputs(checkpoint_label):
+    try:
+        return load_selected_model_with_outputs(checkpoint_label)
+    except Exception as e:
+        outputs = _safe_refresh_task_outputs(DEFAULT_ANALYSIS_OUT)
+        return (
+            None,
+            f"Load failed: {e}",
+            {},
+            64,
+            DEFAULT_ANALYSIS_OUT,
+            _build_flow_markdown(model_loaded=False, inference_ready=False, task_states={}),
+            *outputs,
+        )
+
+
 CUSTOM_CSS = """
 :root {
   --bg1: #f5fbff;
@@ -1258,6 +1279,8 @@ with gr.Blocks(title="Sanskrit Diffusion Model", css=CUSTOM_CSS) as demo:
             )
 
     def refresh_checkpoints():
+        global _CHECKPOINT_CACHE
+        _CHECKPOINT_CACHE = None
         choices = list(checkpoint_map().keys())
         value = default_checkpoint_label() if choices else None
         msg = f"Found {len(choices)} checkpoint(s)." if choices else "No checkpoints found."
@@ -1265,7 +1288,7 @@ with gr.Blocks(title="Sanskrit Diffusion Model", css=CUSTOM_CSS) as demo:
 
     refresh_btn.click(fn=refresh_checkpoints, outputs=[checkpoint_dropdown, load_status])
     load_btn.click(
-        fn=load_selected_model_with_outputs,
+        fn=safe_load_selected_model_with_outputs,
         inputs=[checkpoint_dropdown],
         outputs=[
             model_state,