diff --git a/EasyQuantizationGUI.bat b/EasyQuantizationGUI.bat index 0ce4842..dcdbfd4 100644 --- a/EasyQuantizationGUI.bat +++ b/EasyQuantizationGUI.bat @@ -1,30 +1,79 @@ @echo off -REM Check if pip is installed and show output only if it needs to be installed +setlocal EnableDelayedExpansion + +set "APP_ENTRY=EasyQuantizationGUI.py" +set "VENV_DIR=venv" +set "VENV_ACTIVATE=%VENV_DIR%\Scripts\activate.bat" +set "REQUIREMENTS=requirements.txt" + +REM ── Python availability ──────────────────────────────────────────────────── +python --version >nul 2>&1 +if %ERRORLEVEL% neq 0 ( + echo [ERROR] Python was not found. Please install Python and ensure it is on your PATH. + goto :error +) + +REM ── pip availability ─────────────────────────────────────────────────────── python -m pip --version >nul 2>&1 if %ERRORLEVEL% neq 0 ( - echo Installing pip... + echo [INFO] pip not found. Installing via ensurepip... python -m ensurepip --default-pip + if %ERRORLEVEL% neq 0 ( + echo [ERROR] Failed to install pip. + goto :error + ) ) else ( python -m pip install --upgrade pip >nul 2>&1 ) -REM Check if virtual environment exists, create if it doesn't -if not exist "venv" ( - echo Creating virtual environment... - python -m venv venv - call venv\Scripts\activate - echo Installing requirements... - pip install -r requirements.txt - echo Setup complete! +REM ── Virtual environment setup ────────────────────────────────────────────── +if not exist "%VENV_DIR%\" ( + echo [INFO] Virtual environment not found. Creating "%VENV_DIR%"... + python -m venv "%VENV_DIR%" + if %ERRORLEVEL% neq 0 ( + echo [ERROR] Failed to create virtual environment. + goto :error + ) + + call "%VENV_ACTIVATE%" + if %ERRORLEVEL% neq 0 ( + echo [ERROR] Failed to activate virtual environment. + goto :error + ) + + echo [INFO] Installing dependencies from "%REQUIREMENTS%"... + pip install -r "%REQUIREMENTS%" + if %ERRORLEVEL% neq 0 ( + echo [ERROR] Dependency installation failed. + goto :error + ) + + echo [INFO] Setup complete. echo. ) else ( - call venv\Scripts\activate + call "%VENV_ACTIVATE%" + if %ERRORLEVEL% neq 0 ( + echo [ERROR] Failed to activate virtual environment. + goto :error + ) ) -REM Run the application -python EasyQuantizationGUI.py - -REM Keep the window open if there's an error +REM ── Launch application ───────────────────────────────────────────────────── +echo [INFO] Starting %APP_ENTRY%... +python "%APP_ENTRY%" if %ERRORLEVEL% neq 0 ( - pause -) \ No newline at end of file + echo [ERROR] Application exited with an error (code: %ERRORLEVEL%). + goto :error +) + +goto :end + +:error +echo. +echo Press any key to exit... +pause >nul +exit /b 1 + +:end +endlocal +exit /b 0 \ No newline at end of file diff --git a/EasyQuantizationGUI.py b/EasyQuantizationGUI.py index c56a192..ed7f6b5 100644 --- a/EasyQuantizationGUI.py +++ b/EasyQuantizationGUI.py @@ -1,8 +1,7 @@ -VERSION = "1.11" +VERSION = "1.12" import sys import subprocess -import importlib import os def install(package): @@ -17,338 +16,550 @@ def install(package): import tkinter as tk from tkinter import filedialog, ttk, messagebox -import os import shutil import winsound import tkinter.scrolledtext as scrolledtext +# ── Palette ────────────────────────────────────────────────────────────────── +BG = "#f7fafc" +SURFACE = "#ffffff" +SURFACE2 = "#f0f4f8" +BORDER = "#e2e8f0" +ACCENT = "#2563eb" +ACCENT2 = "#7c3aed" +SUCCESS = "#059669" +WARNING = "#b45309" +DANGER = "#dc2626" +TEXT = "#1d283b" +TEXT_MUTED= "#475569" +TEXT_DIM = "#64748b" +WHITE = "#ffffff" + +FONT_TITLE = ("Segoe UI", 18, "bold") +FONT_SUB = ("Segoe UI", 10) +FONT_LABEL = ("Segoe UI", 9, "bold") +FONT_BODY = ("Segoe UI", 9) +FONT_MONO = ("Consolas", 9) +FONT_BADGE = ("Segoe UI", 8, "bold") + +SUPPORTED_MODELS = [ + "Flux", + "SD3", + "Aurora", + "HiDream", + "Hyvid", + "Wan", + "LTXV", + "SDXL", + "SD1", +] + +QUANTIZE_LEVELS = [ + "Q2_K", "Q2_K_S", + "Q3_K", "Q3_K_L", "Q3_K_M", "Q3_K_S", + "Q4_0", "Q4_1", "Q4_K", "Q4_K_M", "Q4_K_S", + "Q5_0", "Q5_1", "Q5_K", "Q5_K_M", "Q5_K_S", + "Q6_K", "Q8_0", + "F16", "BF16", "F32", +] + +# ── Helpers ─────────────────────────────────────────────────────────────────── +def resource_path(relative_path): + try: + base_path = sys._MEIPASS + except Exception: + base_path = os.path.abspath(".") + return os.path.join(base_path, relative_path) + def scroll_entry_to_end(entry): entry.xview_moveto(1) -def browse_file(entry): - file_path = filedialog.askopenfilename(filetypes=[("Model files", "*.safetensors *.gguf *.sft")]) - if file_path: - file_path = file_path.replace('\\', '/') # Ensure forward slashes - entry.delete(0, tk.END) - entry.insert(0, file_path) - scroll_entry_to_end(entry) - suggest_output_file() # Call this instead of update_output_file - -def suggest_output_file(): - input_file = input_entry.get() - quantize_level = quantize_level_var.get() - if input_file: - input_dir = os.path.dirname(input_file) - input_filename = os.path.basename(input_file) - input_name, _ = os.path.splitext(input_filename) - output_file = f"{input_dir}/{input_name}-{quantize_level}.gguf" - output_entry.delete(0, tk.END) - output_entry.insert(0, output_file) - scroll_entry_to_end(output_entry) - -def browse_output_file(entry): - # Get the current input file and quantization level - input_file = input_entry.get() - quantize_level = quantize_level_var.get() - - # Generate a default output filename - if input_file: - input_dir = os.path.dirname(input_file) - input_filename = os.path.basename(input_file) - input_name, _ = os.path.splitext(input_filename) - default_filename = f"{input_name}-{quantize_level}.gguf" - else: - default_filename = f"output-{quantize_level}.gguf" - input_dir = "/" - - # Open the file dialog with the default filename - file_path = filedialog.asksaveasfilename( - initialdir=input_dir, - initialfile=default_filename, - defaultextension=".gguf", - filetypes=[("GGUF files", "*.gguf")] +# ── Styled widgets ──────────────────────────────────────────────────────────── +def make_entry(parent, **kwargs): + e = tk.Entry( + parent, + bg=SURFACE2, fg=TEXT, + insertbackground=ACCENT, + relief="flat", + highlightthickness=1, + highlightbackground=BORDER, + highlightcolor=ACCENT, + font=FONT_BODY, + **kwargs, + ) + return e + +def make_button(parent, text, command, color=ACCENT, fg=TEXT, width=None, font=None, hover_bg=None): + kw = dict(width=width) if width else {} + + # default hover selection: accent for primary buttons, subtle gray for surface buttons + if hover_bg is None: + hover_bg = BORDER if color == SURFACE2 else ACCENT2 + + b = tk.Button( + parent, + text=text, + command=command, + bg=color, + fg=fg, + disabledforeground=fg, + activebackground=hover_bg, + activeforeground=BG, + relief="flat", + cursor="hand2", + font=font or FONT_LABEL, + padx=12, + pady=6, + bd=0, + **kw, ) - - if file_path: - file_path = file_path.replace('\\', '/') # Ensure forward slashes - entry.delete(0, tk.END) - entry.insert(0, file_path) - scroll_entry_to_end(entry) - -def disable_ui(): - global input_entry, output_entry, input_browse, output_browse, quantize_dropdown, run_button - input_entry.config(state='disabled') - output_entry.config(state='disabled') - input_browse.config(state='disabled') - output_browse.config(state='disabled') - quantize_dropdown.config(state='disabled') - run_button.config(state='disabled') - -def enable_ui(): - global input_entry, output_entry, input_browse, output_browse, quantize_dropdown, run_button - input_entry.config(state='normal') - output_entry.config(state='normal') - input_browse.config(state='normal') - output_browse.config(state='normal') - quantize_dropdown.config(state='readonly') - run_button.config(state='normal') - -def run_llama_quantize(): - input_file = input_entry.get() - output_file = output_entry.get() - quantize_level = quantize_level_var.get() - - if not input_file or not output_file: - messagebox.showerror("Error", "Please select both input and output files.") - return - - # Check if input and output files are the same - if os.path.abspath(input_file) == os.path.abspath(output_file): - messagebox.showerror("Error", "Input and output files cannot be the same.") - return - - output_dir = os.path.dirname(output_file) - required_space = 40_000_000_000 # ~40 GB (a bit more than 36.5 GB) - available_space = shutil.disk_usage(output_dir).free - - if available_space < required_space: - required_gb = required_space / (1024**3) - available_gb = available_space / (1024**3) - messagebox.showerror("Error", f"You need {required_gb:.1f} GB of drive space to continue. Only {available_gb:.1f} GB available.") - return - - disable_ui() - - # Clear previous log - process_text.delete('1.0', tk.END) - root.update() - - is_input_gguf = input_file.lower().endswith(".gguf") - temp_gguf_file = None # Initialize temp_gguf_file - - if not is_input_gguf: - process_text.insert(tk.END, "Starting conversion process (Safetensors/SFT -> GGUF)...\n") - process_text.see(tk.END) - root.update() - - # Convert the input file to GGUF format - convert_py_path = resource_path("convert.py") - output_dir = os.path.dirname(output_file) - # Use a more descriptive temporary file name based on the output file - output_name, _ = os.path.splitext(os.path.basename(output_file)) - temp_gguf_file = os.path.join(output_dir, f"{output_name}_temp_conversion.gguf") - - - # Add cleanup of existing temp file - if os.path.exists(temp_gguf_file): - try: - os.remove(temp_gguf_file) - process_text.insert(tk.END, "Cleaned up existing temporary file.\n") - process_text.see(tk.END) - root.update() - except Exception as e: - process_text.insert(tk.END, f"Error cleaning up temporary file: {e}\n") - process_text.see(tk.END) - root.update() - enable_ui() - return - - try: - startupinfo = subprocess.STARTUPINFO() - startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW - startupinfo.wShowWindow = subprocess.SW_HIDE - - # Get the Python executable path from the current environment - pythonpath = sys.executable - - process = subprocess.Popen([pythonpath, convert_py_path, "--src", input_file, "--dst", temp_gguf_file], - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - bufsize=1, universal_newlines=True, startupinfo=startupinfo) - - for line in process.stdout: - process_text.insert(tk.END, line) - process_text.see(tk.END) - root.update() - - process.wait() - if process.returncode != 0: - raise subprocess.CalledProcessError(process.returncode, process.args) - - process_text.insert(tk.END, "Conversion completed successfully.\n") - process_text.see(tk.END) - root.update() - - except subprocess.CalledProcessError as e: - process_text.insert(tk.END, f"Error converting file: {e}\n") - process_text.insert(tk.END, f"Command: {e.cmd}\n") - process_text.insert(tk.END, f"Return code: {e.returncode}\n") - process_text.see(tk.END) - root.update() - # Clean up the temporary file even if conversion fails - if temp_gguf_file and os.path.exists(temp_gguf_file): - os.remove(temp_gguf_file) - enable_ui() - return - except Exception as e: # Catch other potential errors during conversion - process_text.insert(tk.END, f"An unexpected error occurred during conversion: {e}\n") - process_text.see(tk.END) - root.update() - if temp_gguf_file and os.path.exists(temp_gguf_file): - os.remove(temp_gguf_file) - enable_ui() - return - - # --- End of conversion block --- - else: - process_text.insert(tk.END, "Input is already GGUF. Skipping conversion step.\n") - process_text.see(tk.END) - root.update() - # If input is GGUF, llama-quantize will read directly from it - quantize_input_file = input_file - - # Determine the input file for the quantization step - quantize_input_file = temp_gguf_file if temp_gguf_file else input_file - - # Quantize the file (either the temporary one or the original GGUF) - llama_quantize_path = resource_path("llama-quantize.exe") - process_text.insert(tk.END, "Starting quantization process...\n") - process_text.see(tk.END) - root.update() - - try: - startupinfo = subprocess.STARTUPINFO() - startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW - startupinfo.wShowWindow = subprocess.SW_HIDE - - # Use quantize_input_file determined above - process = subprocess.Popen([llama_quantize_path, quantize_input_file, output_file, quantize_level], - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - bufsize=1, universal_newlines=True, startupinfo=startupinfo) - - for line in process.stdout: - process_text.insert(tk.END, line) - process_text.see(tk.END) - root.update() - - process.wait() - if process.returncode != 0: - # If quantization failed and we used a temp file, report the temp file name - if temp_gguf_file: - process_text.insert(tk.END, f"Quantization command failed on temporary file: {temp_gguf_file}\n") - raise subprocess.CalledProcessError(process.returncode, process.args) - - process_text.insert(tk.END, "Quantization completed successfully.\n") - except subprocess.CalledProcessError as e: - process_text.insert(tk.END, f"Error running llama-quantize: {e}\n") - process_text.insert(tk.END, f"Command: {e.cmd}\n") - process_text.insert(tk.END, f"Return code: {e.returncode}\n") - process_text.see(tk.END) - root.update() - except Exception as e: # Catch other potential errors during quantization - process_text.insert(tk.END, f"An unexpected error occurred during quantization: {e}\n") - process_text.see(tk.END) - root.update() - finally: - # Clean up the temporary file if it was created - if temp_gguf_file and os.path.exists(temp_gguf_file): - try: - os.remove(temp_gguf_file) - process_text.insert(tk.END, "Cleaned up temporary conversion file.\n") - process_text.see(tk.END) - root.update() - except Exception as e: - process_text.insert(tk.END, f"Error cleaning up temporary file {temp_gguf_file}: {e}\n") - process_text.see(tk.END) - root.update() - - process_text.insert(tk.END, "Process finished.\n") # Changed message slightly - process_text.see(tk.END) - root.update() - - enable_ui() - - # Play sound effect - winsound.PlaySound("SystemAsterisk", winsound.SND_ALIAS) - -def main(): - global root, process_text, input_entry, output_entry, quantize_dropdown, run_button, quantize_level_var - global input_browse, output_browse # Add these two variables - root = tk.Tk() - root.title(f"Easy Quantization GUI v{VERSION}") - root.geometry("800x600") - - # Quantize level selection - quantize_frame = tk.Frame(root) - quantize_frame.pack(pady=10, padx=10) - - quantize_label = tk.Label(quantize_frame, text="Quantize Level:") - quantize_label.pack(side=tk.LEFT) - - quantize_levels = ["Q2_K", "Q2_K_S", "Q3_K", "Q3_K_L", "Q3_K_M", "Q3_K_S", "Q4_0", "Q4_1", "Q4_K", "Q4_K_M", "Q4_K_S", "Q5_0", "Q5_1", "Q5_K", "Q5_K_M", "Q5_K_S", "Q6_K", "Q8_0", "F16", "BF16", "F32"] - quantize_level_var = tk.StringVar(root) - quantize_level_var.set("Q8_0") # Set default value to Q8_0 - quantize_dropdown = ttk.Combobox(quantize_frame, textvariable=quantize_level_var, values=quantize_levels, state="readonly") - quantize_dropdown.pack(side=tk.LEFT) - quantize_dropdown.bind("<>", lambda event: suggest_output_file()) + # closures capture current colors to avoid late-binding issues + def _on_enter(event, bg=hover_bg): + b.config(bg=bg) - # Input file selection - input_frame = tk.Frame(root) - input_frame.pack(pady=10, padx=10, fill=tk.X) + def _on_leave(event, bg=color): + b.config(bg=bg) - input_label = tk.Label(input_frame, text="Input File:") - input_label.pack(side=tk.LEFT) + b.bind("", _on_enter) + b.bind("", _on_leave) + return b - input_entry = tk.Entry(input_frame) - input_entry.pack(side=tk.LEFT, expand=True, fill=tk.X) +def make_label(parent, text, font=None, fg=TEXT, **kwargs): + return tk.Label(parent, text=text, bg=BG, fg=fg, font=font or FONT_BODY, **kwargs) - input_browse = tk.Button(input_frame, text="Browse", command=lambda: browse_file(input_entry)) - input_browse.pack(side=tk.RIGHT) +def divider(parent): + tk.Frame(parent, bg=BORDER, height=1).pack(fill="x", pady=8, padx=0) - # Add binding to scroll input entry when it gains focus - input_entry.bind("", lambda event: scroll_entry_to_end(input_entry)) +# ── Main App ────────────────────────────────────────────────────────────────── +class EasyQuantGUI: + def __init__(self): + self.root = tk.Tk() + self.root.title(f"Easy Quantization GUI v{VERSION}") + self.root.geometry("900x680") + self.root.minsize(760, 580) + self.root.configure(bg=BG) + self._set_icon() - # Output file selection - output_frame = tk.Frame(root) - output_frame.pack(pady=10, padx=10, fill=tk.X) + self.quantize_level_var = tk.StringVar(value="Q8_0") + self._build_ui() + self.root.mainloop() - output_label = tk.Label(output_frame, text="Output File:") - output_label.pack(side=tk.LEFT) + def _set_icon(self): + try: + self.root.iconbitmap(resource_path("icon.ico")) + except Exception: + pass + + # ── UI Builder ──────────────────────────────────────────────────────────── + def _build_ui(self): + # ── Header + header = tk.Frame(self.root, bg=SURFACE, pady=14) + header.pack(fill="x") + + tk.Label( + header, text="Easy Quantization GUI", + bg=SURFACE, fg=TEXT, font=FONT_TITLE, + ).pack(side="left", padx=20) + tk.Label( + header, text=f"v{VERSION}", + bg=SURFACE, fg=TEXT_MUTED, font=FONT_SUB, + ).pack(side="left") + + # ── Body (two-column: left=form, right=models) + body = tk.Frame(self.root, bg=BG) + body.pack(fill="both", expand=True, padx=0, pady=0) + body.columnconfigure(0, weight=3) + body.columnconfigure(1, weight=1) + body.rowconfigure(0, weight=1) + + self._build_left(body) + self._build_right(body) + + def _build_left(self, parent): + left = tk.Frame(parent, bg=BG, padx=20, pady=16) + left.grid(row=0, column=0, sticky="nsew") + left.columnconfigure(0, weight=1) + + # ── Section: Files + self._section_label(left, "Files") + + # Input + self._field_row( + left, + label="Input file", + hint=".safetensors · .gguf · .sft", + entry_attr="input_entry", + browse_cmd=self._browse_input, + browse_attr="input_browse", + ) + + # Output + self._field_row( + left, + label="Output file", + hint=".gguf", + entry_attr="output_entry", + browse_cmd=self._browse_output, + browse_attr="output_browse", + ) + + divider(left) + + # ── Section: Quantization + self._section_label(left, "Quantization") + + q_frame = tk.Frame(left, bg=BG) + q_frame.pack(fill="x", pady=(0, 12)) + + make_label(q_frame, "Level:", font=FONT_LABEL, fg=TEXT_DIM).pack(side="left") + + style = ttk.Style() + style.theme_use("clam") + style.configure( + "Dark.TCombobox", + fieldbackground=SURFACE2, background=SURFACE2, + foreground=TEXT, selectforeground=TEXT, + selectbackground=ACCENT, + bordercolor=BORDER, arrowcolor=ACCENT, + relief="flat", + ) + style.map( + "Dark.TCombobox", + fieldbackground=[("readonly", SURFACE2)], + foreground=[("readonly", TEXT)], + selectbackground=[("readonly", SURFACE2)], + ) + + self.quantize_dropdown = ttk.Combobox( + q_frame, + textvariable=self.quantize_level_var, + values=QUANTIZE_LEVELS, + state="readonly", + width=14, + style="Dark.TCombobox", + font=FONT_BODY, + ) + self.quantize_dropdown.pack(side="left", padx=(8, 0)) + self.quantize_dropdown.bind("<>", lambda _: self._suggest_output()) + + # Quant description badge + self.quant_desc_var = tk.StringVar(value=self._quant_hint("Q8_0")) + tk.Label( + q_frame, + textvariable=self.quant_desc_var, + bg=SURFACE2, fg=ACCENT, + font=FONT_BADGE, + padx=8, pady=3, relief="flat", + ).pack(side="left", padx=10) + + self.quantize_level_var.trace_add("write", lambda *_: self.quant_desc_var.set(self._quant_hint(self.quantize_level_var.get()))) + + divider(left) + + # ── Run button + run_row = tk.Frame(left, bg=BG) + run_row.pack(fill="x", pady=(0, 4)) + + self.run_button = make_button( + run_row, "▶ Run Quantization", + self._run, + color=ACCENT, fg=BG, + font=("Segoe UI", 10, "bold"), + ) + self.run_button.pack(side="left") + + self.status_label = tk.Label( + run_row, text="", bg=BG, fg=TEXT_MUTED, font=FONT_BODY, + ) + self.status_label.pack(side="left", padx=14) + + divider(left) + + # ── Process log + self._section_label(left, "Process log") + + log_frame = tk.Frame(left, bg=SURFACE2, relief="flat", highlightthickness=1, highlightbackground=BORDER) + log_frame.pack(fill="both", expand=True, pady=(0, 8)) + + self.process_text = scrolledtext.ScrolledText( + log_frame, + wrap=tk.WORD, + bg=SURFACE2, fg=TEXT_DIM, + insertbackground=ACCENT, + font=FONT_MONO, + relief="flat", + borderwidth=0, + padx=10, pady=8, + ) + self.process_text.pack(fill="both", expand=True) + self.process_text.tag_config("ok", foreground=SUCCESS) + self.process_text.tag_config("err", foreground=DANGER) + self.process_text.tag_config("info", foreground=ACCENT) + self.process_text.tag_config("warn", foreground=WARNING) + + # Clear log button + make_button( + left, "Clear log", + lambda: self.process_text.delete("1.0", tk.END), + color=SURFACE2, fg=TEXT_DIM, font=FONT_BODY, + ).pack(side="right", pady=(0, 2)) + + def _build_right(self, parent): + right = tk.Frame(parent, bg=SURFACE, padx=16, pady=16, width=200) + right.grid(row=0, column=1, sticky="nsew") + right.pack_propagate(False) + + self._section_label(right, "Supported models", bg=SURFACE) + + canvas = tk.Canvas(right, bg=SURFACE, highlightthickness=0, bd=0) + canvas.pack(fill="both", expand=True) + + inner = tk.Frame(canvas, bg=SURFACE) + win_id = canvas.create_window((0, 0), window=inner, anchor="nw") + + def on_configure(e): + canvas.configure(scrollregion=canvas.bbox("all")) + canvas.itemconfig(win_id, width=canvas.winfo_width()) + + inner.bind("", on_configure) + canvas.bind("", on_configure) + + for name in SUPPORTED_MODELS: + row = tk.Frame(inner, bg=SURFACE2, relief="flat") + row.pack(fill="x", pady=2) + + tk.Label(row, text="●", bg=SURFACE2, fg=ACCENT, font=("Segoe UI", 8)).pack(side="left", padx=(8, 4), pady=5) + tk.Label(row, text=name, bg=SURFACE2, fg=TEXT, font=FONT_BODY, anchor="w").pack(side="left", pady=5) + + # bind mousewheel + # Remove scrollbar: keep simple static layout without scroll controls + + # ── Field helper ────────────────────────────────────────────────────────── + def _field_row(self, parent, label, hint, entry_attr, browse_cmd, browse_attr): + tk.Label(parent, text=label, bg=BG, fg=TEXT_DIM, font=FONT_LABEL, anchor="w").pack(fill="x", pady=(6, 2)) + + row = tk.Frame(parent, bg=BG) + row.pack(fill="x", pady=(0, 8)) + row.columnconfigure(0, weight=1) + + e = make_entry(row) + e.grid(row=0, column=0, sticky="ew", ipady=6) + e.bind("", lambda _: scroll_entry_to_end(e)) + e.bind("", lambda _: scroll_entry_to_end(e)) + setattr(self, entry_attr, e) + + btn = make_button(row, "Browse", browse_cmd, color=SURFACE2, fg=TEXT_DIM, font=FONT_BODY) + btn.grid(row=0, column=1, padx=(6, 0)) + setattr(self, browse_attr, btn) + + tk.Label(parent, text=hint, bg=BG, fg=TEXT_MUTED, font=("Segoe UI", 8), anchor="w").pack(fill="x") + + def _section_label(self, parent, text, bg=BG): + tk.Label(parent, text=text.upper(), bg=bg, fg=TEXT_MUTED, font=("Segoe UI", 8, "bold"), anchor="w").pack(fill="x", pady=(0, 6)) + + # ── Quantization hint ──────────────────────────────────────────────────── + def _quant_hint(self, level): + hints = { + "Q2_K": "Extreme compression — very small size, low accuracy", + "Q2_K_S": "Ultra-compact — lowest precision, smallest footprint", + + "Q3_K": "Very small — strong compression, reduced fidelity", + "Q3_K_L": "Low-precision 3-bit — smaller size, lower quality", + "Q3_K_M": "Medium 3-bit — balanced size vs. quality", + "Q3_K_S": "Small 3-bit — slightly better quality than lowest", + + "Q4_0": "4-bit baseline — compact with reasonable accuracy", + "Q4_1": "4-bit variant — slightly different trade-offs", + "Q4_K": "4-bit k-means — efficient with good accuracy", + "Q4_K_M": "4-bit high-quality — improved fidelity", + "Q4_K_S": "4-bit small — optimized for minimal size", + + "Q5_0": "5-bit baseline — better fidelity than 4-bit", + "Q5_1": "5-bit variant — alternate trade-offs", + "Q5_K": "5-bit k-means — higher precision for quality", + "Q5_K_M": "5-bit medium — balanced precision and size", + "Q5_K_S": "5-bit small — space-optimized 5-bit", + + "Q6_K": "6-bit k-means — near-lossless, high quality", + "Q8_0": "8-bit — highest quantized precision, best quality", + + "F16": "Float16 — lower-precision float, good accuracy", + "BF16": "BFloat16 — float variant with wide dynamic range", + "F32": "Float32 — full precision, no quantization", + } + return hints.get(level, "") + + # ── Browse callbacks ────────────────────────────────────────────────────── + def _browse_input(self): + path = filedialog.askopenfilename(filetypes=[("Model files", "*.safetensors *.gguf *.sft")]) + if path: + path = path.replace("\\", "/") + self.input_entry.delete(0, tk.END) + self.input_entry.insert(0, path) + scroll_entry_to_end(self.input_entry) + self._suggest_output() + + def _browse_output(self): + input_file = self.input_entry.get() + level = self.quantize_level_var.get() + if input_file: + idir = os.path.dirname(input_file) + iname = os.path.splitext(os.path.basename(input_file))[0] + default = f"{iname}-{level}.gguf" + else: + idir, default = "/", f"output-{level}.gguf" + + path = filedialog.asksaveasfilename( + initialdir=idir, initialfile=default, + defaultextension=".gguf", filetypes=[("GGUF files", "*.gguf")], + ) + if path: + path = path.replace("\\", "/") + self.output_entry.delete(0, tk.END) + self.output_entry.insert(0, path) + scroll_entry_to_end(self.output_entry) + + def _suggest_output(self): + inp = self.input_entry.get() + level = self.quantize_level_var.get() + if inp: + idir = os.path.dirname(inp) + iname = os.path.splitext(os.path.basename(inp))[0] + out = f"{idir}/{iname}-{level}.gguf" + self.output_entry.delete(0, tk.END) + self.output_entry.insert(0, out) + scroll_entry_to_end(self.output_entry) + + # ── UI state ────────────────────────────────────────────────────────────── + def _disable_ui(self): + for w in (self.input_entry, self.output_entry, self.input_browse, + self.output_browse, self.quantize_dropdown, self.run_button): + w.config(state="disabled") + self.status_label.config(text="⏳ Running…", fg=WARNING) + + def _enable_ui(self): + self.input_entry.config(state="normal") + self.output_entry.config(state="normal") + self.input_browse.config(state="normal") + self.output_browse.config(state="normal") + self.quantize_dropdown.config(state="readonly") + self.run_button.config(state="normal") + self.status_label.config(text="") + + # ── Log helpers ─────────────────────────────────────────────────────────── + def _log(self, text, tag=None): + self.process_text.insert(tk.END, text, tag or "") + self.process_text.see(tk.END) + self.root.update() + + # ── Run ─────────────────────────────────────────────────────────────────── + def _run(self): + input_file = self.input_entry.get().strip() + output_file = self.output_entry.get().strip() + level = self.quantize_level_var.get() + + if not input_file or not output_file: + messagebox.showerror("Missing files", "Please select both input and output files.") + return - output_entry = tk.Entry(output_frame) - output_entry.pack(side=tk.LEFT, expand=True, fill=tk.X) + if os.path.abspath(input_file) == os.path.abspath(output_file): + messagebox.showerror("Same file", "Input and output files cannot be the same.") + return - output_browse = tk.Button(output_frame, text="Browse", command=lambda: browse_output_file(output_entry)) - output_browse.pack(side=tk.RIGHT) + output_dir = os.path.dirname(output_file) + required = 40_000_000_000 + available = shutil.disk_usage(output_dir).free + if available < required: + messagebox.showerror( + "Not enough space", + f"Need {required/1e9:.1f} GB free, only {available/1e9:.1f} GB available.", + ) + return - # Add binding to scroll output entry when it gains focus - output_entry.bind("", lambda event: scroll_entry_to_end(output_entry)) + self._disable_ui() + self.process_text.delete("1.0", tk.END) + self.root.update() - # Run button - run_button = tk.Button(root, text="Run Quantization", command=run_llama_quantize) - run_button.pack(pady=20) + is_gguf = input_file.lower().endswith(".gguf") + temp_gguf = None - # Add process log to bottom of main window - process_frame = tk.Frame(root) - process_frame.pack(pady=10, padx=10, fill=tk.BOTH, expand=True) + # ── Step 1: Convert if needed + if not is_gguf: + self._log("── Step 1/2: Converting to GGUF…\n", "info") + out_name = os.path.splitext(os.path.basename(output_file))[0] + temp_gguf = os.path.join(output_dir, f"{out_name}_temp_conversion.gguf") - process_label = tk.Label(process_frame, text="Process Log:") - process_label.pack(side=tk.TOP, anchor='w') + if os.path.exists(temp_gguf): + try: + os.remove(temp_gguf) + self._log("Cleaned up existing temp file.\n") + except Exception as e: + self._log(f"Error removing temp file: {e}\n", "err") + self._enable_ui() + return - process_text = scrolledtext.ScrolledText(process_frame, wrap=tk.WORD, height=15) - process_text.pack(expand=True, fill=tk.BOTH) + try: + si = subprocess.STARTUPINFO() + si.dwFlags |= subprocess.STARTF_USESHOWWINDOW + si.wShowWindow = subprocess.SW_HIDE + + proc = subprocess.Popen( + [sys.executable, resource_path("convert.py"), "--src", input_file, "--dst", temp_gguf], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1, universal_newlines=True, startupinfo=si, + ) + for line in proc.stdout: + self._log(line) + proc.wait() + if proc.returncode != 0: + raise subprocess.CalledProcessError(proc.returncode, proc.args) + self._log("Conversion complete.\n", "ok") + except Exception as e: + self._log(f"Conversion failed: {e}\n", "err") + if temp_gguf and os.path.exists(temp_gguf): + os.remove(temp_gguf) + self._enable_ui() + return + else: + self._log("Input is GGUF — skipping conversion.\n", "info") - root.mainloop() + # ── Step 2: Quantize + quant_input = temp_gguf if temp_gguf else input_file + step_label = "2/2" if not is_gguf else "1/1" + self._log(f"\n── Step {step_label}: Quantizing ({level})…\n", "info") -def resource_path(relative_path): - """ Get absolute path to resource, works for dev and for PyInstaller """ - try: - # PyInstaller creates a temp folder and stores path in _MEIPASS - base_path = sys._MEIPASS - except Exception: - base_path = os.path.abspath(".") + try: + si = subprocess.STARTUPINFO() + si.dwFlags |= subprocess.STARTF_USESHOWWINDOW + si.wShowWindow = subprocess.SW_HIDE + + proc = subprocess.Popen( + [resource_path("llama-quantize.exe"), quant_input, output_file, level], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1, universal_newlines=True, startupinfo=si, + ) + for line in proc.stdout: + self._log(line) + proc.wait() + if proc.returncode != 0: + raise subprocess.CalledProcessError(proc.returncode, proc.args) + self._log("\nQuantization complete!\n", "ok") + except Exception as e: + self._log(f"Quantization failed: {e}\n", "err") + finally: + if temp_gguf and os.path.exists(temp_gguf): + try: + os.remove(temp_gguf) + self._log("Temp file cleaned up.\n") + except Exception as e: + self._log(f"Could not remove temp file: {e}\n", "warn") + + self._log("\n── Done ──\n", "info") + self.status_label.config(text="✔ Done", fg=SUCCESS) + self._enable_ui() + winsound.PlaySound("SystemAsterisk", winsound.SND_ALIAS) - return os.path.join(base_path, relative_path) if __name__ == "__main__": - main() + EasyQuantGUI() \ No newline at end of file diff --git a/README.md b/README.md index 87b8926..b50516d 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,30 @@ This application basically just simplifies this process: https://github.com/city96/ComfyUI-GGUF/tree/main/tools -![screenshot](https://github.com/user-attachments/assets/11d2315b-9ea4-4caf-a3a0-e211defae7a7) +explorer_FR8IunMN25 + Run `EasyQuantizationGUI.bat` to start the application. Requirements: - [Python](https://www.python.org/downloads/windows/) - Windows (can be adjusted later to support Linux) + +Supported Model Architectures +----------------------------- + +This tool supports quantization for a wide range of model architectures commonly used in image generation and related tasks. The following architectures are supported: + +- Flux +- SD3 +- Aurora +- HiDream +- Hyvid +- Wan +- LTXV +- SDXL +- SD1 + +Usage +----- + +Run `EasyQuantizationGUI.bat` to start the application. diff --git a/convert.py b/convert.py index 4ee893e..76ba3df 100644 --- a/convert.py +++ b/convert.py @@ -18,9 +18,10 @@ class ModelTemplate: keys_detect = [] # list of lists to match in state dict keys_banned = [] # list of keys that should mark model as invalid for conversion keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason + keys_ignore = [] # list of strings to ignore keys by when found def handle_nd_tensor(self, key, data): - raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") + raise NotImplementedError(f"Tensor exceeds maximum supported dimensions by C++ code: key='{key}', shape={data.shape}") class ModelFlux(ModelTemplate): arch = "flux" @@ -60,6 +61,17 @@ class ModelHiDream(ModelTemplate): "img_emb.emb_pos" ] +class CosmosPredict2(ModelTemplate): + arch = "cosmos" + keys_detect = [ + ( + "blocks.0.mlp.layer1.weight", + "blocks.0.adaln_modulation_cross_attn.1.weight", + ) + ] + keys_hiprec = ["pos_embedder"] + keys_ignore = ["_extra_state", "accum_"] + class ModelHyVid(ModelTemplate): arch = "hyvid" keys_detect = [ @@ -73,9 +85,9 @@ def handle_nd_tensor(self, key, data): # hacky but don't have any better ideas path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? if os.path.isfile(path): - raise RuntimeError(f"5D tensor fix file already exists! {path}") + raise RuntimeError(f"5D tensor fix file already exists: {path}") fsd = {key: torch.from_numpy(data)} - tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + logging.warning("5D tensor detected — manual fix required: key='%s', shape=%s", key, data.shape) save_file(fsd, path) class ModelWan(ModelHyVid): @@ -127,8 +139,14 @@ class ModelSD1(ModelTemplate): ), # Non-diffusers ] -# The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] +class ModelLumina2(ModelTemplate): + arch = "lumina2" + keys_detect = [ + ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") + ] + +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, + ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] def is_model_arch(model, state_dict): # check if model is correct @@ -139,7 +157,7 @@ def is_model_arch(model, state_dict): matched = True invalid = any(key in state_dict for key in model.keys_banned) break - assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" + assert not invalid, "Model architecture not allowed for conversion (e.g. reference format instead of diffusers format)" return matched def detect_arch(state_dict): @@ -148,7 +166,7 @@ def detect_arch(state_dict): if is_model_arch(arch, state_dict): model_arch = arch() break - assert model_arch is not None, "Unknown model architecture!" + assert model_arch is not None, "Unknown model architecture" return model_arch def parse_args(): @@ -163,20 +181,32 @@ def parse_args(): return args def strip_prefix(state_dict): - # only keep unet with no prefix! + # prefix for mixed state dict prefix = None for pfx in ["model.diffusion_model.", "model."]: if any([x.startswith(pfx) for x in state_dict.keys()]): prefix = pfx break - sd = {} - for k, v in state_dict.items(): - if prefix and prefix not in k: - continue - if prefix: + # prefix for uniform state dict + if prefix is None: + for pfx in ["net."]: + if all([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # strip prefix if found + if prefix is not None: + logging.info("State dict prefix detected and will be stripped: '%s'", prefix) + sd = {} + for k, v in state_dict.items(): + if prefix not in k: + continue k = k.replace(prefix, "") - sd[k] = v + sd[k] = v + else: + logging.debug("No prefix found in state dict; loading keys as-is") + sd = state_dict return sd @@ -188,7 +218,7 @@ def load_state_dict(path): state_dict = state_dict[subkey] break if len(state_dict) < 20: - raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") + raise RuntimeError(f"Failed to load state dict via subkey — too few keys loaded: {list(state_dict.keys())}") else: state_dict = load_file(path) @@ -205,10 +235,15 @@ def handle_tensors(writer, state_dict, model_arch): max_name_len = name_lengths[0][1] if max_name_len > MAX_TENSOR_NAME_LENGTH: bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) - raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") + raise ValueError(f"Tensor name exceeds maximum allowed length of {MAX_TENSOR_NAME_LENGTH} characters. Offending tensors: {bad_list}") + for key, data in tqdm(state_dict.items()): old_dtype = data.dtype + if any(x in key for x in model_arch.keys_ignore): + logging.debug("Skipping ignored tensor: '%s'", key) + continue + if data.dtype == torch.bfloat16: data = data.to(torch.float32).numpy() # this is so we don't break torch 2.0.X @@ -263,14 +298,14 @@ def handle_tensors(writer, state_dict, model_arch): try: data = gguf.quants.quantize(data, data_qtype) except (AttributeError, gguf.QuantError) as e: - tqdm.write(f"falling back to F16: {e}") + logging.warning("Quantization to %s failed for '%s'; falling back to F16: %s", data_qtype.name, key, e) data_qtype = gguf.GGMLQuantizationType.F16 data = gguf.quants.quantize(data, data_qtype) new_name = key # do we need to rename? shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" - tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") + tqdm.write(f"{f'%-{max_name_len + 4}s' % new_name} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") writer.add_tensor(new_name, data, raw_dtype=data_qtype) @@ -278,7 +313,7 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False): # load & run model detection logic state_dict = load_state_dict(path) model_arch = detect_arch(state_dict) - logging.info(f"* Architecture detected from input: {model_arch.arch}") + logging.info("Architecture detected: '%s'", model_arch.arch) # detect & set dtype for output file dtypes = [x.dtype for x in state_dict.values()] @@ -302,9 +337,9 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False): if os.path.isfile(dst_path) and not overwrite: if interact: - input("Output exists enter to continue or ctrl+c to abort!") + input("Output file already exists. Press Enter to continue or Ctrl+C to abort.") else: - raise OSError("Output exists and overwriting is disabled!") + raise OSError("Output file already exists and overwriting is disabled") # handle actual file writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) @@ -320,11 +355,10 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False): fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" if os.path.isfile(fix): - logging.warning(f"\n### Warning! Fix file found at '{fix}'") - logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") + logging.warning("Fix file detected at '%s' — you likely need to run 'fix_5d_tensors.py' after quantization", fix) return dst_path, model_arch if __name__ == "__main__": args = parse_args() - convert_file(args.src, args.dst) + convert_file(args.src, args.dst) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 9811377..ad84dff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ -torch -tqdm -safetensors -gguf -sentencepiece -pyyaml -numpy +torch==2.7.1 +tqdm==4.67.1 +safetensors==0.4.5 +gguf==0.19.0 +sentencepiece==0.2.0 +PyYAML==6.0.2 +numpy==2.3.1 +packaging==24.2