Upload app.py with huggingface_hub

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 import sys
 import os
+from huggingface_hub import hf_hub_download
 
 sys.path.insert(0, os.path.dirname(__file__))
 
@@ -15,9 +16,14 @@ print(f"Loading on {device}")
 config = NexusConfig()
 model = Nexus(config)
 
-weights_path = os.path.join(os.path.dirname(__file__), "weights", "nexus_instruct.pt")
-if not os.path.exists(weights_path):
-    weights_path = os.path.join(os.path.dirname(__file__), "weights", "nexus_final.pt")
+REPO = "JustScriptzz/nexus-smAll-v1"
+
+weights_local = os.path.join(os.path.dirname(__file__), "weights", "nexus_instruct.pt")
+if os.path.exists(weights_local):
+    weights_path = weights_local
+else:
+    print("Downloading weights from HuggingFace...")
+    weights_path = hf_hub_download(repo_id=REPO, filename="weights/nexus_instruct.pt")
 
 checkpoint = torch.load(weights_path, map_location=device, weights_only=False)
 model.load_state_dict(checkpoint["model_state_dict"])
@@ -25,7 +31,11 @@ model = model.to(device)
 model.eval()
 print("Model loaded")
 
-tokenizer_path = os.path.join(os.path.dirname(__file__), "data", "tokenizer.json")
+tokenizer_local = os.path.join(os.path.dirname(__file__), "data", "tokenizer.json")
+if os.path.exists(tokenizer_local):
+    tokenizer_path = tokenizer_local
+else:
+    tokenizer_path = hf_hub_download(repo_id=REPO, filename="data/tokenizer.json")
 tokenizer = Tokenizer.from_file(tokenizer_path)
 
 bos_id = tokenizer.token_to_id("<bos>") or 1