| import os
|
| import sys
|
| import threading
|
| import time
|
| import logging
|
| from concurrent.futures import ThreadPoolExecutor, as_completed
|
| import tkinter as tk
|
| from tkinter import filedialog, messagebox, ttk, simpledialog
|
| from huggingface_hub import HfApi
|
|
|
|
|
| os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
|
|
|
| os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
|
|
|
| class TextHandler(logging.Handler):
|
| def __init__(self, text_widget):
|
| super().__init__()
|
| self.text_widget = text_widget
|
|
|
| def emit(self, record):
|
| msg = self.format(record)
|
| def append():
|
| self.text_widget.configure(state='normal')
|
| self.text_widget.insert(tk.END, msg + '\n')
|
| self.text_widget.configure(state='disabled')
|
| self.text_widget.see(tk.END)
|
| self.text_widget.after(0, append)
|
|
|
| class HFUploaderApp:
|
| def __init__(self, root):
|
| self.root = root
|
| self.root.title("HF 仓库管理 - 大模型批量上传版")
|
| self.root.geometry("950x700")
|
|
|
| self.token_var = tk.StringVar()
|
| self.repo_id_var = tk.StringVar(value="waveshaker/comfy_models")
|
| self.selected_tree_path = ""
|
| self.upload_queue = {}
|
|
|
|
|
| self.is_uploading = False
|
|
|
| self.setup_ui()
|
| self.load_token()
|
|
|
| def setup_ui(self):
|
|
|
| top_frame = tk.Frame(self.root, pady=10)
|
| top_frame.pack(fill=tk.X, padx=10)
|
| tk.Label(top_frame, text="HF Token:").pack(side=tk.LEFT)
|
| tk.Entry(top_frame, textvariable=self.token_var, show="*", width=50).pack(side=tk.LEFT, padx=5)
|
| tk.Button(top_frame, text="保存 Token", command=self.save_token).pack(side=tk.LEFT)
|
|
|
|
|
| repo_frame = tk.Frame(self.root)
|
| repo_frame.pack(fill=tk.X, padx=10, pady=5)
|
| tk.Label(repo_frame, text="仓库 ID:").pack(side=tk.LEFT)
|
| tk.Entry(repo_frame, textvariable=self.repo_id_var, width=30).pack(side=tk.LEFT, padx=5)
|
|
|
|
|
| paned = tk.PanedWindow(self.root, orient=tk.HORIZONTAL, sashwidth=5)
|
| paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)
|
|
|
|
|
| left_frame = tk.Frame(paned)
|
| paned.add(left_frame, width=550)
|
| tk.Label(left_frame, text="上传队列 (双击目标路径可修改):").pack(anchor=tk.W)
|
|
|
|
|
| cols = ("source", "remote", "status")
|
| self.queue_tree = ttk.Treeview(left_frame, columns=cols, show="headings", height=10)
|
| self.queue_tree.heading("source", text="文件源 (本地/远端链接)")
|
| self.queue_tree.heading("remote", text="远端目标目录")
|
| self.queue_tree.heading("status", text="状态")
|
| self.queue_tree.column("source", width=250)
|
| self.queue_tree.column("remote", width=150)
|
| self.queue_tree.column("status", width=80, anchor="center")
|
| self.queue_tree.tag_configure("done", foreground="gray")
|
|
|
| self.queue_tree.pack(fill=tk.BOTH, expand=True)
|
| self.queue_tree.bind("<Double-1>", self.on_edit_remote_path)
|
|
|
| queue_btns = tk.Frame(left_frame)
|
| queue_btns.pack(fill=tk.X, pady=5)
|
| tk.Button(queue_btns, text="➕ 添加本地文件", command=self.add_to_queue).pack(side=tk.LEFT)
|
| tk.Button(queue_btns, text="❌ 移除选中", command=self.remove_from_queue).pack(side=tk.LEFT, padx=5)
|
|
|
|
|
| self.remote_url_var = tk.StringVar()
|
| tk.Entry(queue_btns, textvariable=self.remote_url_var, width=30).pack(side=tk.LEFT, padx=5)
|
| tk.Button(queue_btns, text="🔗 加入队列", command=self.add_remote_to_queue, ).pack(side=tk.LEFT)
|
|
|
| right_frame = tk.Frame(paned)
|
| paned.add(right_frame, width=350)
|
| tk.Label(right_frame, text="远端仓库目录树:").pack(anchor=tk.W)
|
|
|
| self.tree_view = ttk.Treeview(right_frame, show="tree", height=10)
|
| self.tree_view.pack(fill=tk.BOTH, expand=True)
|
| self.tree_view.bind("<<TreeviewSelect>>", self.on_tree_select)
|
|
|
| tree_btns = tk.Frame(right_frame)
|
| tree_btns.pack(fill=tk.X, pady=5)
|
| tk.Button(tree_btns, text="🔄 拉取远端目录", command=self.fetch_remote_tree).pack(side=tk.LEFT)
|
| tk.Button(tree_btns, text="📁 新建文件夹", command=self.create_remote_folder).pack(side=tk.LEFT, padx=5)
|
| tk.Button(tree_btns, text="⬇️ 设为目标", command=self.set_target_from_tree).pack(side=tk.LEFT)
|
|
|
|
|
| action_frame = tk.Frame(self.root, pady=10)
|
| action_frame.pack(fill=tk.X, padx=10)
|
|
|
| self.upload_btn = tk.Button(action_frame, text="🚀 开始并发上传", command=self.start_upload, font=("Arial", 10, "bold"))
|
| self.upload_btn.pack(side=tk.LEFT, ipadx=20, ipady=5)
|
|
|
|
|
| self.stop_btn = tk.Button(action_frame, text="🛑 停止上传", command=self.stop_upload, font=("Arial", 10, "bold"), state=tk.DISABLED)
|
| self.stop_btn.pack(side=tk.LEFT, padx=10, ipadx=20, ipady=5)
|
|
|
|
|
| self.total_progress_var = tk.DoubleVar()
|
| self.progress_bar = ttk.Progressbar(action_frame, orient="horizontal", length=200, mode="determinate", variable=self.total_progress_var)
|
| self.progress_bar.pack(side=tk.LEFT, padx=15, fill=tk.X, expand=True)
|
|
|
|
|
| self.progress_text = tk.Label(action_frame, text="等待上传...", fg="#2196F3", font=("Arial", 10, "bold"))
|
| self.progress_text.pack(side=tk.LEFT, padx=5)
|
|
|
|
|
| log_frame = tk.Frame(self.root, padx=10, pady=5)
|
| log_frame.pack(fill=tk.BOTH, expand=True)
|
| tk.Label(log_frame, text="操作日志:").pack(anchor=tk.W)
|
|
|
| self.log_text = tk.Text(log_frame, height=12, state='disabled', bg="#f0f0f0")
|
| self.log_text.pack(fill=tk.BOTH, expand=True)
|
|
|
| self.logger = logging.getLogger("HF_Uploader")
|
| self.logger.setLevel(logging.INFO)
|
| handler = TextHandler(self.log_text)
|
| handler.setFormatter(logging.Formatter('[%(asctime)s] %(message)s', '%H:%M:%S'))
|
| self.logger.addHandler(handler)
|
|
|
| def load_token(self):
|
| if os.path.exists("hf_token.txt"):
|
| with open("hf_token.txt", "r") as f:
|
| self.token_var.set(f.read().strip())
|
|
|
| def save_token(self):
|
| with open("hf_token.txt", "w") as f:
|
| f.write(self.token_var.get().strip())
|
| messagebox.showinfo("成功", "Token 已保存到 hf_token.txt")
|
|
|
| def add_to_queue(self):
|
| files = filedialog.askopenfilenames(
|
| title="选择要上传的模型文件",
|
| filetypes=[("SafeTensors Model", "*.safetensors"), ("All Files", "*.*")]
|
| )
|
| if files:
|
| for file_path in files:
|
| if file_path not in self.upload_queue:
|
| self.upload_queue[file_path] = ""
|
| self.refresh_queue()
|
|
|
| def remove_from_queue(self):
|
| selected = self.queue_tree.selection()
|
| for item in selected:
|
| local_path = self.queue_tree.item(item, "values")[0]
|
| if local_path in self.upload_queue:
|
| del self.upload_queue[local_path]
|
| self.refresh_queue()
|
|
|
| def refresh_queue(self):
|
|
|
| existing_status = {}
|
| for item in self.queue_tree.get_children():
|
| values = self.queue_tree.item(item, "values")
|
| if len(values) >= 3 and values[2] == "已完成":
|
| existing_status[values[0]] = "已完成"
|
|
|
|
|
| for item in self.queue_tree.get_children():
|
| self.queue_tree.delete(item)
|
|
|
|
|
| for local_path, remote_target in self.upload_queue.items():
|
|
|
| status = existing_status.get(local_path, "待上传")
|
|
|
| if status == "已完成":
|
| self.queue_tree.insert("", tk.END, values=(local_path, remote_target, status), tags=("done",))
|
| else:
|
| self.queue_tree.insert("", tk.END, values=(local_path, remote_target, status))
|
|
|
|
|
| children = self.queue_tree.get_children()
|
| if children:
|
| last_item = children[-1]
|
| self.queue_tree.selection_set(last_item)
|
| self.queue_tree.focus(last_item)
|
| self.queue_tree.see(last_item)
|
|
|
| def on_edit_remote_path(self, event):
|
| item = self.queue_tree.selection()
|
| if not item: return
|
| item = item[0]
|
| current_local = self.queue_tree.item(item, "values")[0]
|
| current_remote = self.queue_tree.item(item, "values")[1]
|
|
|
| top = tk.Toplevel(self.root)
|
| top.title("修改目标路径")
|
| tk.Label(top, text="输入远端目标文件夹 (例如: moss/ 或留空):").pack(padx=10, pady=5)
|
| entry = tk.Entry(top, width=40)
|
| entry.insert(0, current_remote)
|
| entry.pack(padx=10, pady=5)
|
|
|
| def save():
|
| new_path = entry.get().strip().strip("/")
|
| self.upload_queue[current_local] = new_path
|
| self.refresh_queue()
|
| top.destroy()
|
|
|
| tk.Button(top, text="确定", command=save).pack(pady=10)
|
|
|
| def fetch_remote_tree(self):
|
| token = self.token_var.get().strip()
|
| repo_id = self.repo_id_var.get().strip()
|
| if not token or not repo_id: return
|
|
|
| self.logger.info("正在拉取远端目录结构...")
|
| self.tree_view.delete(*self.tree_view.get_children())
|
|
|
| def fetch():
|
| try:
|
| api = HfApi(token=token)
|
| repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model")
|
| tree_dict = {}
|
| for file_path in repo_files:
|
| parts = file_path.split('/')
|
| current = tree_dict
|
| for part in parts:
|
| if part not in current: current[part] = {}
|
| current = current[part]
|
|
|
| def insert_tree(parent, dict_data, current_path=""):
|
| for key, children in dict_data.items():
|
| path = f"{current_path}/{key}" if current_path else key
|
| node_text = f"📄 {key}" if not children else f"📁 {key}"
|
| node = self.tree_view.insert(parent, tk.END, text=node_text, values=(path,))
|
| if children: insert_tree(node, children, path)
|
|
|
| self.root.after(0, lambda: insert_tree("", tree_dict))
|
| self.logger.info(f"成功加载远端目录结构。")
|
| except Exception as e:
|
| self.root.after(0, lambda err=e: self.logger.error(f"拉取失败: {err}"))
|
| threading.Thread(target=fetch, daemon=True).start()
|
|
|
| def create_remote_folder(self):
|
| token = self.token_var.get().strip()
|
| repo_id = self.repo_id_var.get().strip()
|
| if not token or not repo_id: return
|
|
|
| parent_path = self.selected_tree_path if self.selected_tree_path else ""
|
| folder_name = simpledialog.askstring("新建文件夹", f"在 '{parent_path or '根目录'}' 下新建文件夹:")
|
| if not folder_name: return
|
| folder_name = folder_name.strip().strip("/")
|
|
|
| new_folder_path = f"{parent_path}/{folder_name}" if parent_path else folder_name
|
| gitkeep_path = f"{new_folder_path}/.gitkeep"
|
|
|
| self.logger.info(f"正在新建远端文件夹: {new_folder_path} ...")
|
| def create():
|
| try:
|
| api = HfApi(token=token)
|
| api.upload_file(path_or_fileobj=b"", path_in_repo=gitkeep_path, repo_id=repo_id, repo_type="model")
|
| self.root.after(0, lambda: self.logger.info(f"✅ 文件夹新建成功: {new_folder_path}"))
|
| self.root.after(0, self.fetch_remote_tree)
|
| except Exception as e:
|
| self.root.after(0, lambda err=e: self.logger.error(f"❌ 新建失败: {err}"))
|
| threading.Thread(target=create, daemon=True).start()
|
|
|
| def on_tree_select(self, event):
|
| selected = self.tree_view.selection()
|
| if selected:
|
| self.selected_tree_path = self.tree_view.item(selected[0], "values")[0]
|
| if "." in os.path.basename(self.selected_tree_path):
|
| self.selected_tree_path = os.path.dirname(self.selected_tree_path)
|
|
|
| def set_target_from_tree(self):
|
| if not self.selected_tree_path:
|
| messagebox.showwarning("提示", "请先在右侧树中选中一个文件夹")
|
| return
|
|
|
| queue_selected = self.queue_tree.selection()
|
| if not queue_selected:
|
| messagebox.showwarning("提示", "请先在左侧队列中选中任务")
|
| return
|
|
|
| target_path = self.selected_tree_path.strip("/")
|
| count = 0
|
| for item in queue_selected:
|
| values = self.queue_tree.item(item, "values")
|
|
|
| if len(values) >= 3 and values[2] == "已完成":
|
| continue
|
| local_path = values[0]
|
| self.upload_queue[local_path] = target_path
|
| count += 1
|
|
|
| self.refresh_queue()
|
| self.logger.info(f"✅ 已将 {count} 个文件的目标路径批量设置为: {target_path}")
|
|
|
|
|
| def start_upload(self):
|
| token = self.token_var.get().strip()
|
| repo_id = self.repo_id_var.get().strip()
|
|
|
|
|
| pending_tasks = []
|
| for tree_item in self.queue_tree.get_children():
|
| values = self.queue_tree.item(tree_item, "values")
|
| status = values[2] if len(values) >= 3 else "待上传"
|
| if status != "已完成":
|
| pending_tasks.append(tree_item)
|
|
|
| if not token or not repo_id or not pending_tasks:
|
| messagebox.showwarning("警告", "请检查 Token、仓库 ID，或者当前没有待上传的新文件")
|
| return
|
|
|
| self.is_uploading = True
|
| self.upload_btn.config(state=tk.DISABLED, text="上传中...")
|
| self.stop_btn.config(state=tk.NORMAL)
|
|
|
|
|
| total_files = len(pending_tasks)
|
| self.logger.info(f"🚀 开始并发上传 {total_files} 个待处理文件...")
|
|
|
| def upload_worker():
|
| api = HfApi(token=token)
|
|
|
| tasks = []
|
| for tree_item in pending_tasks:
|
| values = self.queue_tree.item(tree_item, "values")
|
| source = values[0]
|
| remote_dir = values[1]
|
|
|
|
|
| if source.startswith("http://") or source.startswith("https://"):
|
|
|
| filename = source.split("/")[-1]
|
| repo_path = f"{remote_dir}/{filename}" if remote_dir else filename
|
| tasks.append((tree_item, source, repo_path, True))
|
| else:
|
|
|
| filename = os.path.basename(source)
|
| repo_path = f"{remote_dir}/{filename}" if remote_dir else filename
|
| tasks.append((tree_item, source, repo_path, False))
|
|
|
|
|
| completed_files = 0
|
| progress_lock = threading.Lock()
|
|
|
| def update_ui_progress(completed, total):
|
| percent = (completed / total) * 100
|
| self.root.after(0, lambda: self.total_progress_var.set(percent))
|
| self.root.after(0, lambda: self.progress_text.config(text=f"{completed} / {total} 个文件 ({percent:.1f}%)"))
|
|
|
| def upload_one(tree_item, source, repo_path, is_remote):
|
| nonlocal completed_files
|
| if not self.is_uploading:
|
| return False, "用户已取消"
|
|
|
| try:
|
| if is_remote:
|
| import subprocess
|
| import os
|
|
|
| clean_path = source.replace("https://huggingface.co/", "")
|
| parts = clean_path.split("/")
|
| if len(parts) >= 5 and parts[2] in ["blob", "resolve"]:
|
| src_repo_id = f"{parts[0]}/{parts[1]}"
|
| src_file_path = "/".join(parts[4:])
|
|
|
|
|
| current_token = self.token_var.get().strip()
|
|
|
| cmd = [
|
| "hf", "cp",
|
| f"hf://{src_repo_id}/{src_file_path}",
|
| f"hf://{repo_id}/{repo_path}"
|
| ]
|
|
|
|
|
| env = os.environ.copy()
|
| env["HF_TOKEN"] = current_token
|
|
|
| result = subprocess.run(cmd, capture_output=True, text=True, env=env)
|
| if result.returncode != 0:
|
| raise RuntimeError(result.stderr.strip() or result.stdout.strip())
|
| else:
|
|
|
| api.upload_file(
|
| path_or_fileobj=source,
|
| path_in_repo=repo_path,
|
| repo_id=repo_id,
|
| repo_type="model"
|
| )
|
|
|
| with progress_lock:
|
| completed_files += 1
|
| update_ui_progress(completed_files, total_files)
|
|
|
| self.root.after(0, lambda: self.mark_item_done(tree_item))
|
| return True, repo_path
|
| except Exception as e:
|
| return False, f"{repo_path} ({e})"
|
|
|
| self.root.after(0, lambda: self.total_progress_var.set(0))
|
| self.root.after(0, lambda: self.progress_text.config(text=f"0 / {total_files} 个文件", fg="#2196F3"))
|
|
|
| with ThreadPoolExecutor(max_workers=4) as executor:
|
| future_to_task = {executor.submit(upload_one, t, s, r, is_r): r for t, s, r, is_r in tasks}
|
| for future in as_completed(future_to_task):
|
| if not self.is_uploading: break
|
| success, msg = future.result()
|
| if success:
|
| self.root.after(0, lambda m=msg: self.logger.info(f"✅ 成功: {m}"))
|
| elif msg != "用户已取消":
|
| self.root.after(0, lambda m=msg: self.logger.error(f"❌ 失败: {m}"))
|
|
|
| self.root.after(0, self.upload_finished)
|
|
|
| threading.Thread(target=upload_worker, daemon=True).start()
|
|
|
| def update_ui_progress(self, percent, completed, total):
|
| self.total_progress_var.set(percent)
|
| self.progress_text.config(text=f"{completed} / {total} 个文件 ({percent:.1f}%)")
|
|
|
| def mark_item_done(self, tree_item):
|
| """将表格中的某一行标记为已完成（灰色）"""
|
| values = list(self.queue_tree.item(tree_item, "values"))
|
| if len(values) >= 3:
|
| values[2] = "已完成"
|
| else:
|
| values.append("已完成")
|
| self.queue_tree.item(tree_item, values=values, tags=("done",))
|
|
|
| def stop_upload(self):
|
| """强行停止上传"""
|
| if messagebox.askyesno("确认停止", "确定要停止上传吗？\n（已传完的部分已保存在服务器，下次点上传会自动断点续传）"):
|
| self.is_uploading = False
|
| self.logger.warning("🛑 已中断后台请求，界面已重置。等待后台线程安全退出...")
|
| self.upload_btn.config(state=tk.NORMAL, text="🚀 开始并发上传")
|
| self.stop_btn.config(state=tk.DISABLED)
|
| self.progress_text.config(text="已完成 / 已停止", fg="#2E7D32")
|
|
|
| def upload_finished(self):
|
| self.is_uploading = False
|
| self.upload_btn.config(state=tk.NORMAL, text="🚀 开始并发上传")
|
| self.stop_btn.config(state=tk.DISABLED)
|
| self.progress_text.config(text="已完成 / 已停止", fg="#2E7D32")
|
| self.logger.info("🎉 任务处理完毕！")
|
| messagebox.showinfo("完成", "上传任务已结束（或已停止），请查看日志。")
|
|
|
| def add_remote_to_queue(self):
|
| url = self.remote_url_var.get().strip()
|
| if not url:
|
| messagebox.showwarning("提示", "请先粘贴 HF 文件链接")
|
| return
|
|
|
|
|
| if url not in self.upload_queue:
|
| self.upload_queue[url] = ""
|
| self.refresh_queue()
|
| self.remote_url_var.set("")
|
|
|
|
|
| children = self.queue_tree.get_children()
|
| if children:
|
| self.queue_tree.selection_set(children[-1])
|
|
|
| if __name__ == "__main__":
|
| root = tk.Tk()
|
| app = HFUploaderApp(root)
|
| root.mainloop() |
