nodes/thomaskippster-comfymodeldownloader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add Storage Optimizer with SHA-256 hashing and Civitai hash lookup

Browse Source
This commit is contained in:
Thomas
2026-04-26 17:48:33 +02:00
parent a75add1541
commit f271407a8f
9 changed files with 477 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4
View File
@@ -25,9 +25,13 @@ boot.key
# User Model List (internal DB)
uploaded_models.json
pending_downloads.json
model_hashes.json
# Local config & assets
workflow.json
*.json
!pom.xml
!src/**
*.png
*.bat
*.zip
pom.xml
+12
View File
@@ -58,6 +58,18 @@
<version>6.1.5</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>5.11.0</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-junit-jupiter</artifactId>
<version>5.11.0</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.formdev</groupId>
<artifactId>flatlaf</artifactId>
src/main/java/de/tki/comfymodels/Main.java
+117 -38
View File
@@ -39,6 +39,8 @@ import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.ArrayList;
import java.util.Map;
import java.util.HashMap;
import java.util.stream.Collectors;
@Component
@@ -58,6 +60,9 @@ public class Main extends JFrame {
@Autowired
private ModelListService modelListService;
@Autowired
private de.tki.comfymodels.service.impl.ModelHashRegistry hashRegistry;
private JTextField modelsPathField;
private JPasswordField geminiKeyField;
private JPasswordField hfTokenField;
@@ -326,8 +331,19 @@ public class Main extends JFrame {
JButton helpBtn = new JButton("ℹ Help & Security Information");
helpBtn.addActionListener(e -> showHelpDialog());
JButton verifyBtn = new JButton("🔍 Verify Local Models (Corruption Check)");
verifyBtn.addActionListener(e -> verifyLocalModels());
JButton verifyBtn = new JButton("🔍 Fast Verify (Corruption)");
verifyBtn.addActionListener(e -> verifyLocalModels(false));
JButton optimizeBtn = new JButton("👯 Storage Optimizer (Duplicates)");
optimizeBtn.addActionListener(e -> {
int confirm = JOptionPane.showConfirmDialog(this,
"The Storage Optimizer calculates SHA-256 hashes for all local models.\n" +
"This is EXTREMELY slow and resource-intensive for large libraries.\n\n" +
"Do you want to proceed?", "Performance Warning", JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE);
if (confirm == JOptionPane.YES_OPTION) {
verifyLocalModels(true);
}
});
backgroundCheck = new JCheckBox("Stay in Background on Close");
backgroundCheck.addActionListener(e -> configService.setBackgroundModeEnabled(backgroundCheck.isSelected()));
@@ -353,6 +369,7 @@ public class Main extends JFrame {
infoRow.add(helpBtn);
infoRow.add(verifyBtn);
infoRow.add(optimizeBtn);
infoRow.add(backgroundCheck);
infoRow.add(shutdownCheck);
infoRow.add(darkCheck);
@@ -629,7 +646,7 @@ public class Main extends JFrame {
);
}
private void verifyLocalModels() {
private void verifyLocalModels(boolean checkDuplicates) {
String baseDir = modelsPathField.getText();
if (baseDir == null || baseDir.isEmpty()) {
JOptionPane.showMessageDialog(this, "Please set a models directory first.");
@@ -642,10 +659,13 @@ public class Main extends JFrame {
return;
}
statusLabel.setText("Verifying models... please wait.");
String taskName = checkDuplicates ? "Verifying models & checking for duplicates" : "Verifying models (Fast Check)";
statusLabel.setText(taskName + "... please wait.");
new Thread(() -> {
try {
List<IModelValidator.ValidationResult> errors = new ArrayList<>();
Map<String, List<Path>> hashToPaths = new HashMap<>();
List<Path> allFiles = Files.walk(root.toPath())
.filter(Files::isRegularFile)
.filter(p -> {
@@ -663,55 +683,107 @@ public class Main extends JFrame {
IModelValidator.ValidationResult res = modelValidator.validateFile(p.toFile());
if (!res.ok) {
errors.add(res);
} else if (checkDuplicates) {
String hash = hashRegistry.getOrCalculateHash(p.toFile());
if (hash != null) {
hashToPaths.computeIfAbsent(hash, k -> new ArrayList<>()).add(p);
}
}
}
Map<String, List<Path>> duplicates = checkDuplicates ? hashToPaths.entrySet().stream()
.filter(e -> e.getValue().size() > 1)
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)) : new HashMap<>();
SwingUtilities.invokeLater(() -> {
statusLabel.setText("Verification finished. Found " + errors.size() + " issues.");
if (errors.isEmpty()) {
statusLabel.setText("Scan finished. Found " + errors.size() + " issues and " + duplicates.size() + " duplicate sets.");
if (errors.isEmpty() && duplicates.isEmpty()) {
JOptionPane.showMessageDialog(this, "All " + total + " models verified successfully!", "Verification Complete", JOptionPane.INFORMATION_MESSAGE);
} else {
StringBuilder sb = new StringBuilder("The following " + errors.size() + " files appear to be corrupted or invalid:\n\n");
return;
}
StringBuilder sb = new StringBuilder();
if (!errors.isEmpty()) {
sb.append("⚠️ CORRUPTED FILES (").append(errors.size()).append("):\n");
for (IModelValidator.ValidationResult err : errors) {
sb.append("- ").append(new File(err.filePath).getName())
.append(" (").append(err.message).append(")\n")
.append(" Path: ").append(err.filePath).append("\n\n");
sb.append("- ").append(new File(err.filePath).getName()).append(" (").append(err.message).append(")\n");
}
JTextArea textArea = new JTextArea(sb.toString());
textArea.setEditable(false);
JScrollPane scrollPane = new JScrollPane(textArea);
scrollPane.setPreferredSize(new Dimension(800, 500));
Object[] options = {"OK", "Delete All Corrupted Files"};
int choice = JOptionPane.showOptionDialog(this, scrollPane, "Verification Results - Issues Found",
JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE, null, options, options[0]);
if (choice == 1) { // Delete All Corrupted Files
int confirm = JOptionPane.showConfirmDialog(this,
"Are you sure you want to delete these " + errors.size() + " files?\nThis action cannot be undone.",
"Confirm Deletion", JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE);
if (confirm == JOptionPane.YES_OPTION) {
int deletedCount = 0;
for (IModelValidator.ValidationResult err : errors) {
File f = new File(err.filePath);
if (f.exists() && f.delete()) {
deletedCount++;
}
}
JOptionPane.showMessageDialog(this, "Deleted " + deletedCount + " files.");
statusLabel.setText("Cleanup finished. Deleted " + deletedCount + " files.");
sb.append("\n");
}
if (!duplicates.isEmpty()) {
sb.append("👯 DUPLICATES FOUND (").append(duplicates.size()).append(" sets):\n");
for (Map.Entry<String, List<Path>> entry : duplicates.entrySet()) {
sb.append("Hash: ").append(entry.getKey().substring(0, 12)).append("...\n");
for (Path p : entry.getValue()) {
sb.append(" - ").append(p.toFile().getAbsolutePath()).append(" (").append(formatSize(p.toFile().length())).append(")\n");
}
sb.append("\n");
}
}
JTextArea textArea = new JTextArea(sb.toString());
textArea.setEditable(false);
JScrollPane scrollPane = new JScrollPane(textArea);
scrollPane.setPreferredSize(new Dimension(850, 550));
Object[] options = {"OK", "Delete All Corrupted", "Manage Duplicates..."};
int choice = JOptionPane.showOptionDialog(this, scrollPane, "Verification & Duplicate Results",
JOptionPane.DEFAULT_OPTION, JOptionPane.WARNING_MESSAGE, null, options, options[0]);
if (choice == 1) { // Delete All Corrupted
handleDeleteCorrupted(errors);
} else if (choice == 2) { // Manage Duplicates
handleManageDuplicates(duplicates);
}
});
} catch (IOException e) {
SwingUtilities.invokeLater(() -> JOptionPane.showMessageDialog(this, "Error during verification: " + e.getMessage()));
SwingUtilities.invokeLater(() -> JOptionPane.showMessageDialog(this, "Error: " + e.getMessage()));
}
}).start();
}
private void handleDeleteCorrupted(List<IModelValidator.ValidationResult> errors) {
int confirm = JOptionPane.showConfirmDialog(this, "Delete " + errors.size() + " corrupted files?", "Confirm", JOptionPane.YES_NO_OPTION);
if (confirm == JOptionPane.YES_OPTION) {
int deleted = 0;
for (IModelValidator.ValidationResult err : errors) {
File f = new File(err.filePath);
if (f.exists() && f.delete()) {
hashRegistry.unregister(f);
deleted++;
}
}
JOptionPane.showMessageDialog(this, "Deleted " + deleted + " files.");
}
}
private void handleManageDuplicates(Map<String, List<Path>> duplicates) {
long totalSavings = 0;
for (List<Path> paths : duplicates.values()) {
for (int i = 1; i < paths.size(); i++) totalSavings += paths.get(i).toFile().length();
}
int confirm = JOptionPane.showConfirmDialog(this,
"Keep the first file of each set and delete others?\nEstimated savings: " + formatSize(totalSavings),
"Storage Optimizer", JOptionPane.YES_NO_OPTION);
if (confirm == JOptionPane.YES_OPTION) {
int deletedCount = 0;
for (List<Path> paths : duplicates.values()) {
for (int i = 1; i < paths.size(); i++) {
File f = paths.get(i).toFile();
if (f.exists() && f.delete()) {
hashRegistry.unregister(f);
deletedCount++;
}
}
}
JOptionPane.showMessageDialog(this, "Freed " + formatSize(totalSavings) + ".");
}
}
private void showHelpDialog() {
String helpText = "<html><body style='width: 600px; padding: 10px;'>" +
"<h1 style='color: #ff69b4;'>ComfyUI Model Downloader Help</h1>" +
@@ -754,12 +826,19 @@ public class Main extends JFrame {
private void fetchSizeInBackground(ModelInfo info, int rowIndex) {
new Thread(() -> {
long size = searchService.getRemoteSize(info.getUrl());
String s = searchService.formatSize(size);
String s = formatSize(size);
info.setSize(s);
SwingUtilities.invokeLater(() -> tableModel.setValueAt(s, rowIndex, 3));
}).start();
}
private String formatSize(long bytes) {
if (bytes <= 0) return "Unknown";
if (bytes < 1024 * 1024) return (bytes / 1024) + " KB";
if (bytes < 1024 * 1024 * 1024) return String.format("%.1f MB", bytes / (1024.0 * 1024.0));
return String.format("%.2f GB", bytes / (1024.0 * 1024.0 * 1024.0));
}
public static void main(String[] args) {
// Enable modern window decorations BEFORE anything else
com.formdev.flatlaf.FlatLaf.setUseNativeWindowDecorations(true);
src/main/java/de/tki/comfymodels/service/IModelValidator.java
+1
View File
@@ -4,6 +4,7 @@ import java.io.File;
public interface IModelValidator {
ValidationResult validateFile(File file);
String calculateHash(File file);
class ValidationResult {
public final boolean ok;
src/main/java/de/tki/comfymodels/service/impl/ModelHashRegistry.java
+113
View File
@@ -0,0 +1,113 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.service.IModelValidator;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import jakarta.annotation.PostConstruct;
import java.io.File;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
@Service
public class ModelHashRegistry {
private final String HASH_FILE = "model_hashes.json";
private Map<String, CacheEntry> cache = new HashMap<>();
private Map<String, String> hashToPath = new HashMap<>();
@Autowired
private ConfigService configService;
@Autowired
private IModelValidator validator;
private static class CacheEntry {
String hash;
long size;
long lastModified;
CacheEntry(String hash, long size, long lastModified) {
this.hash = hash;
this.size = size;
this.lastModified = lastModified;
}
}
@PostConstruct
public void load() {
try {
File file = configService.getFileInAppData(HASH_FILE);
if (file.exists()) {
String content = Files.readString(file.toPath(), StandardCharsets.UTF_8);
JSONObject json = new JSONObject(content);
for (String path : json.keySet()) {
JSONObject entry = json.getJSONObject(path);
String hash = entry.getString("hash");
long size = entry.optLong("size", 0);
long lm = entry.optLong("lastModified", 0);
if (new File(path).exists()) {
cache.put(path, new CacheEntry(hash, size, lm));
hashToPath.put(hash, path);
}
}
}
} catch (Exception e) {
System.err.println("Error loading hash registry: " + e.getMessage());
}
}
public synchronized void save() {
try {
JSONObject json = new JSONObject();
for (Map.Entry<String, CacheEntry> e : cache.entrySet()) {
JSONObject obj = new JSONObject();
obj.put("hash", e.getValue().hash);
obj.put("size", e.getValue().size);
obj.put("lastModified", e.getValue().lastModified);
json.put(e.getKey(), obj);
}
Files.writeString(configService.getFileInAppData(HASH_FILE).toPath(), json.toString(2), StandardCharsets.UTF_8);
} catch (Exception e) {
System.err.println("Error saving hash registry: " + e.getMessage());
}
}
public synchronized String getOrCalculateHash(File file) {
String path = file.getAbsolutePath();
long currentSize = file.length();
long currentLm = file.lastModified();
if (cache.containsKey(path)) {
CacheEntry entry = cache.get(path);
if (entry.size == currentSize && entry.lastModified == currentLm) {
return entry.hash;
}
}
String hash = validator.calculateHash(file);
if (hash != null) {
cache.put(path, new CacheEntry(hash, currentSize, currentLm));
hashToPath.put(hash, path);
save();
}
return hash;
}
public Optional<String> findPathByHash(String hash) {
return Optional.ofNullable(hashToPath.get(hash));
}
public void unregister(File file) {
String path = file.getAbsolutePath();
CacheEntry entry = cache.remove(path);
if (entry != null) {
hashToPath.remove(entry.hash);
save();
}
}
}
src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java
+41
View File
@@ -83,6 +83,12 @@ public class ModelSearchService implements IModelSearchService {
}
}
@Autowired
private de.tki.comfymodels.service.impl.ModelHashRegistry hashRegistry;
@Autowired
private de.tki.comfymodels.service.IModelValidator modelValidator;
private void performSearch(ModelInfo info, int index, String fileName, String workflowContext,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
@@ -93,6 +99,18 @@ public class ModelSearchService implements IModelSearchService {
if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return;
}
// Priority 2: Civitai Hash Lookup (if file exists locally but we want the source URL/Metadata)
// This is useful for re-downloading or finding updates
String modelsPath = configService.getModelsPath();
java.io.File localFile = new java.io.File(modelsPath, (info.getSave_path() != null ? info.getSave_path() : (info.getType() != null ? info.getType() : "checkpoints")) + java.io.File.separator + info.getName());
if (localFile.exists()) {
onStatusUpdate.accept(index, "🔍 Hashing for Civitai...");
String hash = hashRegistry.getOrCalculateHash(localFile);
if (hash != null) {
if (fetchCivitaiUrlByHash(info, index, hash, onStatusUpdate, onModelFound)) return;
}
}
onStatusUpdate.accept(index, "✨ Gemini Scouting...");
String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext);
if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) {
@@ -275,6 +293,29 @@ public class ModelSearchService implements IModelSearchService {
return false;
}
private boolean fetchCivitaiUrlByHash(ModelInfo info, int rowIndex, String hash,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
try {
String searchUrl = "https://civitai.com/api/v1/model-versions/by-hash/" + hash;
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
JSONObject version = new JSONObject(response.body());
JSONArray files = version.optJSONArray("files");
if (files != null) {
for (int j = 0; j < files.length(); j++) {
JSONObject file = files.getJSONObject(j);
// We check if the name matches or if it's the primary file for this version
if (file.optString("name").equalsIgnoreCase(info.getName()) || file.optBoolean("primary", false)) {
return validateAndSetUrl(info, rowIndex, file.optString("downloadUrl"), "🛡️ Civitai (Hash Match)", onStatusUpdate, onModelFound);
}
}
}
}
} catch (Exception e) {}
return false;
}
private boolean fetchCivitaiUrl(ModelInfo info, int rowIndex, String query,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
src/main/java/de/tki/comfymodels/service/impl/ModelValidator.java
+23
View File
@@ -45,6 +45,29 @@ public class ModelValidator implements IModelValidator {
return new ValidationResult(true, "OK", path);
}
@Override
public String calculateHash(File file) {
if (!file.exists() || !file.isFile()) return null;
try (java.io.InputStream is = new java.io.BufferedInputStream(new java.io.FileInputStream(file), 1024 * 1024)) {
java.security.MessageDigest digest = java.security.MessageDigest.getInstance("SHA-256");
byte[] buffer = new byte[1024 * 1024];
int read;
while ((read = is.read(buffer)) != -1) {
digest.update(buffer, 0, read);
}
byte[] hash = digest.digest();
StringBuilder hexString = new StringBuilder();
for (byte b : hash) {
String hex = Integer.toHexString(0xff & b);
if (hex.length() == 1) hexString.append('0');
hexString.append(hex);
}
return hexString.toString();
} catch (Exception e) {
return null;
}
}
private ValidationResult validateSafetensors(File file) {
try (RandomAccessFile raf = new RandomAccessFile(file, "r")) {
if (raf.length() < 8) return new ValidationResult(false, "File too small to have a header", file.getAbsolutePath());
src/test/java/de/tki/comfymodels/service/ModelHashRegistryTest.java
+119
View File
@@ -0,0 +1,119 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.service.impl.ConfigService;
import de.tki.comfymodels.service.impl.ModelHashRegistry;
import de.tki.comfymodels.service.impl.ModelValidator;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import java.io.File;
import java.lang.reflect.Field;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Optional;
import static org.junit.jupiter.api.Assertions.*;
public class ModelHashRegistryTest {
@TempDir
Path tempDir;
private ConfigService configService;
private IModelValidator validator;
private ModelHashRegistry registry;
private String lastCalculatedHash = "hash123";
private int calculationCount = 0;
@BeforeEach
public void setup() throws Exception {
configService = new ConfigService(null) {
@Override
public File getFileInAppData(String filename) {
return tempDir.resolve(filename).toFile();
}
};
validator = new IModelValidator() {
@Override
public ValidationResult validateFile(File file) { return new ValidationResult(true, "OK", file.getAbsolutePath()); }
@Override
public String calculateHash(File file) {
calculationCount++;
return lastCalculatedHash;
}
};
registry = new ModelHashRegistry();
// Use reflection to set private fields since we aren't using @InjectMocks
setField(registry, "configService", configService);
setField(registry, "validator", validator);
registry.load();
}
private void setField(Object obj, String fieldName, Object value) throws Exception {
Field field = obj.getClass().getDeclaredField(fieldName);
field.setAccessible(true);
field.set(obj, value);
}
@Test
public void testGetOrCalculateHash() throws Exception {
Path file = tempDir.resolve("model.safetensors");
Files.writeString(file, "content");
File modelFile = file.toFile();
lastCalculatedHash = "hash123";
calculationCount = 0;
// First call: should calculate
String hash1 = registry.getOrCalculateHash(modelFile);
assertEquals("hash123", hash1);
assertEquals(1, calculationCount);
// Second call: should use cache because size and timestamp are same
String hash2 = registry.getOrCalculateHash(modelFile);
assertEquals("hash123", hash2);
assertEquals(1, calculationCount);
// Change content -> size stays same, but timestamp changes (usually)
// To be sure for the test, we wait a bit or just assume size change
Files.writeString(file, "content has changed now");
File updatedFile = file.toFile();
lastCalculatedHash = "new_hash";
String hash3 = registry.getOrCalculateHash(updatedFile);
assertEquals("new_hash", hash3);
assertEquals(2, calculationCount); // Re-calculation triggered!
}
@Test
public void testFindPathByHash() throws Exception {
Path file = tempDir.resolve("duplicate.safetensors");
Files.writeString(file, "duplicate content");
File modelFile = file.toFile();
lastCalculatedHash = "duplicate_hash";
registry.getOrCalculateHash(modelFile);
Optional<String> foundPath = registry.findPathByHash("duplicate_hash");
assertTrue(foundPath.isPresent());
assertEquals(modelFile.getAbsolutePath(), foundPath.get());
}
@Test
public void testUnregister() throws Exception {
Path file = tempDir.resolve("to_be_deleted.safetensors");
Files.writeString(file, "soon gone");
File modelFile = file.toFile();
lastCalculatedHash = "temp_hash";
registry.getOrCalculateHash(modelFile);
assertTrue(registry.findPathByHash("temp_hash").isPresent());
registry.unregister(modelFile);
assertFalse(registry.findPathByHash("temp_hash").isPresent());
}
}
src/test/java/de/tki/comfymodels/service/ModelHashTest.java
+47
View File
@@ -0,0 +1,47 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.service.impl.ModelValidator;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import static org.junit.jupiter.api.Assertions.*;
public class ModelHashTest {
@TempDir
Path tempDir;
@Test
public void testCalculateSHA256() throws Exception {
ModelValidator validator = new ModelValidator();
Path testFile = tempDir.resolve("test_model.bin");
byte[] data = "ComfyUI Model Data for Hashing Test".getBytes(java.nio.charset.StandardCharsets.UTF_8);
Files.write(testFile, data);
String hash = validator.calculateHash(testFile.toFile());
assertNotNull(hash);
assertEquals(64, hash.length());
// echo -n "ComfyUI Model Data for Hashing Test" | sha256sum
String expectedHash = "5022ad6b25eac3358b18aa8ef8ee79d06833c43bed545bed1a5d45057d3054e2";
assertEquals(expectedHash, hash);
}
@Test
public void testEmptyFileHash() throws Exception {
ModelValidator validator = new ModelValidator();
Path emptyFile = tempDir.resolve("empty.bin");
Files.createFile(emptyFile);
String hash = validator.calculateHash(emptyFile.toFile());
// SHA-256 for empty string
String expectedEmptyHash = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
assertEquals(expectedEmptyHash, hash);
}
}