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

refactor: implement dependency injection for encryption and enhance model analysis

Browse Source
This commit is contained in:
Thomas
2026-04-25 15:06:56 +02:00
parent 7b9743e490
commit fd1fcaa221
25 changed files with 638 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/java/de/tki/comfymodels/BatchVerification.java
-88
View File
@@ -1,88 +0,0 @@
package de.tki.comfymodels;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.util.List;
import java.util.zip.InflaterInputStream;
public class BatchVerification {
public static void main(String[] args) throws Exception {
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
Path root = Paths.get("C:/pinokio/api/comfy.git/workflows/ComfyUI_examples");
System.out.println("=== BATCH VERIFICATION START ===");
Files.walk(root)
.filter(p -> p.toString().endsWith(".json") || p.toString().endsWith(".png"))
.limit(30) // Stichprobe von 30 verschiedenen Workflows
.forEach(path -> {
try {
System.out.println("\nFILE: " + path.getFileName());
String content;
if (path.toString().endsWith(".png")) {
content = extractWorkflowFromPng(path.toFile());
} else {
content = Files.readString(path);
}
if (content == null) {
System.out.println(" [!] No workflow found in PNG");
return;
}
List<ModelInfo> results = analyzer.analyze(content, path.getFileName().toString());
if (results.isEmpty()) {
System.out.println(" [?] No models detected");
} else {
for (ModelInfo info : results) {
System.out.printf(" [%s] %s | URL: %s | Source: %s\n",
info.getType(), info.getName(),
info.getUrl().equals("MISSING") ? "Missing" : "FOUND",
info.getPopularity());
}
}
} catch (Exception e) {
System.out.println(" [ERROR] " + e.getMessage());
}
});
System.out.println("\n=== BATCH VERIFICATION END ===");
}
private static String extractWorkflowFromPng(File file) {
try (DataInputStream dis = new DataInputStream(new FileInputStream(file))) {
if (dis.readLong() != 0x89504E470D0A1A0AL) return null;
while (true) {
int length = dis.readInt();
byte[] typeBytes = new byte[4];
dis.readFully(typeBytes);
String type = new String(typeBytes, StandardCharsets.US_ASCII);
if ("tEXt".equals(type) || "iTXt".equals(type) || "zTXt".equals(type)) {
byte[] data = new byte[length];
dis.readFully(data);
int nullPos = 0; while (nullPos < data.length && data[nullPos] != 0) nullPos++;
String key = new String(data, 0, nullPos, StandardCharsets.UTF_8);
String value = null;
if ("tEXt".equals(type)) value = new String(data, nullPos + 1, data.length - (nullPos + 1), StandardCharsets.UTF_8);
else if ("zTXt".equals(type)) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, nullPos + 2, data.length - (nullPos + 2)))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
} else if ("iTXt".equals(type)) {
int currentPos = nullPos + 3; int nullCount = 0;
while (currentPos < data.length && nullCount < 2) { if (data[currentPos] == 0) nullCount++; currentPos++; }
if (currentPos < data.length) {
if (data[nullPos + 1] != 0) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, currentPos, data.length - currentPos))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
} else value = new String(data, currentPos, data.length - currentPos, StandardCharsets.UTF_8);
}
}
if ("workflow".equalsIgnoreCase(key) || "prompt".equalsIgnoreCase(key)) return value;
dis.readInt();
} else if ("IEND".equals(type)) break; else dis.skipBytes(length + 4);
}
} catch (Exception e) {}
return null;
}
}
src/main/java/de/tki/comfymodels/Main.java
+55 -22
View File
@@ -101,8 +101,10 @@ public class Main extends JFrame {
} else {
FlatLightLaf.setup();
}
FlatLaf.updateUI();
// Modern UI tweaks
UIManager.put("TitlePane.unifiedBackground", true);
UIManager.put("Button.arc", 8);
UIManager.put("Component.arc", 8);
UIManager.put("TextComponent.arc", 8);
@@ -255,6 +257,9 @@ public class Main extends JFrame {
setTitle("ComfyUIModel-Downloader");
setSize(1450, 1000);
setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE);
// Enable FlatLaf window decorations for this frame
getRootPane().putClientProperty("flatlaf.useWindowDecorations", true);
JPanel mainContainer = new JPanel(new BorderLayout(10, 10));
mainContainer.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10));
@@ -295,7 +300,7 @@ public class Main extends JFrame {
pathRow.add(pathButtons, BorderLayout.EAST);
JPanel apiRow = new JPanel(new BorderLayout());
apiRow.setBorder(BorderFactory.createTitledBorder("Gemini AI API Key"));
apiRow.setBorder(BorderFactory.createTitledBorder("Gemini AI API Key (Optional)"));
geminiKeyField = new JPasswordField();
JButton saveGeminiBtn = new JButton("Save & Discover");
saveGeminiBtn.addActionListener(e -> {
@@ -306,7 +311,7 @@ public class Main extends JFrame {
apiRow.add(saveGeminiBtn, BorderLayout.EAST);
JPanel hfRow = new JPanel(new BorderLayout());
hfRow.setBorder(BorderFactory.createTitledBorder("Hugging Face Access Token"));
hfRow.setBorder(BorderFactory.createTitledBorder("Hugging Face Access Token (Optional)"));
hfTokenField = new JPasswordField();
JButton saveHfBtn = new JButton("Save HF Token");
saveHfBtn.addActionListener(e -> {
@@ -330,6 +335,7 @@ public class Main extends JFrame {
shutdownCheck = new JCheckBox("Shutdown after Queue");
shutdownCheck.addActionListener(e -> configService.setShutdownAfterDownloadEnabled(shutdownCheck.isSelected()));
// Most important feature (irony): Dark mode for Reddit... or at least for this tool.
darkCheck = new JCheckBox("Dark Mode");
darkCheck.addActionListener(e -> {
boolean isDark = darkCheck.isSelected();
@@ -411,7 +417,7 @@ public class Main extends JFrame {
JTable modelTable = new JTable(tableModel);
modelTable.putClientProperty("terminateEditOnFocusLost", Boolean.TRUE);
// Sofortiger Abbruch des Editors bei Klick, um Wert zu übernehmen
// Immediate termination of the editor on click to apply the value
modelTable.addMouseListener(new MouseAdapter() {
@Override
public void mouseReleased(MouseEvent e) {
@@ -566,13 +572,25 @@ public class Main extends JFrame {
String status = info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good";
boolean isSelected = true;
Path localPath = Paths.get(baseModelsPath, info.getSave_path() != null ? info.getSave_path() : type, info.getName());
if (Files.exists(localPath)) {
boolean existsLocally = Files.exists(localPath);
if (existsLocally) {
status = "✅ Already exists";
isSelected = false; // Deselect if already present
try {
long localSizeBytes = Files.size(localPath);
info.setSize(searchService.formatSize(localSizeBytes));
} catch (IOException e) {
info.setSize("Local (Error)");
}
}
tableModel.addRow(new Object[]{isSelected, info.getType(), info.getName(), info.getSize(), info.getPopularity(), targetSubDir, info.getUrl(), status});
if (!info.getUrl().equals("MISSING")) fetchSizeInBackground(info, i);
// Only fetch remote size if NOT already existing locally
if (!existsLocally && !info.getUrl().equals("MISSING")) {
fetchSizeInBackground(info, i);
}
}
downloadButton.setEnabled(!modelsToDownload.isEmpty());
}
@@ -695,26 +713,38 @@ public class Main extends JFrame {
}
private void showHelpDialog() {
String helpText = "<html><body style='width: 500px; padding: 10px;'>" +
"<h1>ComfyUIModel-Downloader Help</h1>" +
"<h3>Core Features</h3>" +
String helpText = "<html><body style='width: 600px; padding: 10px;'>" +
"<h1 style='color: #ff69b4;'>ComfyUI Model Downloader Help</h1>" +
"<p>This tool automates the process of finding and downloading models required for your ComfyUI workflows.</p>" +
"<h3>🚀 Core Features</h3>" +
"<ul>" +
" <li><b>Workflow Analysis:</b> Drag and drop ComfyUI JSON or PNG files to extract required models.</li>" +
" <li><b>Deep Search:</b> Uses AI and multi-platform searching (Hugging Face, Civitai) to find missing model URLs.</li>" +
" <li><b>Model List Integration:</b> Match models against your own known-good local lists.</li>" +
" <li><b>Queue Management:</b> Download multiple models simultaneously with pause/resume support.</li>" +
" <li><b>Smart Extraction:</b> Drag and drop <b>JSON</b> or <b>PNG</b> workflow files. The tool extracts models even from PNG metadata.</li>" +
" <li><b>Deep Search:</b> Automatically finds URLs on <b>Hugging Face</b> (Official & Community) and <b>Civitai</b>.</li>" +
" <li><b>AI Scouting:</b> Uses Gemini AI to understand the workflow context and find the exact right model versions.</li>" +
" <li><b>Auto-Organization:</b> Automatically sorts models into correct subfolders (<code>checkpoints</code>, <code>loras</code>, <code>vae</code>, <code>controlnet</code>, etc.).</li>" +
" <li><b>Existing File Detection:</b> Detects models you already have, calculates their local size, and prevents redundant downloads.</li>" +
"</ul>" +
"<h3>API Keys & Security</h3>" +
"<h3>🛡️ Quality & Security</h3>" +
"<ul>" +
" <li><b>Gemini API Key:</b> Required for the 'AI Scouting' feature. Gemini analyzes your workflow context " +
" to predict the best repositories for obscure models.</li>" +
" <li><b>Hugging Face Token:</b> Recommended to avoid rate limits and to access gated repositories " +
" (like some FLUX or StabilityAI models).</li>" +
" <li><b>Local Storage (Vault):</b> All keys and settings are stored <b>locally</b> on your machine " +
" in an encrypted file (<code>settings.vault</code>). They are never sent to any server except " +
" the official Google/Hugging Face APIs during requests.</li>" +
" <li><b>Password Protection:</b> Your vault is protected by the password you choose at startup. " +
" Without this password, the API keys cannot be decrypted.</li>" +
" <li><b>Model Verification:</b> Use the <b>Verify</b> button to scan your local models for corrupted files (e.g., interrupted downloads or broken Safetensors).</li>" +
" <li><b>Encrypted Vault:</b> Your optional API keys are stored in <code>settings.vault</code> using <b>AES-256</b> encryption, protected by your master password.</li>" +
" <li><b>Privacy First:</b> No data is ever sent to third-party servers except for official API requests to Google, Hugging Face, or Civitai.</li>" +
"</ul>" +
"<h3>⚙️ Advanced Settings</h3>" +
"<ul>" +
" <li><b>Background Mode:</b> Close the window to keep the downloader running in the <b>System Tray</b>.</li>" +
" <li><b>Auto-Shutdown:</b> Automatically shut down your PC after the download queue is finished.</li>" +
" <li><b>Dark Mode:</b> Modern dark interface for better workflow integration.</li>" +
" <li><b>Model List Import:</b> Import your own JSON/CSV lists to prioritize your trusted sources.</li>" +
"</ul>" +
"<h3>🔑 API Keys (Optional)</h3>" +
"<ul>" +
" <li><b>Gemini Key:</b> Recommended for the best 'AI Scouting' results and repository discovery.</li>" +
" <li><b>HF Token:</b> Required to download <b>Gated Models</b> (like FLUX.1-dev or certain SD3 versions).</li>" +
"</ul>" +
"</body></html>";
@@ -731,6 +761,9 @@ public class Main extends JFrame {
}
public static void main(String[] args) {
// Enable modern window decorations BEFORE anything else
com.formdev.flatlaf.FlatLaf.setUseNativeWindowDecorations(true);
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class);
Main main = context.getBean(Main.class);
main.launch(args);
src/main/java/de/tki/comfymodels/PngTest.java
-49
View File
@@ -1,49 +0,0 @@
package de.tki.comfymodels;
import java.io.*;
import java.nio.charset.StandardCharsets;
public class PngTest {
public static void main(String[] args) {
File file = new File("flux_canny_model_example.png");
String workflow = extractWorkflowFromPng(file);
if (workflow != null) {
System.out.println("SUCCESS! Extracted workflow length: " + workflow.length());
System.out.println("Snippet: " + workflow.substring(0, Math.min(200, workflow.length())));
} else {
System.out.println("FAILED! No workflow found.");
}
}
private static String extractWorkflowFromPng(File file) {
try (DataInputStream dis = new DataInputStream(new FileInputStream(file))) {
if (dis.readLong() != 0x89504E470D0A1A0AL) return null;
while (true) {
int length = dis.readInt();
byte[] typeBytes = new byte[4];
dis.readFully(typeBytes);
String type = new String(typeBytes, StandardCharsets.US_ASCII);
if ("tEXt".equals(type) || "iTXt".equals(type) || "zTXt".equals(type)) {
byte[] data = new byte[length];
dis.readFully(data);
String content = new String(data, StandardCharsets.UTF_8);
if (content.contains("workflow") || content.contains("prompt")) {
for (int i = 0; i < data.length; i++) {
if (data[i] == 0) {
return new String(data, i + 1, data.length - i - 1, StandardCharsets.UTF_8).trim();
}
}
}
dis.readInt(); // CRC
} else if ("IEND".equals(type)) {
break;
} else {
dis.skipBytes(length + 4);
}
}
} catch (Exception e) {
e.printStackTrace();
}
return null;
}
}
src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java
+55 -74
View File
@@ -14,7 +14,7 @@ import java.util.regex.Pattern;
@Service
public class ComfyModelAnalyzer implements IModelAnalyzer {
@Autowired(required = false)
@Autowired
private LocalAIService aiService;
@Autowired
@@ -23,14 +23,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
@Autowired(required = false)
private GeminiAIService geminiService;
private static final class ExpertPatterns {
static final Pattern FILE = Pattern.compile("([a-zA-Z0-9_\\-\\.\\/]+\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_VAE = Pattern.compile("(vae|tokenizer|autoencoder)", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_CLIP = Pattern.compile("(clip|t5|text_encoder|llama|gemma|mistral|t5xxl|clip_l|clip_g|qwen)", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_UNET = Pattern.compile("(unet|diffusion|model|transformer|base|transformer|upscale|esrgan|z_image|lumina|flux|dit|mochi|ltx|cosmos|hunyuan|hy|wan|hidream|aura|svd)", Pattern.CASE_INSENSITIVE);
static final Pattern MD_LINK = Pattern.compile("\\[([^\\]]+)\\]\\((https?://[^\\)]+)\\)", Pattern.CASE_INSENSITIVE);
static final Pattern NAKED_URL = Pattern.compile("(https?://[a-zA-Z0-9\\-\\.\\/\\?\\=\\&\\%]+(?:\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))(?:\\?[^\\s\\\"\\']*)?)", Pattern.CASE_INSENSITIVE);
}
private final Pattern FILE_PATTERN = Pattern.compile("([a-zA-Z0-9_\\-\\.\\/]+\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))", Pattern.CASE_INSENSITIVE);
private final Pattern COMP_VAE_PATTERN = Pattern.compile("(vae|tokenizer|autoencoder|encoder|decoder)", Pattern.CASE_INSENSITIVE);
private final Pattern COMP_CLIP_PATTERN = Pattern.compile("(clip|text.encoder|t5|llama|gemma|mistral|qwen|bert|vit|embed)", Pattern.CASE_INSENSITIVE);
private final Pattern COMP_UNET_PATTERN = Pattern.compile("(unet|diffusion|transformer|dit|model|base|upscale|esrgan|resnet|sampling)", Pattern.CASE_INSENSITIVE);
private final Pattern MD_LINK_PATTERN = Pattern.compile("\\[([^\\]]+)\\]\\((https?://[^\\)]+)\\)", Pattern.CASE_INSENSITIVE);
private final Pattern NAKED_URL_PATTERN = Pattern.compile("(https?://[a-zA-Z0-9\\-\\.\\/\\?\\=\\&\\%]+(?:\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))(?:\\?[^\\s\\\"\\']*)?)", Pattern.CASE_INSENSITIVE);
@Override
public List<ModelInfo> analyze(String jsonText, String fileName) {
@@ -40,13 +38,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
JSONObject jo = parseJson(jsonText);
if (jo == null) return results;
// NEW: Globale Kontext-Erkennung vor der Analyse
// Global context detection based on frequency/filename rather than hardcoded list
String globalContext = extractGlobalContext(jo, fileName);
findModelsMetadata(jo, results);
scanForModelFiles(jo, null, results);
// Wende globalen Kontext auf alle Ergebnisse an
applyGlobalContext(results, globalContext);
for (ModelInfo info : results) {
@@ -59,13 +56,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
info.setSave_path(match.getSave_path());
info.setPopularity("📂 MODEL LIST MATCH");
} else {
// Fallback to AI prediction if not in list
if (aiService == null) aiService = new LocalAIService();
LocalAIService.Prediction prediction = aiService.predictProvider(info.getName());
info.setPopularity(prediction.getLabel());
if (aiService != null) {
LocalAIService.Prediction prediction = aiService.predictProvider(info.getName());
info.setPopularity(prediction.getLabel());
}
// Falls die Popularität noch generisch ist, füge den Kontext hinzu
if (globalContext != null && !globalContext.isEmpty() && info.getPopularity().contains("Community")) {
if (globalContext != null && !globalContext.isEmpty() && info.getPopularity() != null && info.getPopularity().contains("Community")) {
info.setPopularity("✨ Context: " + globalContext + " (" + info.getPopularity() + ")");
}
}
@@ -75,26 +71,22 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
}
private String extractGlobalContext(JSONObject jo, String fileName) {
Map<String, Integer> hints = new HashMap<>();
// 1. Aus Dateiname (höchste Prio)
// Generic approach: extract the most significant prefix from filename or content
if (fileName != null && !fileName.toLowerCase().startsWith("workflow")) {
String base = fileName.split("[_\\- ]")[0].toLowerCase();
if (base.length() >= 3) return base;
}
// 2. Scan den ganzen JSON nach signifikanten Schlüsselwörtern
// Scan for frequently occurring specific terms in class_types (that are not standard keywords)
Map<String, Integer> counts = new HashMap<>();
String fullText = jo.toString().toLowerCase();
String[] keywords = {"hidream", "flux", "auraflow", "sd3", "hunyuan", "mochi", "ltxv", "cosmos", "wan2.1"};
for (String kw : keywords) {
if (fullText.contains(kw)) hints.put(kw, hints.getOrDefault(kw, 0) + 5);
// Look for common architecture names as hints
String[] archHints = {"flux", "sdxl", "sd15", "sd3", "hunyuan", "mochi", "ltx", "cosmos", "wan", "hidream", "aura", "lumina"};
for (String hint : archHints) {
if (fullText.contains(hint)) counts.put(hint, counts.getOrDefault(hint, 0) + 1);
}
// 3. Scan Node-Typen
String nodesText = fullText;
if (nodesText.contains("hidream")) hints.put("hidream", hints.getOrDefault("hidream", 0) + 10);
return hints.entrySet().stream()
return counts.entrySet().stream()
.max(Map.Entry.comparingByValue())
.map(Map.Entry::getKey)
.orElse(null);
@@ -103,10 +95,10 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
private void applyGlobalContext(List<ModelInfo> res, String ctx) {
if (ctx == null) return;
for (ModelInfo i : res) {
// Wenn der Dateiname den Kontext enthält oder der Dateiname sehr generisch ist (wie ae.safetensors)
String n = i.getName().toLowerCase();
if (n.contains(ctx) || n.equals("ae.safetensors") || n.contains("clip_l") || n.contains("t5xxl")) {
if (i.getPopularity() == null || i.getPopularity().isEmpty()) {
// If filename matches context or is a generic model file
if (n.contains(ctx) || n.startsWith("ae.") || n.contains("clip_") || n.contains("t5") || n.contains("diffusion")) {
if (i.getPopularity() == null || i.getPopularity().isEmpty() || i.getPopularity().contains("Community")) {
i.setPopularity("✨ Context: " + ctx);
}
}
@@ -133,7 +125,6 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
if (m != null) {
String n = m.optString("name", m.optString("filename"));
if (n != null && !n.isEmpty()) {
// Strikte Übernahme des Pfades aus den Metadaten
String rawPath = m.optString("directory", m.optString("type", "checkpoints"));
ModelInfo info = new ModelInfo(rawPath, n, m.optString("url", "MISSING"));
info.setSave_path(rawPath);
@@ -153,6 +144,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
if (obj instanceof JSONObject) {
JSONObject jo = (JSONObject) obj;
String node = jo.optString("class_type", jo.optString("type", null));
// Ignore information from markdown notes as requested
if (node != null && (node.equalsIgnoreCase("MarkdownNote") || node.equalsIgnoreCase("Note"))) {
return;
}
String curCtx = node != null ? (inferTypeFromNode(node) != null ? inferTypeFromNode(node) : ctx) : ctx;
if (jo.has("widgets_values")) {
JSONArray ws = jo.optJSONArray("widgets_values");
@@ -174,24 +171,19 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
if (v == null || v.trim().isEmpty()) return;
String vLow = v.toLowerCase().trim();
// Ignoriere Standardwerte, die keine Dateinamen sind
if (vLow.equals("none") || vLow.equals("default") || vLow.equals("undefined") || vLow.length() < 4) return;
// 1. Suche nach Markdown Links: [Name](URL)
Matcher mdMatcher = ExpertPatterns.MD_LINK.matcher(v);
Matcher mdMatcher = MD_LINK_PATTERN.matcher(v);
while (mdMatcher.find()) {
String name = mdMatcher.group(1);
String url = fixUrl(mdMatcher.group(2));
if (name.contains(".")) { // Nur wenn der Linktext wie ein Dateiname aussieht
if (name.contains(".")) {
String type = hint != null ? inferTypeFromKey(hint) : null;
if ((type == null || type.equals("checkpoints")) && url.contains("/resolve/main/split_files/")) {
String sub = url.substring(url.indexOf("/split_files/") + 13);
if (sub.contains("/")) type = sub.substring(0, sub.indexOf("/"));
} else if ((type == null || type.equals("checkpoints")) && url.contains("/resolve/main/")) {
if ((type == null || type.equals("checkpoints")) && url.contains("/resolve/main/")) {
String sub = url.substring(url.indexOf("/resolve/main/") + 14);
if (sub.contains("/")) {
String folder = sub.substring(0, sub.indexOf("/")).toLowerCase();
if (folder.equals("loras") || folder.equals("vae") || folder.equals("checkpoints") || folder.equals("text_encoders") || folder.equals("diffusion_models") || folder.equals("controlnet")) {
if (folder.matches("(loras|vae|checkpoints|text_encoders|diffusion_models|controlnet|embeddings|upscale_models)")) {
type = folder;
}
}
@@ -201,8 +193,7 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
}
}
// 2. Suche nach nackten URLs
Matcher nakedMatcher = ExpertPatterns.NAKED_URL.matcher(v);
Matcher nakedMatcher = NAKED_URL_PATTERN.matcher(v);
while (nakedMatcher.find()) {
String url = fixUrl(nakedMatcher.group(1));
String name = url;
@@ -216,33 +207,26 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
}
}
Matcher m = ExpertPatterns.FILE.matcher(v);
Matcher m = FILE_PATTERN.matcher(v);
while (m.find()) {
String match = m.group(1);
// Ignoriere offensichtliche URLs (wir wollen nur Dateinamen)
if (match.startsWith("http") || match.startsWith("//")) continue;
// Extrahiere nur den reinen Dateinamen, falls ein Pfad gematcht wurde
String f = match;
if (f.contains("/")) f = f.substring(f.lastIndexOf("/") + 1);
if (f.contains("\\")) f = f.substring(f.lastIndexOf("\\") + 1);
if (f.length() > 3) {
// Erneuter Check nach Pfad-Extraktion
String fLow = f.toLowerCase();
if (fLow.equals("none.safetensors") || fLow.equals("default.safetensors")) continue;
String type = hint != null ? inferTypeFromKey(hint) : null;
if (type == null) type = ctx;
// Verfeinerung basierend auf Dateinamen (Höchste Priorität für Spezialmodelle)
if (fLow.contains("qwen-image-lightning")) type = "unet";
else if (fLow.contains("vae_approx") || fLow.contains("vae-approx")) type = "vae_approx";
else if (ExpertPatterns.COMP_VAE.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "vae";
else if (fLow.contains("clip_vision")) type = "clip_vision";
else if (ExpertPatterns.COMP_CLIP.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "clip";
else if (ExpertPatterns.COMP_UNET.matcher(fLow).find() && (type == null || type.equals("checkpoints") || (type.equals("clip") && !fLow.contains("clip_l") && !fLow.contains("clip_g")))) type = "unet";
// Generic refinement based on structural keywords
if (COMP_VAE_PATTERN.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "vae";
else if (COMP_CLIP_PATTERN.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "clip";
else if (COMP_UNET_PATTERN.matcher(fLow).find() && (type == null || type.equals("checkpoints") || (type.equals("clip") && !fLow.contains("clip_l") && !fLow.contains("clip_g")))) type = "unet";
addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", f, "MISSING"));
}
@@ -252,13 +236,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
private String inferTypeFromNode(String nt) {
nt = nt.toLowerCase();
if (nt.contains("lora")) return "loras";
if (nt.contains("vae_approx") || nt.contains("vae-approx")) return "vae_approx";
if (nt.contains("vae")) return "vae";
if (nt.contains("controlnet") || nt.contains("t2i_adapter")) return "controlnet";
if (nt.contains("upscale") || nt.contains("esrgan")) return "upscale_models";
if (nt.contains("latent_upscale")) return "latent_upscale_models";
if (nt.contains("clip_vision")) return "clip_vision";
if (nt.contains("clip") || nt.contains("text_encoder") || nt.contains("cliploader") || nt.contains("llama") || nt.contains("mistral") || nt.contains("t5") || nt.contains("qwen")) return "clip";
if (nt.contains("clip_vision") || nt.contains("clipvision")) return "clip_vision";
if (nt.contains("clip") || nt.contains("text_encoder") || nt.contains("tokenizer")) return "clip";
if (nt.contains("diffusion_model")) return "diffusion_models";
if (nt.contains("unet") || nt.contains("diffusion")) return "unet";
if (nt.contains("gligen")) return "gligen";
@@ -276,12 +259,11 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
private String inferTypeFromKey(String k) {
k = k.toLowerCase();
if (k.contains("lora")) return "loras";
if (k.contains("vae_approx") || k.contains("vae-approx")) return "vae_approx";
if (k.contains("vae")) return "vae";
if (k.contains("clip_vision")) return "clip_vision";
if (k.contains("clip") || k.contains("text_encoder") || k.contains("llama") || k.contains("mistral") || k.contains("t5") || k.contains("qwen")) return "clip";
if (k.contains("clip_vision") || k.contains("clipvision")) return "clip_vision";
if (k.contains("clip") || k.contains("text_encoder") || k.contains("tokenizer")) return "clip";
if (k.contains("diffusion_model")) return "diffusion_models";
if (k.contains("unet") || k.contains("diffusion") || k.contains("hidream") || k.contains("aura") || k.contains("mochi") || k.contains("svd") || k.contains("z_image") || k.contains("lumina") || k.contains("flux") || k.contains("dit")) return "unet";
if (k.contains("unet") || k.contains("diffusion")) return "unet";
if (k.contains("control") && k.contains("net")) return "controlnet";
if (k.contains("upscale") || k.contains("esrgan")) return "upscale_models";
if (k.contains("embedding")) return "embeddings";
@@ -297,18 +279,17 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
private void addModelInfo(List<ModelInfo> res, ModelInfo info) {
for (ModelInfo e : res) {
if (e.getName().equalsIgnoreCase(info.getName())) {
if (!info.getUrl().equals("MISSING")) {
e.setUrl(info.getUrl());
// If the existing entry is generic "checkpoints" but we found a more specific type
if ("checkpoints".equals(e.getType()) && !"checkpoints".equals(info.getType())) {
e.setType(info.getType());
e.setSave_path(info.getSave_path());
} else if (e.getUrl().equals("MISSING")) {
boolean isEWeak = e.getType() == null || e.getType().equals("checkpoints") || e.getType().equals("clip");
boolean isInfoStrong = info.getType() != null && (info.getType().equals("unet") || info.getType().equals("checkpoints"));
if (isEWeak && isInfoStrong) {
e.setType(info.getType());
e.setSave_path(info.getSave_path());
}
if (info.getSave_path() != null) e.setSave_path(info.getSave_path());
}
// Always take URL if the current one is MISSING
if (e.getUrl().equals("MISSING") && !info.getUrl().equals("MISSING")) {
e.setUrl(info.getUrl());
if (info.getSave_path() != null) e.setSave_path(info.getSave_path());
if (info.getType() != null && !"checkpoints".equals(info.getType())) e.setType(info.getType());
}
return;
}
src/main/java/de/tki/comfymodels/service/impl/ConfigService.java
+14 -6
View File
@@ -1,6 +1,7 @@
package de.tki.comfymodels.service.impl;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.io.File;
@@ -9,11 +10,18 @@ import java.nio.file.Files;
@Service
public class ConfigService {
private static final String CONFIG_FILE = "app_settings.json";
private static final String VAULT_FILE = "settings.vault";
private final String CONFIG_FILE = "app_settings.json";
private final String VAULT_FILE = "settings.vault";
private JSONObject settings = new JSONObject();
private String masterPassword = null;
private final EncryptionUtils encryptionUtils;
@Autowired
public ConfigService(EncryptionUtils encryptionUtils) {
this.encryptionUtils = encryptionUtils;
}
/**
* Use the current working directory for application data.
*/
@@ -32,7 +40,7 @@ public class ConfigService {
if (vault.exists()) {
String encrypted = Files.readString(vault.toPath(), StandardCharsets.UTF_8);
try {
String decrypted = EncryptionUtils.decrypt(encrypted, password);
String decrypted = encryptionUtils.decrypt(encrypted, password);
JSONObject decryptedJson = new JSONObject(decrypted);
this.settings = decryptedJson;
this.masterPassword = password;
@@ -67,7 +75,7 @@ public class ConfigService {
return;
}
try {
String encrypted = EncryptionUtils.encrypt(settings.toString(), masterPassword);
String encrypted = encryptionUtils.encrypt(settings.toString(), masterPassword);
Files.writeString(getFileInAppData(VAULT_FILE).toPath(), encrypted, StandardCharsets.UTF_8);
System.out.println("Vault saved successfully to: " + getFileInAppData(VAULT_FILE).getAbsolutePath());
} catch (Exception e) {
@@ -91,7 +99,7 @@ public class ConfigService {
public boolean isShutdownAfterDownloadEnabled() { return settings.optBoolean("shutdown_after_download", false); }
public void setShutdownAfterDownloadEnabled(boolean enabled) { settings.put("shutdown_after_download", enabled); save(); }
public boolean isDarkMode() { return settings.optBoolean("dark_mode", true); }
public boolean isDarkMode() { return settings.optBoolean("dark_mode", false); }
public void setDarkMode(boolean enabled) { settings.put("dark_mode", enabled); save(); }
public void savePendingDownloads(String json) {
@@ -117,4 +125,4 @@ public class ConfigService {
public boolean isUnlocked() {
return masterPassword != null;
}
}
}
src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java
+1 -1
View File
@@ -261,7 +261,7 @@ public class DefaultDownloadManager implements IDownloadManager {
}
}
} catch (Exception e) {
// Falls es ein Timing-Problem ist (Flagge wird gesetzt während Exception fliegt)
// In case of a timing issue (flag set while exception is thrown)
try { Thread.sleep(50); } catch (InterruptedException ignored) {}
if (isStopped || !isSelected(index) || Thread.currentThread().isInterrupted()) {
src/main/java/de/tki/comfymodels/service/impl/EncryptionUtils.java
+9 -6
View File
@@ -1,5 +1,7 @@
package de.tki.comfymodels.service.impl;
import org.springframework.stereotype.Component;
import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
@@ -11,12 +13,13 @@ import java.security.SecureRandom;
import java.security.spec.KeySpec;
import java.util.Base64;
@Component
public class EncryptionUtils {
private static final int ITERATIONS = 65536;
private static final int KEY_LENGTH = 256;
private static final byte[] SALT = "ComfyUI-Vault-Salt-2026".getBytes(StandardCharsets.UTF_8);
private final int ITERATIONS = 65536;
private final int KEY_LENGTH = 256;
private final byte[] SALT = "ComfyUI-Vault-Salt-2026".getBytes(StandardCharsets.UTF_8);
public static String encrypt(String strToEncrypt, String password) throws Exception {
public String encrypt(String strToEncrypt, String password) throws Exception {
byte[] iv = new byte[16];
new SecureRandom().nextBytes(iv);
IvParameterSpec ivspec = new IvParameterSpec(iv);
@@ -37,7 +40,7 @@ public class EncryptionUtils {
return Base64.getEncoder().encodeToString(combined);
}
public static String decrypt(String strToDecrypt, String password) throws Exception {
public String decrypt(String strToDecrypt, String password) throws Exception {
if (strToDecrypt == null || strToDecrypt.trim().isEmpty()) {
throw new Exception("Encrypted data is empty");
}
@@ -62,4 +65,4 @@ public class EncryptionUtils {
cipher.init(Cipher.DECRYPT_MODE, secretKey, ivspec);
return new String(cipher.doFinal(encrypted), StandardCharsets.UTF_8);
}
}
}
src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java
+8 -8
View File
@@ -65,14 +65,14 @@ public class GeminiAIService {
String context = metadataContext != null ? metadataContext : "No context";
String shortContext = context.length() > 25000 ? context.substring(0, 25000) : context;
String prompt = "Websuche: Ermittle das offizielle Hugging Face Repository für die Datei '" + modelName + "'.\n\n" +
"KONTEXT:\n" +
"Workflow-Datei: " + fileName + "\n" +
"Workflow-Daten: " + shortContext + "\n\n" +
"ANWEISUNG:\n" +
"1. Identifiziere das exakte Hugging Face Repository (z.B. black-forest-labs/FLUX.1-schnell).\n" +
"2. Antworte NUR mit der Repository ID (Format: creator/repo) oder 'UNKNOWN'.\n" +
"3. Wenn du einen direkten Download-Link findest, gib diesen stattdessen aus.";
String prompt = "Web search: Determine the official Hugging Face repository for the file '" + modelName + "'.\n\n" +
"CONTEXT:\n" +
"Workflow file: " + fileName + "\n" +
"Workflow data: " + shortContext + "\n\n" +
"INSTRUCTION:\n" +
"1. Identify the exact Hugging Face repository (e.g., black-forest-labs/FLUX.1-schnell).\n" +
"2. Respond ONLY with the repository ID (format: creator/repo) or 'UNKNOWN'.\n" +
"3. If you find a direct download link, output it instead.";
JSONObject payload = new JSONObject();
JSONArray contents = new JSONArray();
src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java
+9 -7
View File
@@ -1,15 +1,17 @@
package de.tki.comfymodels.service.impl;
import org.springframework.stereotype.Service;
import jakarta.annotation.PostConstruct;
import java.util.*;
@Service
public class LocalAIService {
private static final Map<String, String[]> KNOWLEDGE_BASE = new LinkedHashMap<>();
private static final Map<String, Double> GLOBAL_IDF = new HashMap<>();
private final Map<String, String[]> KNOWLEDGE_BASE = new LinkedHashMap<>();
private final Map<String, Double> GLOBAL_IDF = new HashMap<>();
static {
@PostConstruct
public void init() {
// --- TOP RESEARCH LABS ---
add("black-forest-labs", "flux, flux1, flux2, bfl, schnell, dev, pro, mistral");
add("stabilityai", "sdxl, sd15, sd21, sd3, svd, cascade, turbo, lightning, stable, diffusion");
@@ -32,11 +34,11 @@ public class LocalAIService {
calculateIdf();
}
private static void add(String provider, String keywords) {
private void add(String provider, String keywords) {
KNOWLEDGE_BASE.put(provider, keywords.split(", "));
}
private static void calculateIdf() {
private void calculateIdf() {
Set<String> allTerms = new HashSet<>();
for (String[] terms : KNOWLEDGE_BASE.values()) allTerms.addAll(Arrays.asList(terms));
int docCount = KNOWLEDGE_BASE.size();
@@ -48,7 +50,7 @@ public class LocalAIService {
public Prediction predictProvider(String fileName) {
String name = fileName.toLowerCase();
// Spezial-Check für komplexe Dateinamen
// Special check for complex filenames
if (name.contains("mistral") && name.contains("flux2")) return new Prediction("black-forest-labs", 0.95);
if (name.contains("pony") || name.contains("v6")) return new Prediction("PonyDiffusion", 0.85);
@@ -78,7 +80,7 @@ public class LocalAIService {
return new Prediction(bestProvider, Math.min(0.99, maxScore / 3.5));
}
public static class Prediction {
public class Prediction {
public final String provider;
public final double confidence;
public Prediction(String provider, double confidence) { this.provider = provider; this.confidence = confidence; }
src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java
+2 -2
View File
@@ -221,8 +221,8 @@ public class ModelSearchService implements IModelSearchService {
public String formatSize(long bytes) {
if (bytes == -401) return "🔒 Auth Required";
if (bytes <= 0) return "Unknown";
if (bytes > 1024L * 1024L * 1024L) return String.format("%.2f GB", bytes / 1024.0 / 1024.0 / 1024.0);
return String.format("%.2f MB", bytes / 1024.0 / 1024.0);
if (bytes >= 1024L * 1024L * 1024L) return String.format(Locale.US, "%.2f GB", bytes / 1024.0 / 1024.0 / 1024.0);
return String.format(Locale.US, "%.2f MB", bytes / 1024.0 / 1024.0);
}
private boolean fetchHuggingFaceUrlInSpecificRepo(ModelInfo info, int rowIndex, String repoId,