From ca4274e3848fb66e2db698fbe35fb6bc97fd2aef Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Sat, 17 May 2025 00:31:54 +0800
Subject: [PATCH] added size info into HF searcher

---
 koboldcpp.py | 38 +++++++++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/koboldcpp.py b/koboldcpp.py
index d0626b0bb..b792fd106 100644
--- a/koboldcpp.py
+++ b/koboldcpp.py
@@ -4182,30 +4182,44 @@ def show_gui():
 
     def model_searcher():
         searchbox1 = None
+        searchbox2 = None
         modelsearch1_var = ctk.StringVar(value="")
         modelsearch2_var = ctk.StringVar(value="")
+        fileinfotxt_var = ctk.StringVar(value="")
         # Create popup window
         popup = ctk.CTkToplevel(root)
         popup.title("Model File Browser")
         popup.geometry("400x400")
+        searchedmodels = []
+        searchedsizes = []
 
         def confirm_search_model_choice():
-            nonlocal modelsearch1_var, modelsearch2_var, model_var
+            nonlocal modelsearch1_var, modelsearch2_var, model_var, fileinfotxt_var
             if modelsearch1_var.get()!="" and modelsearch2_var.get()!="":
                 model_var.set(f"https://huggingface.co/{modelsearch1_var.get()}/resolve/main/{modelsearch2_var.get()}")
             popup.destroy()
+        def update_search_quant_file_size(a,b,c):
+            nonlocal modelsearch1_var, modelsearch2_var, fileinfotxt_var, searchedmodels, searchedsizes, searchbox2
+            try:
+                selected_index = searchbox2.cget("values").index(modelsearch2_var.get())
+                pickedsize = searchedsizes[selected_index]
+                fileinfotxt_var.set(f"Size: {round(pickedsize/1024/1024/1024,2)} GB")
+            except Exception:
+                fileinfotxt_var.set("")
         def fetch_search_quants(a,b,c):
-            nonlocal modelsearch1_var, modelsearch2_var
+            nonlocal modelsearch1_var, modelsearch2_var, fileinfotxt_var, searchedmodels, searchedsizes
             try:
                 if modelsearch1_var.get()=="":
                     return
                 searchedmodels = []
-                resp = make_url_request(f"https://huggingface.co/api/models/{modelsearch1_var.get()}",None,'GET',{},10)
-                for m in resp["siblings"]:
-                    if ".gguf" in m["rfilename"]:
-                        if "-of-0" in m["rfilename"] and "00001" not in m["rfilename"]:
+                searchedsizes = []
+                resp = make_url_request(f"https://huggingface.co/api/models/{modelsearch1_var.get()}/tree/main?recursive=true",None,'GET',{},10)
+                for m in resp:
+                    if m["type"]=="file" and ".gguf" in m["path"]:
+                        if "-of-0" in m["path"] and "00001" not in m["path"]:
                             continue
-                        searchedmodels.append(m["rfilename"])
+                        searchedmodels.append(m["path"])
+                        searchedsizes.append(m["size"])
                 searchbox2.configure(values=searchedmodels)
                 if len(searchedmodels)>0:
                     quants = ["q4k","q4_k","q4", "q3", "q5", "q6", "q8"] #autopick priority
@@ -4220,19 +4234,23 @@ def show_gui():
                         if found_good:
                             break
                     modelsearch2_var.set(chosen_model)
+                    update_search_quant_file_size(1,1,1)
                 else:
                     modelsearch2_var.set("")
+                    fileinfotxt_var.set("")
             except Exception as e:
                 modelsearch1_var.set("")
                 modelsearch2_var.set("")
+                fileinfotxt_var.set("")
                 print(f"Error: {e}")
 
         def fetch_search_models():
             from tkinter import messagebox
-            nonlocal searchbox1, searchbox2, modelsearch1_var, modelsearch2_var
+            nonlocal searchbox1, searchbox2, modelsearch1_var, modelsearch2_var, fileinfotxt_var
             try:
                 modelsearch1_var.set("")
                 modelsearch2_var.set("")
+                fileinfotxt_var.set("")
                 searchbox1.configure(values=[])
                 searchbox2.configure(values=[])
                 searchedmodels = []
@@ -4259,6 +4277,7 @@ def show_gui():
             except Exception as e:
                 modelsearch1_var.set("")
                 modelsearch2_var.set("")
+                fileinfotxt_var.set("")
                 print(f"Error: {e}")
 
         ctk.CTkLabel(popup, text="Enter Search String:").pack(pady=(10, 0))
@@ -4275,7 +4294,8 @@ def show_gui():
         searchbox2 = ctk.CTkComboBox(popup, values=[], width=340, variable=modelsearch2_var, state="readonly")
         searchbox2.pack(pady=5)
         modelsearch1_var.trace("w", fetch_search_quants)
-
+        modelsearch2_var.trace("w", update_search_quant_file_size)
+        ctk.CTkLabel(popup, text="", textvariable=fileinfotxt_var, text_color="#ffff00").pack(pady=(10, 0))
         ctk.CTkButton(popup, text="Confirm Selection", command=confirm_search_model_choice).pack(pady=5)
 
         popup.transient(root)