Upload folder using huggingface_hub

app.py

@@ -4,9 +4,11 @@ import os
 import subprocess
 import sys
 from datetime import datetime
+from pathlib import Path
 
 import gradio as gr
 import torch
+from huggingface_hub import hf_hub_download
 
 from config import CONFIG
 from inference import _resolve_device, load_model, run_inference, _decode_clean, _decode_with_cleanup
@@ -17,6 +19,24 @@ RESULTS_DIR = "generated_results"
 DEFAULT_ANALYSIS_OUT = "analysis/outputs"
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
+HF_DEFAULT_MODEL_REPO = os.environ.get("HF_DEFAULT_MODEL_REPO", "bhsinghgrid/DevaFlow")
+HF_DEFAULT_MODEL_FILE = os.environ.get("HF_DEFAULT_MODEL_FILE", "best_model.pt")
+
+
+def _download_hf_default_checkpoint():
+    try:
+        cache_dir = Path(".hf_model_cache")
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        ckpt = hf_hub_download(
+            repo_id=HF_DEFAULT_MODEL_REPO,
+            filename=HF_DEFAULT_MODEL_FILE,
+            local_dir=str(cache_dir),
+            local_dir_use_symlinks=False,
+        )
+        return ckpt
+    except Exception:
+        return None
+
 
 def discover_checkpoints():
     found = []
@@ -35,6 +55,17 @@ def discover_checkpoints():
                     "root": root,
                 }
             )
+    # Space-safe fallback: always expose one downloadable checkpoint option.
+    hf_ckpt = _download_hf_default_checkpoint()
+    if hf_ckpt and os.path.exists(hf_ckpt):
+        found.append(
+            {
+                "label": f"HF default  [{HF_DEFAULT_MODEL_REPO}]",
+                "path": hf_ckpt,
+                "experiment": "hf_default",
+                "root": "hf",
+            }
+        )
     return found
 
 
@@ -253,6 +284,20 @@ def generate_from_ui(
 
 def _run_analysis_cmd(task, ckpt_path, output_dir, input_text="dharmo rakṣati rakṣitaḥ", phase="analyze"):
     os.makedirs(output_dir, exist_ok=True)
+    # Space-safe Task4 fallback: if ablation models don't exist, bootstrap them
+    # from currently selected checkpoint so Task4 can still execute end-to-end.
+    if str(task) == "4" and phase == "analyze":
+        for t in (4, 8, 16, 32, 64):
+            t_dir = Path("ablation_results") / f"T{t}"
+            t_dir.mkdir(parents=True, exist_ok=True)
+            dst = t_dir / "best_model.pt"
+            if not dst.exists():
+                try:
+                    os.symlink(os.path.abspath(ckpt_path), str(dst))
+                except Exception:
+                    import shutil
+                    shutil.copy2(ckpt_path, str(dst))
+
     cmd = [
         sys.executable,
         "analysis/run_analysis.py",