|
|
| import json
|
| import os
|
| import sys
|
| import time
|
| import threading
|
| import traceback
|
| import tkinter as tk
|
| from tkinter import ttk, messagebox, filedialog
|
| from tkinter.scrolledtext import ScrolledText
|
| import urllib.request
|
| import urllib.error
|
| from pathlib import Path
|
| import subprocess
|
| import tempfile
|
| import pandas as pd
|
| from datasets import Dataset, Features, Image as HFImage, Value
|
| from huggingface_hub import HfApi, DatasetCard
|
| from PIL import Image
|
| import os
|
| for k in ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"]:
|
| os.environ.pop(k, None)
|
|
|
|
|
| SITE_CONFIG_FILE = "site_config.json"
|
| SPACE_ID = "waveshaker/loras_gallery"
|
| MODEL_REPO_ID = "waveshaker/loras"
|
| DATASET_REPO_ID = "waveshaker/dataset_collections"
|
| DATASET_BASE = "https://huggingface.co/datasets/waveshaker/dataset_collections/tree/main"
|
|
|
| OUTPUT_FILE = "models.json"
|
| TOKEN_FILE = "hf_token.txt"
|
| LOG_FILE = "models_manager.log"
|
| HTTP_TIMEOUT = 20
|
|
|
|
|
| TOP_ORDER = {"krea2": 1, "anima": 2, "flux1_dev": 3, "pony": 4, "sd15": 5}
|
| SUB_ORDER = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}
|
| IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
|
|
|
| def custom_sort_key(file_path):
|
| parts = file_path.split('/')
|
| top_dir = parts[0] if len(parts) > 0 else ""
|
| sub_dir = parts[1] if len(parts) > 1 else ""
|
| return (TOP_ORDER.get(top_dir, 999), SUB_ORDER.get(sub_dir, 999), file_path)
|
|
|
|
|
| try:
|
| from tkinterdnd2 import DND_FILES, TkinterDnD
|
| _HAS_DND = True
|
| except Exception:
|
| DND_FILES = None
|
| TkinterDnD = None
|
| _HAS_DND = False
|
|
|
| class ModelsManager(tk.Tk):
|
| def __init__(self):
|
| super().__init__()
|
| self.title("LoRA 仓库管理工具")
|
| self.geometry("1250x850")
|
|
|
|
|
| self.token_var = tk.StringVar(value=self._load_token())
|
|
|
|
|
| self.site_config_data = {"models": [], "categories": []}
|
|
|
|
|
| self.items = {}
|
| self.remote_files = []
|
| self.new_paths = set()
|
| self.stale_paths = set()
|
| self._search_job = None
|
| self._loading = False
|
| self.api = None
|
|
|
|
|
| if getattr(sys, 'frozen', False):
|
| current_dir = os.path.dirname(sys.executable)
|
| else:
|
| current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
| default_models_dir = os.path.join(current_dir, "models")
|
| if not os.path.exists(default_models_dir):
|
| os.makedirs(default_models_dir, exist_ok=True)
|
|
|
| self.local_folder_var = tk.StringVar(value=default_models_dir)
|
| self.remote_dir_var = tk.StringVar()
|
|
|
|
|
| self.token_var = tk.StringVar(value=self._load_token())
|
|
|
|
|
| self._build_global_ui()
|
| self._build_json_tab()
|
| self._build_model_tab()
|
| self._build_dataset_tab()
|
|
|
| self._refresh_api()
|
| self._setup_dnd()
|
| self.log_json("程序启动")
|
|
|
|
|
|
|
|
|
| def log_json(self, msg):
|
| self._log_to_tab(msg, getattr(self, 'log_json_text', None))
|
|
|
| def log_model(self, msg):
|
| self._log_to_tab(msg, getattr(self, 'log_model_text', None))
|
|
|
| def _log_to_tab(self, msg, target_text_widget):
|
| if not target_text_widget: return
|
| ts = time.strftime("%H:%M:%S")
|
| line = f"[{ts}] {msg}"
|
| try:
|
| with open(LOG_FILE, "a", encoding="utf-8") as f:
|
| f.write(line + "\n")
|
| except Exception: pass
|
|
|
| def _append():
|
| try:
|
| target_text_widget.configure(state="normal")
|
| target_text_widget.insert("end", line + "\n")
|
| target_text_widget.see("end")
|
| target_text_widget.configure(state="disabled")
|
| except Exception: pass
|
|
|
| if threading.current_thread() is threading.main_thread():
|
| _append()
|
| else:
|
| self.after(0, _append)
|
|
|
| def log_exc(self, prefix="", is_model=False):
|
| tb = traceback.format_exc()
|
| msg = f"{prefix}异常:\n{tb}"
|
| if is_model:
|
| self.log_model(msg)
|
| else:
|
| self.log_json(msg)
|
|
|
| def _load_token(self):
|
| env = os.environ.get("HF_TOKEN", "").strip()
|
| if env: return env
|
| try:
|
| with open(TOKEN_FILE, "r", encoding="utf-8") as f:
|
| return f.read().strip()
|
| except FileNotFoundError:
|
| return ""
|
|
|
| def _save_token(self, token):
|
| if token:
|
| with open(TOKEN_FILE, "w", encoding="utf-8") as f: f.write(token)
|
| else:
|
| try: os.remove(TOKEN_FILE)
|
| except FileNotFoundError: pass
|
|
|
| def _refresh_api(self):
|
| token = self.token_var.get().strip()
|
| self.api = HfApi(token=token) if token else HfApi()
|
|
|
| def _ensure_token(self):
|
| token = self.token_var.get().strip()
|
| if not token:
|
| messagebox.showwarning("需要 Token", "请先在上方填写 HF Token 再操作。")
|
| self.log_json("操作中止：未填写 Token")
|
| return None
|
| self._refresh_api()
|
| return token
|
|
|
| def _build_global_ui(self):
|
|
|
| token_row = ttk.Frame(self)
|
| token_row.pack(fill="x", padx=8, pady=(8, 2))
|
| ttk.Label(token_row, text="HF Token:").pack(side="left")
|
| ttk.Entry(token_row, textvariable=self.token_var, width=60, show="*").pack(side="left", padx=6)
|
| ttk.Button(token_row, text="📥 保存 Token", command=self._on_save_token).pack(side="left")
|
| ttk.Button(token_row, text="❌ 清除", command=self._on_clear_token).pack(side="left", padx=4)
|
|
|
|
|
| style = ttk.Style()
|
|
|
| if sys.platform.startswith("win"):
|
| style.theme_use('vista')
|
|
|
|
|
| style.layout("TNotebook.Tab", [
|
| ('Notebook.tab', {'sticky': 'nswe', 'children': [
|
| ('Notebook.padding', {'side': 'top', 'sticky': 'nswe', 'children': [
|
| ('Notebook.label', {'side': 'top', 'sticky': ''})
|
| ]})
|
| ]})
|
| ])
|
|
|
|
|
| style.configure("TNotebook.Tab",
|
| padding=[12, 6],
|
| font=("Microsoft YaHei", 9),
|
| borderwidth=0,
|
| background="#f0f0f0",
|
| foreground="#333333"
|
| )
|
|
|
|
|
| style.map("TNotebook.Tab",
|
| padding=[("selected", [12, 6]), ("active", [12, 6]), ("!selected", [12, 6])],
|
| borderwidth=[("selected", 0), ("active", 0), ("!selected", 0)],
|
|
|
| focuscolor=[("selected", "none"), ("active", "none"), ("!selected", "none")],
|
|
|
| background=[("selected", "#e1f5fe"), ("active", "#f5f5f5")],
|
| foreground=[("selected", "#0277bd")]
|
| )
|
|
|
|
|
| self.notebook = ttk.Notebook(self)
|
| self.notebook.pack(fill="both", expand=True, padx=8, pady=4)
|
|
|
| self.tab_json = ttk.Frame(self.notebook)
|
| self.tab_model = ttk.Frame(self.notebook)
|
| self.tab_dataset = ttk.Frame(self.notebook)
|
|
|
|
|
| self.notebook.add(self.tab_json, text=" 🚀 Space - JSON 管理 ")
|
| self.notebook.add(self.tab_model, text=" 🧊 Models - 模型上传 ")
|
| self.notebook.add(self.tab_dataset, text=" 📚 Dataset - 数据集上传 ")
|
|
|
| def _on_save_token(self):
|
| self._save_token(self.token_var.get().strip())
|
| self._refresh_api()
|
| self.log_json("Token 已保存")
|
| messagebox.showinfo("完成", "Token 已保存到本地")
|
|
|
| def _on_clear_token(self):
|
| self.token_var.set("")
|
| self._save_token("")
|
| self._refresh_api()
|
| self.log_json("Token 已清除")
|
| messagebox.showinfo("完成", "Token 已清除")
|
|
|
| def open_log_file(self):
|
| try: os.startfile(LOG_FILE)
|
| except Exception:
|
| try:
|
| import subprocess
|
| subprocess.Popen(["notepad", LOG_FILE])
|
| except Exception as e:
|
| messagebox.showerror("错误", f"无法打开日志: {e}")
|
|
|
| def _setup_dnd(self):
|
| if not _HAS_DND: return
|
| try:
|
| TkinterDnD._require(self)
|
| self.tree.drop_target_register(DND_FILES)
|
| self.tree.dnd_bind("<<Drop>>", self._on_drop_json)
|
| except Exception as e:
|
| self.log_json(f"DnD 初始化失败: {e}")
|
|
|
|
|
|
|
|
|
|
|
| def _build_json_tab(self):
|
| parent = self.tab_json
|
|
|
| top = ttk.Frame(parent)
|
| top.pack(fill="x", padx=8, pady=4)
|
| self.btn_load_remote = ttk.Button(top, text="🔄 拉取 models.json", command=self.load_remote)
|
| self.btn_load_remote.pack(side="left")
|
| ttk.Button(top, text="📂 加载本地 models.json", command=self.load_local).pack(side="left", padx=6)
|
| ttk.Button(top, text="💾 保存 models.json 到本地", command=self.save_local).pack(side="left")
|
| ttk.Button(top, text="🢁 上传 models.json", command=self.upload_models_json).pack(side="left", padx=6)
|
|
|
| ttk.Button(top, text="🌐 打开 Dataset Repo", compound="left", command=self.open_dataset_repo).pack(side="right", padx=6)
|
| ttk.Button(top, text="⚙️ 编辑 site_config", command=self.open_site_config_editor).pack(side="right", padx=12)
|
|
|
|
|
| filt = ttk.Frame(parent)
|
| filt.pack(fill="x", padx=8)
|
| ttk.Label(filt, text="🔍 搜索:").pack(side="left")
|
| self.search_var = tk.StringVar()
|
| self.search_var.trace_add("write", lambda *_: self._debounced_refresh())
|
| ttk.Entry(filt, textvariable=self.search_var, width=30).pack(side="left", padx=4)
|
|
|
| self.only_new = tk.BooleanVar(value=False)
|
| ttk.Checkbutton(filt, text="只看新增", variable=self.only_new, command=self.refresh_tree).pack(side="left", padx=6)
|
| self.only_empty_trigger = tk.BooleanVar(value=False)
|
| ttk.Checkbutton(filt, text="只看 trigger 为空", variable=self.only_empty_trigger, command=self.refresh_tree).pack(side="left")
|
|
|
| body = ttk.Panedwindow(parent, orient="vertical")
|
| body.pack(fill="both", expand=True, padx=8, pady=6)
|
|
|
| tree_frame = ttk.Frame(body)
|
| cols = ("title", "trigger", "dataset")
|
| self.tree = ttk.Treeview(tree_frame, columns=cols, show="tree headings", height=20)
|
| self.tree.heading("#0", text="路径")
|
| self.tree.heading("title", text="title")
|
| self.tree.heading("trigger", text="trigger")
|
| self.tree.heading("dataset", text="dataset")
|
| self.tree.column("#0", width=380, stretch=True)
|
| self.tree.column("title", width=180)
|
| self.tree.column("trigger", width=260)
|
| self.tree.column("dataset", width=240)
|
| vsb = ttk.Scrollbar(tree_frame, orient="vertical", command=self.tree.yview)
|
| self.tree.configure(yscrollcommand=vsb.set)
|
| self.tree.pack(side="left", fill="both", expand=True)
|
| vsb.pack(side="left", fill="y")
|
| self.tree.tag_configure("new", background="#d4f5d4")
|
| self.tree.tag_configure("empty_trigger", background="#fff3c4")
|
| self.tree.tag_configure("stale", foreground="#999999")
|
| self.tree.bind("<Double-1>", self.on_double_click)
|
|
|
| log_frame = ttk.LabelFrame(body, text="JSON 操作日志")
|
| self.log_json_text = ScrolledText(log_frame, height=10, state="disabled", wrap="word")
|
| self.log_json_text.pack(fill="both", expand=True)
|
|
|
| body.add(tree_frame, weight=3)
|
| body.add(log_frame, weight=1)
|
|
|
| def open_dataset_repo(self):
|
| import webbrowser
|
| webbrowser.open(f"https://huggingface.co/datasets/{DATASET_REPO_ID}/tree/main")
|
|
|
| def _debounced_refresh(self):
|
| if self._search_job: self.after_cancel(self._search_job)
|
| self._search_job = self.after(250, self.refresh_tree)
|
|
|
| def load_local(self):
|
| self.log_json(f"加载本地 {OUTPUT_FILE}")
|
| try:
|
| with open(OUTPUT_FILE, "r", encoding="utf-8") as f:
|
| data = json.load(f)
|
| except FileNotFoundError:
|
| messagebox.showwarning("提示", f"本地没有 {OUTPUT_FILE}")
|
| return
|
| except Exception as e:
|
| self.log_json(f"读取失败: {e}")
|
| return
|
|
|
| self.items = data.get("items", {})
|
| if not self.items and "titles" in data:
|
| self.items = {k: {"title": v, "trigger": "", "dataset": ""} for k, v in data["titles"].items()}
|
| self.new_paths = set()
|
| self.stale_paths = set()
|
| self.refresh_tree()
|
| self.log_json(f"本地加载完成，{len(self.items)} 条")
|
|
|
| def load_remote(self):
|
| if self._loading: return
|
| self._loading = True
|
| self.btn_load_remote.configure(state="disabled")
|
| token = self.token_var.get().strip()
|
| self.log_json(f"开始拉取远端…")
|
| threading.Thread(target=self._load_remote_worker, args=(token,), daemon=True).start()
|
|
|
| def _load_remote_worker(self, token):
|
| try:
|
| api = HfApi(token=token) if token else HfApi()
|
| t0 = time.time()
|
| all_files = list(api.list_repo_files(repo_id=MODEL_REPO_ID, repo_type="model"))
|
| remote_files = [f for f in all_files if f.endswith(".safetensors")]
|
| self.log_json(f"步骤1: 远端文件 {len(all_files)}，safetensors {len(remote_files)}，{time.time()-t0:.2f}s")
|
|
|
| space_subdomain = SPACE_ID.replace("/", "-").replace("_", "-").lower()
|
| url = f"https://{space_subdomain}.static.hf.space/models.json"
|
| remote_data = None
|
| try:
|
| req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
|
| with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as response:
|
| remote_data = json.loads(response.read().decode())
|
| self.log_json(f"步骤2: JSON 获取成功")
|
| except Exception as e:
|
| self.log_json(f"步骤2 失败: {type(e).__name__}: {e}")
|
|
|
| existing_items = {}
|
| if remote_data:
|
| existing_items = remote_data.get("items", {})
|
| if not existing_items and "titles" in remote_data:
|
| existing_items = {k: {"title": v, "trigger": "", "dataset": ""} for k, v in remote_data["titles"].items()}
|
|
|
| merged, new_paths = {}, set()
|
| remote_set = set(remote_files)
|
| for path in remote_files:
|
| if path in existing_items:
|
| merged[path] = existing_items[path]
|
| else:
|
| merged[path] = {"title": "", "trigger": "", "dataset": ""}
|
| new_paths.add(path)
|
| stale_paths = set(existing_items.keys()) - remote_set
|
| for p in stale_paths:
|
| merged[p] = existing_items[p]
|
|
|
| items = {k: merged[k] for k in sorted(merged.keys(), key=custom_sort_key)}
|
| self.log_json(f"步骤3: 合并 {len(items)} 条，新增 {len(new_paths)}，失效 {len(stale_paths)}")
|
|
|
| self.after(0, lambda: self._apply_remote(items, remote_files, new_paths, stale_paths))
|
| except Exception:
|
| self.log_exc("拉取失败")
|
| self.after(0, lambda: messagebox.showerror("错误", "拉取失败，详见日志"))
|
| finally:
|
| self.after(0, self._finish_loading)
|
|
|
| def _apply_remote(self, items, remote_files, new_paths, stale_paths):
|
| self.items = items
|
| self.remote_files = remote_files
|
| self.new_paths = new_paths
|
| self.stale_paths = stale_paths
|
| self.refresh_tree()
|
|
|
| def _finish_loading(self):
|
| self._loading = False
|
| self.btn_load_remote.configure(state="normal")
|
|
|
| def refresh_tree(self):
|
| self.tree.delete(*self.tree.get_children())
|
| kw = self.search_var.get().strip().lower()
|
| only_new = self.only_new.get()
|
| only_empty = self.only_empty_trigger.get()
|
|
|
| for path, meta in self.items.items():
|
| title = meta.get("title", "")
|
| trigger = meta.get("trigger", "")
|
| if kw and kw not in path.lower() and kw not in str(title).lower(): continue
|
| if only_new and path not in self.new_paths: continue
|
| if only_empty and trigger not in ("", [], None): continue
|
|
|
| tags = []
|
| if path in self.new_paths: tags.append("new")
|
| elif path in self.stale_paths: tags.append("stale")
|
| if trigger in ("", [], None): tags.append("empty_trigger")
|
|
|
| trig_display = trigger if isinstance(trigger, str) else " | ".join(map(str, trigger))
|
| self.tree.insert("", "end", iid=path, text=path, values=(title, trig_display, meta.get("dataset", "")), tags=tags)
|
|
|
| def on_double_click(self, event):
|
| sel = self.tree.focus()
|
| if sel and sel in self.items:
|
| self.open_manager(sel)
|
|
|
| def open_manager(self, path):
|
| meta = self.items[path]
|
| win = tk.Toplevel(self)
|
| win.title(f"编辑: {path}")
|
| win.geometry("720x560")
|
| win.transient(self); win.grab_set()
|
|
|
| ttk.Label(win, text="路径").pack(anchor="w", padx=10, pady=(10, 0))
|
| p_entry = ttk.Entry(win, width=90)
|
| p_entry.insert(0, path); p_entry.configure(state="readonly"); p_entry.pack(padx=10, fill="x")
|
|
|
| ttk.Label(win, text="title").pack(anchor="w", padx=10, pady=(8, 0))
|
| title_var = tk.StringVar(value=meta.get("title", ""))
|
| ttk.Entry(win, textvariable=title_var, width=90).pack(padx=10, fill="x")
|
|
|
| ttk.Label(win, text="trigger").pack(anchor="w", padx=10, pady=(8, 0))
|
| trig_text = tk.Text(win, height=8, wrap="word"); trig_text.pack(padx=10, fill="both")
|
| trig = meta.get("trigger", "")
|
| trig_text.insert("1.0", "\n".join(map(str, trig)) if isinstance(trig, list) else trig)
|
|
|
| ttk.Label(win, text="dataset").pack(anchor="w", padx=10, pady=(8, 0))
|
| ds_var = tk.StringVar(value=meta.get("dataset", ""))
|
| ttk.Entry(win, textvariable=ds_var, width=90).pack(padx=10, fill="x")
|
|
|
| def auto_dataset():
|
| parts = path.split("/")
|
| if len(parts) >= 3:
|
| top, sub = parts[0], parts[1]
|
| base = Path(parts[2]).stem
|
| for noise in ("_onetrainer", "_Flux_512", "_Flux", "_512"):
|
| base = base.replace(noise, "")
|
| ds_var.set(f"{DATASET_BASE}/{top}/{sub}/{base}")
|
|
|
| ttk.Button(win, text="自动生成 dataset", command=auto_dataset).pack(anchor="w", padx=10, pady=4)
|
|
|
| def save_and_close():
|
| title = title_var.get().strip()
|
| raw = trig_text.get("1.0", "end").strip()
|
| trigger_val = ([l.strip() for l in raw.splitlines() if l.strip()] if "\n" in raw else raw)
|
| self.items[path] = {"title": title, "trigger": trigger_val, "dataset": ds_var.get().strip()}
|
| self.refresh_tree()
|
| win.destroy()
|
|
|
| btn = ttk.Frame(win); btn.pack(fill="x", pady=10, padx=10)
|
| ttk.Button(btn, text="保存", command=save_and_close).pack(side="right")
|
| ttk.Button(btn, text="取消", command=win.destroy).pack(side="right", padx=6)
|
|
|
| def save_local(self):
|
| data_items = {k: v for k, v in self.items.items() if k not in self.stale_paths}
|
| sorted_items = {k: data_items[k] for k in sorted(data_items.keys(), key=custom_sort_key)}
|
| with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
|
| json.dump({"items": sorted_items}, f, indent=2, ensure_ascii=False)
|
| self.log_json(f"已保存本地 {OUTPUT_FILE}，{len(sorted_items)} 条")
|
| messagebox.showinfo("完成", f"已保存到 {OUTPUT_FILE}")
|
|
|
| def upload_models_json(self):
|
| if not self._ensure_token(): return
|
| if not messagebox.askyesno("确认", f"上传 models.json 到 {SPACE_ID}？"): return
|
|
|
| data_items = {k: v for k, v in self.items.items() if k not in self.stale_paths}
|
| sorted_items = {k: data_items[k] for k in sorted(data_items.keys(), key=custom_sort_key)}
|
|
|
| tmp = "models.upload.json"
|
| with open(tmp, "w", encoding="utf-8") as f:
|
| json.dump({"items": sorted_items}, f, indent=2, ensure_ascii=False)
|
|
|
| self.log_json(f"开始上传 JSON ({len(sorted_items)} 条)")
|
| api = self.api
|
|
|
| def work():
|
| try:
|
| api.upload_file(path_or_fileobj=tmp, path_in_repo="models.json", repo_id=SPACE_ID, repo_type="space")
|
| self.log_json("models.json 上传成功")
|
| self.after(0, lambda: messagebox.showinfo("完成", "JSON 上传成功"))
|
| except Exception:
|
| self.log_exc("JSON 上传失败")
|
| self.after(0, lambda: messagebox.showerror("失败", "详见日志"))
|
| finally:
|
| try: os.remove(tmp)
|
| except OSError: pass
|
|
|
| threading.Thread(target=work, daemon=True).start()
|
|
|
|
|
|
|
|
|
| def open_site_config_editor(self):
|
| if not self._ensure_token(): return
|
|
|
| self.log_json("正在从 Space 拉取 site_config.json...")
|
|
|
| def fetch_worker():
|
| try:
|
|
|
| file_path = self.api.hf_hub_download(
|
| repo_id=SPACE_ID,
|
| filename=SITE_CONFIG_FILE,
|
| repo_type="space"
|
| )
|
| with open(file_path, "r", encoding="utf-8") as f:
|
| config_data = json.load(f)
|
|
|
| self.after(0, lambda: self._show_config_window(config_data))
|
| except Exception as e:
|
| self.log_json(f"拉取 site_config 失败: {e}")
|
|
|
| self.after(0, lambda: self._show_config_window({"models": [], "categories": []}))
|
|
|
| threading.Thread(target=fetch_worker, daemon=True).start()
|
|
|
| def _show_config_window(self, config_data):
|
| win = tk.Toplevel(self)
|
| win.title("编辑 site_config.json (保存后自动上传)")
|
| win.geometry("600x550")
|
| win.transient(self); win.grab_set()
|
|
|
| ttk.Label(win, text="💡 提示：修改顺序即可改变前端左侧导航栏的展示顺序。", foreground="gray").pack(pady=5)
|
|
|
|
|
| frame_models = ttk.LabelFrame(win, text="Base Models (按顺序)")
|
| frame_models.pack(fill="both", expand=True, padx=10, pady=5)
|
|
|
| text_models = tk.Text(frame_models, height=10, wrap="none")
|
| text_models.pack(fill="both", expand=True, padx=5, pady=5)
|
|
|
| text_models.insert("1.0", json.dumps(config_data.get("models", []), indent=2, ensure_ascii=False))
|
|
|
| frame_cats = ttk.LabelFrame(win, text="Categories (按顺序)")
|
| frame_cats.pack(fill="both", expand=True, padx=10, pady=5)
|
|
|
| text_cats = tk.Text(frame_cats, height=10, wrap="none")
|
| text_cats.pack(fill="both", expand=True, padx=5, pady=5)
|
| text_cats.insert("1.0", json.dumps(config_data.get("categories", []), indent=2, ensure_ascii=False))
|
|
|
| def save_and_upload():
|
| try:
|
|
|
| new_models = json.loads(text_models.get("1.0", "end").strip())
|
| new_cats = json.loads(text_cats.get("1.0", "end").strip())
|
|
|
| new_config = {
|
| "models": new_models,
|
| "categories": new_cats
|
| }
|
|
|
|
|
| with open(SITE_CONFIG_FILE, "w", encoding="utf-8") as f:
|
| json.dump(new_config, f, indent=2, ensure_ascii=False)
|
|
|
| self.log_json("正在上传 site_config.json 到 Space...")
|
| self.api.upload_file(
|
| path_or_fileobj=SITE_CONFIG_FILE,
|
| path_in_repo=SITE_CONFIG_FILE,
|
| repo_id=SPACE_ID,
|
| repo_type="space"
|
| )
|
| self.log_json("🎉 site_config.json 上传成功！")
|
| messagebox.showinfo("成功", "配置已更新并上传到 Space！")
|
| win.destroy()
|
| except json.JSONDecodeError as e:
|
| messagebox.showerror("JSON 格式错误", f"请检查括号、逗号是否漏了。\n报错信息: {e}")
|
| except Exception as e:
|
| self.log_json(f"上传失败: {e}")
|
| messagebox.showerror("错误", f"上传失败: {e}")
|
|
|
| btn_frame = ttk.Frame(win)
|
| btn_frame.pack(fill="x", pady=10, padx=10)
|
| ttk.Button(btn_frame, text="保存并上传", command=save_and_upload).pack(side="right")
|
| ttk.Button(btn_frame, text="取消", command=win.destroy).pack(side="right", padx=6)
|
|
|
| def _on_drop_json(self, event):
|
| raw = event.data.strip("{}")
|
| p = Path(raw)
|
| if p.is_dir():
|
| self.notebook.select(self.tab_model)
|
| self.local_folder_var.set(str(p))
|
|
|
|
|
|
|
|
|
|
|
| def _build_model_tab(self):
|
| parent = self.tab_model
|
|
|
|
|
| frame = ttk.Frame(parent)
|
| frame.pack(fill="x", padx=10, pady=10)
|
|
|
|
|
| ttk.Label(frame, text="本地文件夹:").grid(row=0, column=0, sticky="w", pady=4)
|
| ttk.Entry(frame, textvariable=self.local_folder_var, width=70).grid(row=0, column=1, padx=6, sticky="we")
|
| ttk.Button(frame, text="📁 浏览...", width=18, command=self._browse_model_local).grid(row=0, column=2)
|
|
|
|
|
| ttk.Label(frame, text="上传到 HF 目录:").grid(row=1, column=0, sticky="nw", pady=4)
|
|
|
| tree_frame = ttk.Frame(frame)
|
| tree_frame.grid(row=1, column=1, padx=6, sticky="nsew")
|
|
|
|
|
| vsb = ttk.Scrollbar(tree_frame, orient="vertical")
|
| vsb.pack(side="right", fill="y")
|
|
|
|
|
| self.remote_tree = ttk.Treeview(tree_frame, height=6, show="tree",
|
| yscrollcommand=vsb.set, selectmode="browse")
|
| self.remote_tree.pack(side="left", fill="both", expand=True)
|
| vsb.config(command=self.remote_tree.yview)
|
|
|
|
|
| self.remote_tree.bind("<<TreeviewSelect>>", self._on_remote_tree_select)
|
| self.remote_tree.bind("<Double-1>", lambda e: self._start_model_upload())
|
|
|
|
|
| ttk.Button(frame, text="🔄 拉取 Models 目录", width=18, command=self._fetch_hf_dirs).grid(row=1, column=2, sticky="n")
|
|
|
|
|
| btn_row = ttk.Frame(frame)
|
| btn_row.grid(row=2, column=0, columnspan=3, pady=15, sticky="w")
|
|
|
|
|
| self.btn_upload = ttk.Button(btn_row, text="🢁 开始上传", command=self._start_model_upload)
|
| self.btn_upload.pack(side="left")
|
|
|
|
|
| ttk.Button(btn_row, text="📄 打开日志文件", command=self.open_log_file).pack(side="left", padx=10)
|
|
|
|
|
| self.status_model_var = tk.StringVar(value="就绪")
|
| ttk.Label(btn_row, textvariable=self.status_model_var).pack(side="left", padx=10)
|
|
|
|
|
| self.progress_var = tk.DoubleVar()
|
| self.progress_bar = ttk.Progressbar(btn_row, variable=self.progress_var, maximum=100, length=200, mode='determinate')
|
| self.progress_bar.pack(side="left", padx=10)
|
|
|
| frame.columnconfigure(1, weight=1)
|
|
|
| frame.rowconfigure(1, weight=1)
|
|
|
|
|
| log_frame = ttk.LabelFrame(parent, text="模型上传日志")
|
| log_frame.pack(fill="both", expand=True, padx=10, pady=6)
|
| self.log_model_text = ScrolledText(log_frame, height=10, state="disabled", wrap="word")
|
| self.log_model_text.pack(fill="both", expand=True)
|
|
|
|
|
| def _browse_model_local(self):
|
|
|
| if getattr(sys, 'frozen', False):
|
| current_dir = os.path.dirname(sys.executable)
|
| else:
|
| current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
| folder = filedialog.askdirectory(title="选择包含 .safetensors 的文件夹", initialdir=current_dir)
|
| if folder:
|
| self.local_folder_var.set(folder)
|
|
|
| def _fetch_hf_dirs(self):
|
| if not self._ensure_token(): return
|
| self.status_model_var.set("正在读取远端目录...")
|
| self.log_model("正在读取远端仓库目录列表...")
|
|
|
| def work():
|
| try:
|
| api = self.api
|
| all_files = list(api.list_repo_files(repo_id=MODEL_REPO_ID, repo_type="model"))
|
| dirs = set()
|
| for f in all_files:
|
| d = str(Path(f).parent)
|
| if d != ".":
|
| dirs.add(d)
|
|
|
| sorted_dirs = sorted(list(dirs))
|
| self.after(0, lambda: self._update_remote_tree(sorted_dirs))
|
| except Exception:
|
| self.log_exc("读取远端目录失败", is_model=True)
|
| finally:
|
| self.after(0, lambda: self.status_model_var.set("就绪"))
|
|
|
| threading.Thread(target=work, daemon=True).start()
|
|
|
| def _update_remote_tree(self, dirs):
|
|
|
| for item in self.remote_tree.get_children():
|
| self.remote_tree.delete(item)
|
|
|
| tree_nodes = {}
|
|
|
| for d in dirs:
|
|
|
| safe_d = d.replace("\\", "/")
|
| parts = safe_d.split("/")
|
|
|
| parent_id = ""
|
| current_path = ""
|
|
|
| for i, part in enumerate(parts):
|
| current_path = f"{current_path}/{part}" if current_path else part
|
|
|
|
|
| if current_path not in tree_nodes:
|
|
|
| is_leaf = (i == len(parts) - 1)
|
| display_text = f"📁 {part}" if is_leaf else part
|
|
|
| node_id = self.remote_tree.insert(
|
| parent_id, "end",
|
| text=display_text,
|
| values=(current_path,)
|
| )
|
| tree_nodes[current_path] = node_id
|
|
|
|
|
| parent_id = tree_nodes[current_path]
|
|
|
|
|
| for root_node in self.remote_tree.get_children(""):
|
| self.remote_tree.item(root_node, open=True)
|
|
|
| self.log_model(f"已加载 {len(dirs)} 个远端目录结构")
|
|
|
| def _on_remote_tree_select(self, event):
|
|
|
| selection = self.remote_tree.selection()
|
| if selection:
|
| item_id = selection[0]
|
|
|
| values = self.remote_tree.item(item_id, "values")
|
| if values:
|
| self.remote_dir_var.set(values[0])
|
|
|
|
|
| def _start_model_upload(self):
|
| local_dir = self.local_folder_var.get().strip()
|
|
|
| remote_dir = self.remote_dir_var.get().strip().strip("/")
|
|
|
| if not local_dir or not os.path.exists(local_dir):
|
| messagebox.showwarning("错误", "请选择有效的本地文件夹。")
|
| return
|
| if not remote_dir:
|
| if not messagebox.askyesno("确认", "远端目录为空，将上传到仓库根目录，继续？"):
|
| return
|
| if not self._ensure_token(): return
|
|
|
| self.status_model_var.set("正在扫描并比对文件...")
|
| self.log_model(f"开始扫描: {local_dir} -> {MODEL_REPO_ID}/{remote_dir}")
|
|
|
| def scan_worker():
|
| try:
|
| api = self.api
|
|
|
|
|
| local_files = [f for f in Path(local_dir).iterdir() if f.is_file()]
|
| for f in local_files:
|
| if f.suffix.lower() in [".png", ".webp"]:
|
| jpg_path = f.with_suffix(".jpg")
|
| try:
|
| self.log_model(f"清洗图片并转码: {f.name} -> {jpg_path.name}")
|
| with Image.open(f) as img:
|
| if img.mode in ("RGBA", "P"):
|
| img = img.convert("RGB")
|
| img.save(jpg_path, format="JPEG", quality=95)
|
| os.remove(f)
|
| except Exception as e:
|
| self.log_model(f"转码失败 {f.name}: {e}")
|
|
|
|
|
| local_files = [f for f in Path(local_dir).iterdir() if f.is_file()]
|
| if not local_files:
|
| self.after(0, lambda: messagebox.showwarning("提示", "本地文件夹是空的。"))
|
| return
|
|
|
|
|
| remote_all = set(api.list_repo_files(repo_id=MODEL_REPO_ID, repo_type="model"))
|
|
|
| to_upload = []
|
| for f in local_files:
|
| if f.suffix.lower() not in [".safetensors"] + list(IMG_EXTS):
|
| continue
|
|
|
| remote_path = f"{remote_dir}/{f.name}" if remote_dir else f.name
|
| if remote_path not in remote_all:
|
| to_upload.append((f, remote_path))
|
|
|
| if not to_upload:
|
| self.after(0, lambda: messagebox.showinfo("完成", "远端已存在所有文件，无需上传。"))
|
| self.after(0, lambda: self.status_model_var.set("无需上传"))
|
| self.log_model("扫描完毕，无新增文件。")
|
| return
|
|
|
| msg = f"本地清理完毕。找到 {len(to_upload)} 个新文件待上传。\n\n目标目录: {remote_dir}\n\n继续？"
|
| self.after(0, lambda: self._do_model_upload_confirm(msg, to_upload, api))
|
|
|
| except Exception:
|
| self.log_exc("扫描失败", is_model=True)
|
| self.after(0, lambda: self.status_model_var.set("扫描失败"))
|
|
|
| threading.Thread(target=scan_worker, daemon=True).start()
|
|
|
| def _do_model_upload_confirm(self, msg, to_upload, api):
|
| if not messagebox.askyesno("确认上传", msg):
|
| self.status_model_var.set("已取消")
|
| return
|
|
|
| def upload_worker():
|
| try:
|
| from huggingface_hub import CommitOperationAdd
|
|
|
|
|
| self.after(0, lambda: self.btn_upload.configure(state="disabled"))
|
| self.after(0, lambda: self.progress_var.set(0))
|
| self.after(0, lambda: self.status_model_var.set("正在打包并上传..."))
|
|
|
|
|
| operations = []
|
| for local_file, remote_path in to_upload:
|
|
|
| safe_remote_path = remote_path.replace("\\", "/")
|
| operations.append(
|
| CommitOperationAdd(
|
| path_or_fileobj=str(local_file),
|
| path_in_repo=safe_remote_path
|
| )
|
| )
|
|
|
|
|
| remote_dir = self.remote_dir_var.get().strip().replace("\\", "/").strip("/")
|
| commit_msg = f"Upload files to {remote_dir}" if remote_dir else "Upload files to root"
|
|
|
| self.log_model(f"开始打包 {len(operations)} 个文件，一次性 Commit...")
|
| self.log_model(f"Commit message: {commit_msg}")
|
|
|
|
|
| is_uploading = [True]
|
| def update_progress():
|
| if not is_uploading[0]: return
|
| current = self.progress_var.get()
|
|
|
| if current < 95:
|
| self.progress_var.set(current + 2)
|
| self.after(300, update_progress)
|
|
|
| self.after(100, update_progress)
|
|
|
|
|
| api.create_commit(
|
| repo_id=MODEL_REPO_ID,
|
| repo_type="model",
|
| operations=operations,
|
| commit_message=commit_msg
|
| )
|
|
|
|
|
| is_uploading[0] = False
|
| self.after(0, lambda: self.progress_var.set(100))
|
| self.after(0, lambda: self.status_model_var.set("上传完成 ✅"))
|
| self.after(0, lambda: self.btn_upload.configure(state="normal"))
|
|
|
| self.log_model("打包上传成功！")
|
| self.after(0, lambda: messagebox.showinfo("完成", "上传成功！"))
|
| self.after(0, self.load_remote)
|
|
|
| except Exception as e:
|
|
|
| is_uploading[0] = False
|
| self.after(0, lambda: self.progress_var.set(0))
|
| self.after(0, lambda: self.btn_upload.configure(state="normal"))
|
| self.after(0, lambda: self.status_model_var.set("上传失败 ❌"))
|
|
|
| self.log_exc("打包上传失败", is_model=True)
|
| self.after(0, lambda: messagebox.showerror("失败", "详见日志"))
|
|
|
| threading.Thread(target=upload_worker, daemon=True).start()
|
|
|
|
|
|
|
|
|
| def _build_dataset_tab(self):
|
| parent = self.tab_dataset
|
|
|
| frame = ttk.Frame(parent)
|
| frame.pack(fill="x", padx=10, pady=10)
|
|
|
|
|
| ttk.Label(frame, text="本地母文件夹:").grid(row=0, column=0, sticky="w", pady=4)
|
| self.ds_local_folder_var = tk.StringVar()
|
| ttk.Entry(frame, textvariable=self.ds_local_folder_var, width=70).grid(row=0, column=1, columnspan=2, padx=6, sticky="we")
|
| ttk.Button(frame, text="📁 浏览...", width=18, command=self._browse_ds_local).grid(row=0, column=3)
|
|
|
|
|
| ttk.Label(frame, text="Region 前缀:").grid(row=1, column=0, sticky="w", pady=8)
|
|
|
| prefix_frame = ttk.Frame(frame)
|
| prefix_frame.grid(row=1, column=1, columnspan=3, sticky="w", padx=6, pady=8)
|
|
|
|
|
| self.ds_region_var = tk.StringVar(value="KR")
|
| self.ds_custom_region_var = tk.StringVar()
|
|
|
|
|
| ttk.Button(prefix_frame, text="🔄 拉取 datasets 目录", command=self._fetch_dataset_dirs).pack(side="left", padx=(0, 10))
|
|
|
|
|
| self.ds_radio_container = ttk.Frame(prefix_frame)
|
| self.ds_radio_container.pack(side="left")
|
|
|
|
|
| self._render_dataset_radios(["CN", "KR", "JP", "Fictional", "WE"])
|
|
|
|
|
| btn_row = ttk.Frame(frame)
|
| btn_row.grid(row=2, column=0, columnspan=4, pady=15, sticky="w")
|
| self.btn_ds_upload = ttk.Button(btn_row, text="🢁 开始 5合1 处理并上传", command=self._start_dataset_pipeline)
|
| self.btn_ds_upload.pack(side="left")
|
|
|
| self.status_ds_var = tk.StringVar(value="就绪")
|
| ttk.Label(btn_row, textvariable=self.status_ds_var).pack(side="left", padx=10)
|
|
|
| frame.columnconfigure(1, weight=1)
|
| frame.columnconfigure(2, weight=1)
|
|
|
|
|
| log_frame = ttk.LabelFrame(parent, text="Dataset 自动化流水线日志")
|
| log_frame.pack(fill="both", expand=True, padx=10, pady=6)
|
| self.log_dataset_text = ScrolledText(log_frame, height=15, state="disabled", wrap="word")
|
| self.log_dataset_text.pack(fill="both", expand=True)
|
|
|
| def _fetch_dataset_dirs(self):
|
| if not self._ensure_token(): return
|
| self.log_dataset("正在读取 Dataset 仓库根目录...")
|
|
|
| def work():
|
| try:
|
| files = self.api.list_repo_files(repo_id=DATASET_REPO_ID, repo_type="dataset")
|
| dirs = set()
|
| for f in files:
|
| parts = f.split("/")
|
|
|
| if len(parts) > 1 and parts[0] not in ["tool", ".git"]:
|
| dirs.add(parts[0])
|
|
|
| sorted_dirs = sorted(list(dirs))
|
|
|
|
|
| def update_ui():
|
| self._render_dataset_radios(sorted_dirs)
|
| self.log_dataset(f"✅ 远端目录刷新成功！共发现 {len(sorted_dirs)} 个分类: {', '.join(sorted_dirs)}")
|
|
|
| self.after(0, update_ui)
|
|
|
| except Exception as e:
|
| self.log_exc("读取 Dataset 目录失败", is_model=True)
|
|
|
| self.after(0, lambda: self.status_ds_var.set("读取目录失败"))
|
|
|
| threading.Thread(target=work, daemon=True).start()
|
|
|
| def _render_dataset_radios(self, dirs):
|
|
|
| for widget in self.ds_radio_container.winfo_children():
|
| widget.destroy()
|
|
|
|
|
| current_val = self.ds_region_var.get()
|
| if current_val not in dirs and current_val != "Custom" and dirs:
|
| self.ds_region_var.set(dirs[0])
|
|
|
|
|
| for region in dirs:
|
| ttk.Radiobutton(self.ds_radio_container, text=region, variable=self.ds_region_var, value=region).pack(side="left", padx=(0, 10))
|
|
|
|
|
| ttk.Radiobutton(self.ds_radio_container, text="自定义:", variable=self.ds_region_var, value="Custom").pack(side="left", padx=(10, 2))
|
| self.ds_custom_entry = ttk.Entry(self.ds_radio_container, textvariable=self.ds_custom_region_var, width=15, state="disabled")
|
| self.ds_custom_entry.pack(side="left")
|
|
|
|
|
| def on_region_change(*args):
|
| if self.ds_region_var.get() == "Custom":
|
| self.ds_custom_entry.configure(state="normal")
|
| else:
|
| self.ds_custom_entry.configure(state="disabled")
|
|
|
|
|
| if hasattr(self, '_region_trace_id'):
|
| try:
|
| self.ds_region_var.trace_remove("write", self._region_trace_id)
|
| except Exception:
|
| pass
|
|
|
| self._region_trace_id = self.ds_region_var.trace_add("write", on_region_change)
|
|
|
|
|
| def log_dataset(self, msg):
|
| self._log_to_tab(msg, getattr(self, 'log_dataset_text', None))
|
|
|
| def _browse_ds_local(self):
|
|
|
| current_dir = os.path.dirname(os.path.abspath(__file__))
|
| folder = filedialog.askdirectory(title="选择包含子文件夹的数据集母目录", initialdir=current_dir)
|
| if folder:
|
| self.ds_local_folder_var.set(folder)
|
|
|
| def _start_dataset_pipeline(self):
|
| local_dir = self.ds_local_folder_var.get().strip()
|
|
|
|
|
| selected_region = self.ds_region_var.get()
|
| if selected_region == "Custom":
|
| region_prefix = self.ds_custom_region_var.get().strip().replace(" ", "_").replace("-", "_")
|
| if not region_prefix:
|
| messagebox.showwarning("错误", "您选择了自定义前缀，但输入框为空。")
|
| return
|
| else:
|
| region_prefix = selected_region
|
|
|
|
|
| if not local_dir or not os.path.exists(local_dir):
|
| messagebox.showwarning("错误", "请选择有效的本地母文件夹。")
|
| return
|
| if not self._ensure_token(): return
|
|
|
| if not messagebox.askyesno("确认", f"即将对以下目录执行 5合1 流水线：\n{local_dir}\n前缀：{region_prefix}\n\n继续？"):
|
| return
|
|
|
| self.btn_ds_upload.configure(state="disabled")
|
| self.status_ds_var.set("处理中...")
|
| self.log_dataset(f"================ 开始 5合1 流水线 ================")
|
| self.log_dataset(f"目标目录: {local_dir} | 前缀: {region_prefix}")
|
|
|
| def worker():
|
| api = self.api
|
| try:
|
|
|
| self._clean_txt_newlines(local_dir)
|
|
|
|
|
| self._remove_mask_metadata(local_dir)
|
|
|
|
|
| self._convert_png_to_jpg(local_dir)
|
|
|
|
|
| temp_readme_path, valid_folders = self._generate_parquet_and_update_yaml(api, local_dir, region_prefix)
|
|
|
|
|
| self._upload_files_and_readme(api, local_dir, region_prefix, temp_readme_path, valid_folders)
|
|
|
| self.log_dataset("🎉 5合1 全流程处理完毕！")
|
| self.after(0, lambda: messagebox.showinfo("完成", "Dataset 流水线执行成功！"))
|
| self.after(0, lambda: self.status_ds_var.set("上传完成"))
|
|
|
|
|
| self.after(0, self.load_remote)
|
|
|
| except Exception as e:
|
| self.log_exc("Dataset 流水线异常", is_model=True)
|
| self.after(0, lambda: self.status_ds_var.set("执行失败"))
|
| finally:
|
| self.after(0, lambda: self.btn_ds_upload.configure(state="normal"))
|
|
|
| threading.Thread(target=worker, daemon=True).start()
|
|
|
|
|
|
|
| def _clean_txt_newlines(self, current_dir):
|
| import glob
|
| self.log_dataset("🧹 [任务 1/5] 正在清理所有 TXT 文件中的换行符...")
|
| txt_files = glob.glob(os.path.join(current_dir, "**", "*.txt"), recursive=True)
|
| if not txt_files:
|
| self.log_dataset(" ✅ 没有找到需要处理的 TXT 文件。")
|
| return
|
|
|
| cleaned_count = 0
|
| for txt_path in txt_files:
|
| if os.path.dirname(txt_path) == current_dir:
|
| continue
|
| try:
|
| with open(txt_path, 'r', encoding='utf-8') as f:
|
| content = ''.join(line.strip() for line in f)
|
| with open(txt_path, 'w', encoding='utf-8') as f:
|
| f.write(content)
|
| cleaned_count += 1
|
| except Exception as e:
|
| self.log_dataset(f" ⚠️ 处理 {os.path.basename(txt_path)} 失败: {e}")
|
| self.log_dataset(f" 🎉 TXT 清理完成！共处理了 {cleaned_count} 个文件。")
|
|
|
| def _remove_mask_metadata(self, current_dir):
|
| import glob
|
| self.log_dataset("🛡️ [任务 2/5] 正在剥离 Mask 蒙版图的元数据...")
|
| png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
|
| mask_files = [f for f in png_files if "masklabel" in os.path.basename(f).lower()]
|
|
|
| if not mask_files:
|
| self.log_dataset(" ✅ 没有找到 masklabel.png 文件，跳过此步骤。")
|
| return
|
|
|
| success_count = 0
|
| for file_path in mask_files:
|
| try:
|
| with Image.open(file_path) as img:
|
| data = list(img.getdata())
|
| clean_img = Image.new(img.mode, img.size)
|
| clean_img.putdata(data)
|
| clean_img.save(file_path, "PNG")
|
| success_count += 1
|
| except Exception as e:
|
| self.log_dataset(f" ⚠️ 处理失败 [{os.path.basename(file_path)}]: {e}")
|
| self.log_dataset(f" 🎉 元数据清理完成！共处理了 {success_count} 张 Mask 图。")
|
|
|
| def _convert_png_to_jpg(self, current_dir):
|
| import glob
|
| self.log_dataset("🎬 [任务 3/5] 正在批量转换常规 PNG 为 JPG...")
|
| png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
|
| if not png_files:
|
| self.log_dataset(" ✅ 没有找到需要转换的 PNG 文件。")
|
| return
|
|
|
| converted_count = 0
|
| for png_path in png_files:
|
| if "masklabel" in png_path.lower():
|
| continue
|
| jpg_path = os.path.splitext(png_path)[0] + ".jpg"
|
| try:
|
| subprocess.run(
|
| ["ffmpeg", "-y", "-i", png_path, "-q:v", "2", jpg_path],
|
| stdout=subprocess.DEVNULL,
|
| stderr=subprocess.DEVNULL,
|
| check=True
|
| )
|
| os.remove(png_path)
|
| converted_count += 1
|
| self.log_dataset(f" 🔄 转换: {os.path.basename(png_path)} -> {os.path.basename(jpg_path)}")
|
| except subprocess.CalledProcessError:
|
| self.log_dataset(f" ⚠️ 转换失败: {png_path}")
|
| except FileNotFoundError:
|
| self.log_dataset(" ❌ 错误: 找不到 ffmpeg！请确保已安装并配置了环境变量。")
|
| return
|
| self.log_dataset(f" 🎉 转换完成！共转换 {converted_count} 张图片。")
|
|
|
| def _generate_parquet_and_update_yaml(self, api, current_dir, region_prefix):
|
| import glob, tempfile, pandas as pd
|
| from datasets import Dataset, Features, Image as HFImage, Value
|
| from huggingface_hub import DatasetCard
|
| from PIL import Image as PILImage
|
|
|
| self.log_dataset("📦 [任务 4/5] 正在本地打包 Parquet 并更新 YAML 配置...")
|
| local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
|
|
|
| self.log_dataset(f" ☁️ 正在读取云端 {DATASET_REPO_ID} 的现有配置...")
|
| try:
|
|
|
| remote_readme_path = api.hf_hub_download(repo_id=DATASET_REPO_ID, filename="README.md", repo_type="dataset")
|
| with open(remote_readme_path, "r", encoding="utf-8") as f:
|
| content = f.read()
|
| card = DatasetCard(content)
|
| self.log_dataset(" ✅ 成功读取云端 README。")
|
| except Exception as e:
|
| self.log_dataset(f" ⚠️ 找不到云端 README，将新建: {e}")
|
| card = DatasetCard("")
|
|
|
| yaml_data = card.data if card.data else {}
|
| if "configs" not in yaml_data: yaml_data["configs"] = []
|
|
|
| yaml_data["dataset_info"] = {
|
| "features": [
|
| {"name": "image", "dtype": "image"},
|
| {"name": "file_name", "dtype": "string"},
|
| {"name": "text", "dtype": "string"}
|
| ]
|
| }
|
|
|
| existing_config_names = [c.get("config_name") for c in yaml_data["configs"]]
|
| new_configs_added = False
|
| valid_folders = []
|
|
|
| for sub_folder_name in os.listdir(current_dir):
|
| sub_folder_path = os.path.join(current_dir, sub_folder_name)
|
| if os.path.isdir(sub_folder_path) and not sub_folder_name.startswith('.'):
|
| jpg_files = glob.glob(os.path.join(sub_folder_path, "*.jpg")) + glob.glob(os.path.join(sub_folder_path, "*.jpeg"))
|
| if not jpg_files:
|
| self.log_dataset(f" ⚠️ 跳过子文件夹 [{sub_folder_name}]: 未找到任何 .jpg 图片。")
|
| continue
|
|
|
| self.log_dataset(f" 📦 正在打包: {sub_folder_name}")
|
| data_rows = []
|
|
|
| for jpg_path in jpg_files:
|
| file_name = os.path.basename(jpg_path)
|
| base_name = os.path.splitext(file_name)[0]
|
| txt_path = os.path.join(sub_folder_path, f"{base_name}.txt")
|
|
|
| text_content = ""
|
| if os.path.exists(txt_path):
|
| with open(txt_path, "r", encoding="utf-8") as f:
|
| text_content = f.read().strip()
|
|
|
| try:
|
| with open(jpg_path, "rb") as f:
|
| image_bytes = {"bytes": f.read(), "path": None}
|
| except Exception as e:
|
| self.log_dataset(f" ⚠️ 跳过损坏图片 {file_name}: {e}")
|
| continue
|
|
|
| data_rows.append({
|
| "image": image_bytes,
|
| "file_name": file_name,
|
| "text": text_content
|
| })
|
|
|
| if data_rows:
|
| df = pd.DataFrame(data_rows)
|
| features = Features({
|
| "image": HFImage(),
|
| "file_name": Value("string"),
|
| "text": Value("string")
|
| })
|
| ds = Dataset.from_pandas(df, features=features)
|
| ds.to_parquet(os.path.join(sub_folder_path, "train.parquet"))
|
| valid_folders.append(sub_folder_name)
|
| self.log_dataset(f" ✅ 生成 train.parquet ({len(data_rows)} 条数据)")
|
|
|
| config_name = f"{local_folder_name}_{sub_folder_name}".lower()
|
| path_str = f"{region_prefix}/{local_folder_name}/{sub_folder_name}/train.parquet"
|
|
|
| if config_name not in existing_config_names:
|
| yaml_data["configs"].append({
|
| "config_name": config_name,
|
| "data_files": [{"split": "train", "path": path_str}]
|
| })
|
| new_configs_added = True
|
| self.log_dataset(f" ➕ 已添加配置: {config_name}")
|
|
|
| def get_priority(config):
|
| data_files = config.get("data_files", [])
|
| path = data_files[0].get("path", "") if data_files else ""
|
| priority_dict = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}
|
| path_prefix = path.split('/')[0] if path else ""
|
| return priority_dict.get(path_prefix, 99)
|
|
|
| yaml_data["configs"] = sorted(yaml_data["configs"], key=lambda c: (get_priority(c), c.get("config_name", "")))
|
| self.log_dataset(" 🔄 已根据 [CN > KR > JP > Fictional > WE] 的顺序重排 YAML。")
|
|
|
| if new_configs_added:
|
| card.data = yaml_data
|
| with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', suffix='.md', delete=False) as temp_readme:
|
| temp_readme.write(str(card))
|
| return temp_readme.name, valid_folders
|
| return None, valid_folders
|
|
|
| def _upload_files_and_readme(self, api, current_dir, region_prefix, temp_readme_path, valid_folders):
|
| self.log_dataset("🚀 [任务 5/5] 正在上传文件到 Hugging Face...")
|
| local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
|
|
|
| for sub_folder_name in valid_folders:
|
| sub_folder_path = os.path.join(current_dir, sub_folder_name)
|
| path_in_repo = f"{region_prefix}/{local_folder_name}/{sub_folder_name}"
|
|
|
| self.log_dataset(f" 📦 开始上传: {sub_folder_name} -> {path_in_repo}")
|
| try:
|
| api.upload_folder(
|
| folder_path=sub_folder_path,
|
| repo_id=DATASET_REPO_ID,
|
| repo_type="dataset",
|
| path_in_repo=path_in_repo,
|
| commit_message=f"Update files for {path_in_repo}",
|
| allow_patterns=["*.jpg", "*.jpeg", "*.txt", "*.png", "*.parquet"]
|
| )
|
| self.log_dataset(f" 🎉 {sub_folder_name} 上传完成！")
|
| except Exception as e:
|
| self.log_dataset(f" ❌ {sub_folder_name} 上传失败: {e}")
|
|
|
| if temp_readme_path:
|
| self.log_dataset("\n 📝 正在推送最新的 README.md...")
|
| try:
|
| api.upload_file(
|
| path_or_fileobj=temp_readme_path,
|
| path_in_repo="README.md",
|
| repo_id=DATASET_REPO_ID,
|
| repo_type="dataset",
|
| commit_message="Auto-update README YAML configs"
|
| )
|
| self.log_dataset(" 🎉 README.md 已成功更新到云端！")
|
| except Exception as e:
|
| self.log_dataset(f" ❌ README.md 上传失败: {e}")
|
| finally:
|
| os.remove(temp_readme_path)
|
|
|
| if __name__ == "__main__":
|
| app = ModelsManager()
|
| app.mainloop() |
