Update app.py

app.py

@@ -1,25 +1,67 @@
 import os
+import traceback
+import gradio as gr
 import torch
 from huggingface_hub import hf_hub_download
 
-MODEL_REPO_ID = "farrell236/CephVIT"
+# Change this
+MODEL_REPO_ID = "your-username/your-private-model-repo"
 MODEL_FILENAME = "best.pt"
 
-hf_token = os.getenv("HF_TOKEN")
-assert hf_token, "HF_TOKEN missing"
 
-local_path = hf_hub_download(
-    repo_id=MODEL_REPO_ID,
-    filename=MODEL_FILENAME,
-    token=hf_token,
-)
+def inspect_checkpoint():
+    try:
+        hf_token = os.getenv("HF_TOKEN")
+        if not hf_token:
+            return "ERROR: HF_TOKEN is missing. Add it in Space Settings -> Secrets."
+
+        local_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID,
+            filename=MODEL_FILENAME,
+            token=hf_token,
+        )
+
+        lines = []
+        lines.append("Download successful.")
+        lines.append(f"Local path: {local_path}")
+
+        ckpt = torch.load(local_path, map_location="cpu")
+
+        lines.append("")
+        lines.append(f"Top-level object type: {type(ckpt).__name__}")
+
+        if isinstance(ckpt, dict):
+            top_keys = list(ckpt.keys())
+            lines.append(f"Top-level key count: {len(top_keys)}")
+            lines.append("Top-level keys:")
+            for k in top_keys[:50]:
+                lines.append(f"  - {k}")
 
-print("Downloaded to:", local_path)
+            if "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
+                sd_keys = list(ckpt["state_dict"].keys())
+                lines.append("")
+                lines.append(f"state_dict key count: {len(sd_keys)}")
+                lines.append("First 20 state_dict keys:")
+                for k in sd_keys[:20]:
+                    lines.append(f"  - {k}")
+
+        else:
+            lines.append("Checkpoint is not a dict, so no keys to print.")
+
+        return "\n".join(lines)
+
+    except Exception as e:
+        return f"ERROR:\n{type(e).__name__}: {e}\n\n{traceback.format_exc()}"
+
+
+demo = gr.Interface(
+    fn=inspect_checkpoint,
+    inputs=None,
+    outputs=gr.Textbox(label="Checkpoint inspection", lines=30),
+    title="Private checkpoint test",
+    description="Checks whether best.pt can be downloaded from a private Hugging Face repo and inspected.",
+)
 
-ckpt = torch.load(local_path, map_location="cpu")
-print("Type:", type(ckpt))
 
-if isinstance(ckpt, dict):
-    print("Top-level keys:", list(ckpt.keys())[:20])
-    if "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
-        print("state_dict keys:", list(ckpt["state_dict"].keys())[:20])
+if __name__ == "__main__":
+    demo.launch()