mirror of https://github.com/leafspark/AutoGGUF
94 lines
3.4 KiB
Python
94 lines
3.4 KiB
Python
from PyQt6.QtWidgets import *
|
|
from PyQt6.QtCore import *
|
|
from PyQt6.QtGui import *
|
|
import os
|
|
import sys
|
|
import psutil
|
|
import subprocess
|
|
import time
|
|
import signal
|
|
import json
|
|
import platform
|
|
import requests
|
|
import zipfile
|
|
import traceback
|
|
from datetime import datetime
|
|
from imports_and_globals import open_file_safe
|
|
|
|
class QuantizationThread(QThread):
|
|
# Define custom signals for communication with the main thread
|
|
output_signal = pyqtSignal(str)
|
|
status_signal = pyqtSignal(str)
|
|
finished_signal = pyqtSignal()
|
|
error_signal = pyqtSignal(str)
|
|
model_info_signal = pyqtSignal(dict)
|
|
|
|
def __init__(self, command, cwd, log_file):
|
|
super().__init__()
|
|
self.command = command
|
|
self.cwd = cwd
|
|
self.log_file = log_file
|
|
self.process = None
|
|
self.model_info = {}
|
|
|
|
def run(self):
|
|
try:
|
|
# Start the subprocess
|
|
self.process = subprocess.Popen(
|
|
self.command,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
cwd=self.cwd,
|
|
)
|
|
# Open log file and process output
|
|
with open_file_safe(self.log_file, "w") as log:
|
|
for line in self.process.stdout:
|
|
line = line.strip()
|
|
self.output_signal.emit(line)
|
|
log.write(line + "\n")
|
|
log.flush()
|
|
self.status_signal.emit("In Progress")
|
|
self.parse_model_info(line)
|
|
|
|
# Wait for process to complete
|
|
self.process.wait()
|
|
if self.process.returncode == 0:
|
|
self.status_signal.emit("Completed")
|
|
self.model_info_signal.emit(self.model_info)
|
|
else:
|
|
self.error_signal.emit(
|
|
f"Process exited with code {self.process.returncode}"
|
|
)
|
|
self.finished_signal.emit()
|
|
except Exception as e:
|
|
self.error_signal.emit(str(e))
|
|
|
|
def parse_model_info(self, line):
|
|
# Parse output for model information
|
|
if "llama_model_loader: loaded meta data with" in line:
|
|
parts = line.split()
|
|
self.model_info["kv_pairs"] = parts[6]
|
|
self.model_info["tensors"] = parts[9]
|
|
elif "general.architecture" in line:
|
|
self.model_info["architecture"] = line.split("=")[-1].strip()
|
|
elif line.startswith("llama_model_loader: - kv"):
|
|
key = line.split(":")[2].strip()
|
|
value = line.split("=")[-1].strip()
|
|
self.model_info.setdefault("kv_data", {})[key] = value
|
|
elif line.startswith("llama_model_loader: - type"):
|
|
parts = line.split(":")
|
|
if len(parts) > 1:
|
|
quant_type = parts[1].strip()
|
|
tensors = parts[2].strip().split()[0]
|
|
self.model_info.setdefault("quantization_type", []).append(
|
|
f"{quant_type}: {tensors} tensors"
|
|
)
|
|
|
|
def terminate(self):
|
|
# Terminate the subprocess if it's still running
|
|
if self.process:
|
|
os.kill(self.process.pid, signal.SIGTERM)
|
|
self.process.wait(timeout=5)
|
|
if self.process.poll() is None:
|
|
os.kill(self.process.pid, signal.SIGKILL) |